From acce8d1432f8d26a2b6b836f7d619ca86b8c5fca Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 12 Oct 2026 21:33:39 -0500 Subject: [PATCH 01/79] feat: add shared lib Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 7 +- tests/models/common/__init__.py | 2 + tests/models/common/mono/__init__.py | 2 + tests/models/common/mono/kit.py | 290 +++++++++++++++++++++ tests/models/common/mono/reference.py | 252 ++++++++++++++++++ tests/models/common/mono/test_reference.py | 116 +++++++++ vllm/envs.py | 12 + vllm/models/common/mono/__init__.py | 13 + vllm/models/common/mono/abi.py | 41 +++ vllm/models/common/mono/failstop.py | 71 +++++ vllm/models/common/mono/guards.py | 183 +++++++++++++ vllm/models/common/mono/layer.py | 87 +++++++ vllm/models/common/mono/peer.py | 140 ++++++++++ vllm/models/common/mono/routing.py | 209 +++++++++++++++ vllm/models/common/mono/spec.py | 173 ++++++++++++ 15 files changed, 1596 insertions(+), 2 deletions(-) create mode 100644 tests/models/common/__init__.py create mode 100644 tests/models/common/mono/__init__.py create mode 100644 tests/models/common/mono/kit.py create mode 100644 tests/models/common/mono/reference.py create mode 100644 tests/models/common/mono/test_reference.py create mode 100644 vllm/models/common/mono/__init__.py create mode 100644 vllm/models/common/mono/abi.py create mode 100644 vllm/models/common/mono/failstop.py create mode 100644 vllm/models/common/mono/guards.py create mode 100644 vllm/models/common/mono/layer.py create mode 100644 vllm/models/common/mono/peer.py create mode 100644 vllm/models/common/mono/routing.py create mode 100644 vllm/models/common/mono/spec.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 552910e85a96..838fe71f3879 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -70,8 +70,9 @@ steps: - tests/models/test_deepseek_v41_replay_batch.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py + - tests/models/common/mono/ 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 + - 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/common/mono mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -92,8 +93,10 @@ steps: - 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 + - vllm/models/common/mono/ + - tests/models/common/mono/ 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_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 + - 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_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/common/mono - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/common/__init__.py b/tests/models/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/models/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/common/mono/__init__.py b/tests/models/common/mono/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/models/common/mono/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py new file mode 100644 index 000000000000..758fe2319fa8 --- /dev/null +++ b/tests/models/common/mono/kit.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The MonoKernel conformance kit (RFC #60904). + +Every model that ships a ``MonoSpec`` subclasses ``MonoConformance`` and fills +in the five hooks, which gives it the contract the library relies on: refusals +are complete, the spec module is importable without a kernel backend, the +capture sizes are the spec's widths, one forward decides once and alike on +every rank, caches re-bind when they move, attach does not move the model's +weights, cache checks report every mismatch, a failed step fails, served steps +are counted, and every width is warmed up before capture. + +The kit never launches a kernel: it drives the library through the spec's own +callbacks, so a spec whose kernels need a GPU still conforms on CPU. +""" + +import subprocess +import sys +from collections import defaultdict +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import replace +from typing import Any + +import pytest +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import ForwardContext, override_forward_context +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry +from vllm.models.common.mono import failstop +from vllm.models.common.mono.guards import ( + configure, + load_specs, + refuse_config, + select, + set_capture_sizes, +) +from vllm.models.common.mono.layer import attach, forward_layer +from vllm.models.common.mono.spec import MonoUnsupported + +KERNEL_BACKENDS = ("aiter", "flydsl") +"""Backends a spec module must not pull in: it is imported at config time, on +every platform, before anything has decided the spec can run.""" + + +@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 + + +def _spy(spec, field: str, calls: list): + """``spec`` with one callback recording its calls.""" + inner = getattr(spec, field) + + def recorded(*args): + calls.append(args) + return inner(*args) + + return replace(spec, **{field: recorded}) + + +class MonoConformance: + """What every ``MonoSpec`` must satisfy. Subclass it in the model's test + module and fill in the hooks below.""" + + specs_path: str + """``"package.module:MONO_SPECS"``, as the model passes it to ``attach``.""" + + 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 decode_metadata(self, config, model) -> Any: + """Attention metadata for a decode step the kernels take.""" + raise NotImplementedError + + def skipped_metadata(self, config, model) -> Any: + """Attention metadata for a step ``skip_step`` sends to vLLM's path.""" + raise NotImplementedError + + def bad_caches(self, config, model) -> dict[str, Any]: + """Caches that break at least two of ``bind``'s checks at once.""" + 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 mono(self, config, model): + return attach(model, config, self.specs_path) + + @pytest.fixture(autouse=True) + def _counters(self): + failstop.STEPS.clear() + failstop.INCIDENTS.clear() + yield + failstop.STEPS.clear() + failstop.INCIDENTS.clear() + + def spec(self, config): + spec, refusals = select(load_specs(self.specs_path), config) + assert spec is not None, refusals + return spec + + def test_refusals_report_every_reason(self, config): + config.lora_config = object() + config.speculative_config = object() + config.use_v2_model_runner = False + for spec in load_specs(self.specs_path): + reasons = refuse_config(spec, config) + assert sum("lora" in r for r in reasons) == 1, reasons + assert sum("speculative_decoding" in r for r in reasons) == 1, reasons + assert sum("Model Runner V2" in r for r in reasons) == 1, reasons + + def test_refusal_names_every_variant(self, config): + config.use_v2_model_runner = False + with pytest.raises(MonoUnsupported) as excinfo: + configure(config, self.specs_path) + for spec in load_specs(self.specs_path): + assert spec.name in str(excinfo.value) + + def test_spec_module_imports_no_kernel_backend(self): + module = self.specs_path.split(":")[0] + source = ( + "import sys; sys.path.insert(0, '.'); import importlib;" + f"importlib.import_module({module!r});" + "print('\\n'.join(sys.modules))" + ) + 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_capture_sizes_are_the_widths(self, config): + spec = self.spec(config) + widths = spec.widths(config) + config.compilation_config.cudagraph_capture_sizes = list( + range(1, max(widths) * 2 + 1) + ) + set_capture_sizes(spec, config) + sizes = config.compilation_config.cudagraph_capture_sizes + assert [s for s in sizes if s <= max(widths)] == sorted(widths) + + def test_decided_once_a_forward(self, config, model, mono): + asked: list = [] + mono.spec = _spy(mono.spec, "skip_step", asked) + with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): + self._run_layers(mono, model) + assert len(asked) == 1 + assert len(mono.kernels) > 1 + + def test_every_rank_decides_alike(self, config, model): + decisions = [] + for rank in range(2): + other = attach(self.make_model(config), config, self.specs_path) + other.ctx.rank = rank + with forward( + config, CUDAGraphMode.NONE, self.decode_metadata(config, model) + ): + decisions.append(other.begin_step(self._positions(other))) + assert decisions[0] == decisions[1] is None + + def test_skipped_step_runs_vllm_layers(self, config, model, mono): + with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): + mono.begin_once(self._positions(mono)) + assert mono.step is None + assert mono.last is not None + + def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): + off = max(mono.widths) + 1 + with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): + assert mono.begin_step(torch.zeros(off, dtype=torch.long)) == "width" + + def test_piecewise_runs_vllm_layers(self, config, model, mono): + meta = self.decode_metadata(config, model) + with forward(config, CUDAGraphMode.PIECEWISE, meta): + assert mono.begin_step(self._positions(mono)) == "piecewise" + + def test_caches_rebind_when_they_move(self, config, model, mono): + binds: list = [] + mono.spec = _spy(mono.spec, "bind", binds) + self._decode(config, model, mono) + assert len(binds) == 1 + self._decode(config, model, mono) + assert len(binds) == 1 + sfc = config.compilation_config.static_forward_context + for name in mono.cache_layers: + sfc[name].kv_cache = sfc[name].kv_cache.clone() + self._decode(config, model, mono) + assert len(binds) == 2 + + def test_bind_reports_every_mismatch(self, config, model, mono): + reasons = mono.spec.bind(mono.ctx, self.bad_caches(config, model)) + assert len(reasons) >= 2, reasons + + def test_mismatched_caches_refuse_the_step(self, config, model, mono): + sfc = config.compilation_config.static_forward_context + bad = self.bad_caches(config, model) + for name, cache in zip(mono.cache_layers, bad.values(), strict=False): + sfc[name].kv_cache = cache + with pytest.raises(MonoUnsupported): + self._decode(config, model, mono) + + def test_attach_does_not_move_weights(self, config): + model = self.make_model(config) + before = {n: p.data_ptr() for n, p in model.named_parameters()} + attach(model, config, self.specs_path) + after = {n: p.data_ptr() for n, p in model.named_parameters()} + assert before == after + + def test_fail_stop_raises(self, mono): + words = torch.zeros_like(mono.health()) + words[0] = 1 + with pytest.raises(RuntimeError, match="MonoKernel step failed"): + failstop.check(words) + assert failstop.INCIDENTS["0"] == 1 + + def test_fail_stop_warns_when_unsafe(self, mono, monkeypatch): + monkeypatch.setattr("vllm.envs.VLLM_MONO_UNSAFE_WARN", True) + words = torch.zeros_like(mono.health()) + words[0] = 1 + failstop.check(words) + assert failstop.INCIDENTS["0"] == 1 + + def test_healthy_step_passes(self, mono): + failstop.check(torch.zeros_like(mono.health())) + assert not failstop.INCIDENTS + + def test_counters_follow_full_replay(self, config, model, mono): + width = min(mono.widths) + with forward(config, CUDAGraphMode.FULL, self.decode_metadata(config, model)): + mono.begin_step(torch.zeros(width, dtype=torch.long)) + for _ in range(3): + mono.account(width, CUDAGraphMode.FULL) + assert failstop.drain() == {f"steps:{mono.spec.name}:mono": 3} + + def test_counters_name_the_skipped_path(self, config, model, mono): + width = min(mono.widths) + with forward(config, CUDAGraphMode.FULL, self.skipped_metadata(config, model)): + why = mono.begin_step(torch.zeros(width, dtype=torch.long)) + mono.account(width, CUDAGraphMode.FULL) + assert failstop.drain() == {f"steps:{mono.spec.name}:{why}": 1} + + def test_warmup_registered_for_every_width(self, config): + registry = JitWarmupRegistry(config) + with registry.activate(): + mono = attach(self.make_model(config), config, self.specs_path) + assert len(registry) >= 1 + by_width: dict[int, set] = defaultdict(set) + for (_, width), kernel in mono.kernels.items(): + by_width[width] |= set(kernel.warmup_keys()) + assert set(by_width) == set(mono.widths) + assert all(by_width.values()) + registry.warmup() + + def _positions(self, mono) -> torch.Tensor: + return torch.zeros(min(mono.widths), dtype=torch.long) + + def _decode(self, config, model, mono) -> None: + with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): + mono.begin_step(self._positions(mono)) + + def _run_layers(self, mono, model) -> None: + hidden = torch.zeros(min(mono.widths), model.config.hidden_size) + positions = self._positions(mono) + for i, layer in enumerate(model.layers): + hidden = forward_layer(mono, i, layer, positions, hidden) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py new file mode 100644 index 000000000000..ef34117ba546 --- /dev/null +++ b/tests/models/common/mono/reference.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The reference MonoKernel: a ``LayerKernel`` that calls the vLLM layer it +replaces, and the ``MONO_SPECS`` that declare it. + +It computes nothing of its own and launches nothing on a device, so it runs on +CPU in the conformance kit (``kit``). What it is for is the contract: a model +porting to MonoKernels can read one complete spec, and the library's own +behaviour can be tested without a GPU or a kernel backend. +""" + +from dataclasses import dataclass, replace +from types import SimpleNamespace +from typing import Any + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.models.common.mono.spec import CacheReq, MonoContext, MonoSpec, MonoStep + +CACHE = CacheReq(group="kv", dtype="bfloat16", layout="2nhd") + +_LAYOUT_DIMS = {"2nhd": 4} + + +class ReferenceLayer(torch.nn.Module): + """A decoder layer the reference kernel can replace: one weight, one + attention layer name, and ``dense`` for the layers a spec skips.""" + + def __init__(self, index: int, hidden: int, dense: bool = False): + super().__init__() + self.index = index + self.dense = dense + self.attn_name = f"layers.{index}.attn" + self.weight = torch.nn.Parameter(torch.ones(hidden) * index) + + def forward(self, positions: torch.Tensor, hidden: torch.Tensor): + return hidden + self.weight + + +class ReferenceModel(torch.nn.Module): + """A model shaped like the layer loop a MonoKernel hooks into.""" + + def __init__(self, layers: int = 4, hidden: int = 8): + super().__init__() + self.config = SimpleNamespace(num_hidden_layers=layers, hidden_size=hidden) + self.layers = torch.nn.ModuleList( + [ReferenceLayer(i, hidden, dense=i == 0) for i in range(layers)] + ) + self.start_layer = 0 + self.end_layer = layers + self.mono = None + + def step_health(self) -> torch.Tensor | None: + from vllm.models.common.mono import health + + return health(self.mono) + + +@dataclass +class Packed: + """What ``pack`` builds once at attach: views of the model's weights, never + a copy, and the error words every kernel of this rank shares.""" + + weights: dict[int, torch.Tensor] + words: torch.Tensor + + +class ReferenceKernel: + """One layer's launch at one width: calls the layer it replaces.""" + + def __init__(self, ctx: MonoContext, packed: Packed, layer, width: int): + self.cache_layers = (layer.attn_name,) + self.ctx = ctx + self.packed = packed + self.layer = layer + self.width = width + self.launched: list[int] = [] + self.compiled: list[Any] = [] + + def launch(self, step: MonoStep, positions: torch.Tensor, *args, **kwargs): + self.launched.append(step.num_tokens) + return self.layer(positions, *args, **kwargs) + + def health(self) -> torch.Tensor: + return self.packed.words + + def warmup_keys(self) -> list[Any]: + return [(self.layer.index, self.width)] + + def compile(self, key: Any) -> None: + self.compiled.append(key) + + def fail(self, rank: int = 0) -> None: + """Reports an in-kernel wait that gave up, as a launch would.""" + self.packed.words[rank] = 1 + + +def layers(config) -> list[int]: + return list(range(config.num_hidden_layers)) + + +def widths(vllm_config) -> tuple[int, ...]: + most = vllm_config.scheduler_config.max_num_seqs + return tuple(w for w in (1, 2, 4, 8) if w <= most) + + +def refuse_config(vllm_config) -> list[str]: + reasons = [] + if not widths(vllm_config): + reasons.append("max_num_seqs below the smallest width") + if vllm_config.model_config.dtype is not torch.bfloat16: + reasons.append(f"dtype {vllm_config.model_config.dtype}: needs bfloat16") + return reasons + + +def refuse_model(model) -> list[str]: + return [] if isinstance(model, ReferenceModel) else ["not a ReferenceModel"] + + +def skip_layer(ctx: MonoContext, layer) -> str | None: + return "dense layer" if layer.dense else None + + +def skip_step(ctx: MonoContext, step: MonoStep) -> str | None: + return "prefill" if step.attn_metadata.get("prefill") else None + + +def pack(model, ctx: MonoContext) -> Packed: + _, world, _ = _tp() + return Packed( + weights={i: layer.weight.data for i, layer in enumerate(model.layers)}, + words=torch.zeros(world, dtype=torch.int32, device=ctx.device), + ) + + +def repack(model, ctx: MonoContext, packed: Packed) -> None: + for i, layer in enumerate(model.layers): + packed.weights[i].copy_(layer.weight.data) + + +def bind(ctx: MonoContext, caches) -> list[str]: + reasons = [] + for name, cache in caches.items(): + if str(cache.dtype).removeprefix("torch.") != CACHE.dtype: + reasons.append(f"{name}: dtype {cache.dtype}, needs {CACHE.dtype}") + if cache.ndim != _LAYOUT_DIMS[CACHE.layout]: + reasons.append(f"{name}: {cache.ndim}d, needs layout {CACHE.layout}") + return reasons + + +def build(ctx: MonoContext, packed: Packed, layer, width: int) -> ReferenceKernel: + return ReferenceKernel(ctx, packed, layer, width) + + +def _tp(): + from vllm.models.common.mono.peer import tp_group + + return tp_group() + + +REFERENCE = MonoSpec( + name="reference", + supports=frozenset(), + caches={"attn": CACHE}, + weight_source="views", + layers=layers, + widths=widths, + refuse_config=refuse_config, + refuse_model=refuse_model, + skip_layer=skip_layer, + skip_step=skip_step, + pack=pack, + repack=repack, + bind=bind, + build=build, +) + +REFERENCE_WIDE = replace( + REFERENCE, + name="reference-wide", + refuse_config=lambda c: refuse_config(c) + + (["needs max_num_seqs >= 8"] if c.scheduler_config.max_num_seqs < 8 else []), +) + +MONO_SPECS = (REFERENCE_WIDE, REFERENCE) + + +def make_config(**overrides) -> Any: + """A config shaped like the parts of ``VllmConfig`` the library reads. + + Args: + **overrides: Top-level attributes to replace. Nested ones are set on + the returned object (``config.parallel_config.enable_eplb = True``). + + """ + config = SimpleNamespace( + kernel_config=SimpleNamespace(monokernel=True), + model_config=SimpleNamespace(dtype=torch.bfloat16, enable_sleep_mode=False), + scheduler_config=SimpleNamespace(max_num_seqs=4), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=False, + enable_eplb=False, + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + offload_config=SimpleNamespace( + uva=SimpleNamespace(cpu_offload_gb=0), + prefetch=SimpleNamespace(offload_group_size=0), + ), + aux_output_config=SimpleNamespace(enable_return_routed_experts=False), + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.FULL, + cudagraph_capture_sizes=[1, 2, 3, 4, 16], + static_forward_context={}, + ), + lora_config=None, + kv_transfer_config=None, + speculative_config=None, + weight_transfer_config=None, + use_v2_model_runner=True, + ) + for key, value in overrides.items(): + setattr(config, key, value) + return config + + +def make_model(config, layers: int = 4) -> ReferenceModel: + """A loaded reference model, with its KV caches already in the config's + static forward context.""" + model = ReferenceModel(layers=layers) + set_caches(config, model) + return model + + +def set_caches(config, model, **kwargs) -> dict[str, torch.Tensor]: + """Puts one cache a layer in the static forward context, and returns them.""" + caches = { + layer.attn_name: make_cache(**kwargs) + for layer in model.layers + if not layer.dense + } + config.compilation_config.static_forward_context = { + name: SimpleNamespace(kv_cache=cache) for name, cache in caches.items() + } + return caches + + +def make_cache(dtype=torch.bfloat16, ndim: int = 4) -> torch.Tensor: + return torch.zeros((2,) + (1,) * (ndim - 1), dtype=dtype) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py new file mode 100644 index 000000000000..c622fdbf8407 --- /dev/null +++ b/tests/models/common/mono/test_reference.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The conformance kit against the reference spec: what the kit asserts is the +library's behaviour, so this is also the library's own test.""" + +from dataclasses import replace + +import pytest +import torch + +from tests.models.common.mono import reference +from tests.models.common.mono.kit import MonoConformance, forward +from vllm.config.compilation import CUDAGraphMode +from vllm.models.common.mono.guards import enabled, select +from vllm.models.common.mono.layer import attach, forward_layer +from vllm.models.common.mono.spec import MonoUnsupported + +SPECS_PATH = "tests.models.common.mono.reference:MONO_SPECS" + + +class TestReferenceSpec(MonoConformance): + specs_path = SPECS_PATH + + def make_config(self, **overrides): + return reference.make_config(**overrides) + + def make_model(self, config): + return reference.make_model(config) + + def decode_metadata(self, config, model): + return {"prefill": False} + + def skipped_metadata(self, config, model): + return {"prefill": True} + + def bad_caches(self, config, model): + return { + "wrong": reference.make_cache(dtype=torch.float16, ndim=2), + } + + +def test_switch_off_means_no_attach(): + config = reference.make_config() + config.kernel_config.monokernel = False + assert not enabled(config) + assert attach(reference.make_model(config), config, SPECS_PATH) is None + + +def test_the_first_servable_variant_wins(): + config = reference.make_config() + config.scheduler_config.max_num_seqs = 4 + spec, refusals = select(reference.MONO_SPECS, config) + assert spec is reference.REFERENCE + assert refusals["reference-wide"] == ["needs max_num_seqs >= 8"] + + config.scheduler_config.max_num_seqs = 8 + spec, refusals = select(reference.MONO_SPECS, config) + assert spec is reference.REFERENCE_WIDE + assert refusals == {} + + +def test_skipped_layers_keep_vllm_path(): + config = reference.make_config() + model = reference.make_model(config) + mono = attach(model, config, SPECS_PATH) + assert {i for i, _ in mono.kernels} == {1, 2, 3} + + +def test_kernel_layers_follow_the_pipeline_slice(): + config = reference.make_config() + model = reference.make_model(config) + model.start_layer, model.end_layer = 2, 4 + mono = attach(model, config, SPECS_PATH) + assert {i for i, _ in mono.kernels} == {2, 3} + + +def test_kernels_run_the_layers_they_replace(): + config = reference.make_config() + model = reference.make_model(config) + mono = attach(model, config, SPECS_PATH) + width = min(mono.widths) + positions = torch.zeros(width, dtype=torch.long) + hidden = torch.zeros(width, model.config.hidden_size) + with forward(config, CUDAGraphMode.NONE, {"prefill": False}): + for i, layer in enumerate(model.layers): + hidden = forward_layer(mono, i, layer, positions, hidden) + assert torch.equal(hidden, torch.full_like(hidden, 0 + 1 + 2 + 3)) + assert [k.launched for k in mono.kernels.values() if k.launched] == [[width]] * 3 + + +def test_repack_writes_in_place(): + config = reference.make_config() + model = reference.make_model(config) + mono = attach(model, config, SPECS_PATH) + before = {i: w.data_ptr() for i, w in mono.packed.weights.items()} + with torch.no_grad(): + for layer in model.layers: + layer.weight.copy_(torch.full_like(layer.weight, 7)) + mono.repack(model) + assert {i: w.data_ptr() for i, w in mono.packed.weights.items()} == before + assert all((w == 7).all() for w in mono.packed.weights.values()) + + +def test_repack_is_refused_when_the_spec_declares_none(): + config = reference.make_config() + model = reference.make_model(config) + mono = attach(model, config, SPECS_PATH) + mono.spec = replace(mono.spec, repack=None) + with pytest.raises(MonoUnsupported, match="weights cannot move"): + mono.repack(model) + + +def test_model_refusal_is_raised(): + config = reference.make_config() + with pytest.raises(MonoUnsupported, match="not a ReferenceModel"): + attach(torch.nn.Linear(2, 2), config, SPECS_PATH) diff --git a/vllm/envs.py b/vllm/envs.py index a2b46030772f..dc0751cebdbd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -157,6 +157,8 @@ VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False VLLM_ROCM_MONO_DECODE: bool = False + VLLM_MONO_UNSAFE_WARN: bool = False + VLLM_MONO_FAULT_INJECT: str = "" VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1404,6 +1406,16 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ROCM_MONO_DECODE": lambda: ( os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1") ), + # Warn instead of failing the step when a mono decode kernel reports an + # error word (an in-kernel wait gave up). For debugging only: the step's + # output is not trustworthy. By default is disabled. + "VLLM_MONO_UNSAFE_WARN": lambda: ( + os.getenv("VLLM_MONO_UNSAFE_WARN", "False").lower() in ("true", "1") + ), + # Make one mono decode launch report a fault, to test the fail-stop path: + # "rank:epoch:us" sets the rank, the step epoch and how long the injected + # stall lasts. By default is disabled. + "VLLM_MONO_FAULT_INJECT": lambda: os.getenv("VLLM_MONO_FAULT_INJECT", ""), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py new file mode 100644 index 000000000000..524fe5008c19 --- /dev/null +++ b/vllm/models/common/mono/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MonoKernels (RFC #60904): persistent kernels that run most or all of one +decoder layer's decode step in a single launch. A model brings a ``MonoSpec`` +(``spec``); this package owns the rest: refusals (``guards``), the per-step +decision and cache re-binding (``routing``), the layer-loop hook (``layer``), +fail-stop and metrics (``failstop``), peer memory (``peer``) and the +error-word layout (``abi``).""" + +from vllm.models.common.mono.guards import configure +from vllm.models.common.mono.layer import account_step, attach, forward_layer, health + +__all__ = ["account_step", "attach", "configure", "forward_layer", "health"] diff --git a/vllm/models/common/mono/abi.py b/vllm/models/common/mono/abi.py new file mode 100644 index 000000000000..d684416843b6 --- /dev/null +++ b/vllm/models/common/mono/abi.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The layout every MonoKernel's error words and fault hook follow. + +A rank's health words: a header whose word FAILED is set once any of its +in-kernel waits gave up (sticky: a failed rank stays failed, and its later waits +exit at once), then one record a CTA of the first wait it gave up on (flag, +region id, pair index, tag, seen tag, CTA, clock, pad). + +The output rank's flags: one int32 word a TP rank in the first PEER_HEAD bytes +of the output rank's peer buffer. A rank whose wait gives up writes its word +there, drains the store, then sets FAILED, so nothing it stores later can +overtake the flag. The runner reads these words after every step. + +The fault hook (tests only, ``VLLM_MONO_FAULT_INJECT=rank:epoch:us``): a +build-time stall of one CTA of rank ``rank`` for ``us`` microseconds, once, in +the step at ``epoch``; no code when unset. +""" + +from vllm import envs + +MONO_ABI_VERSION = 1 + +RECORD_WORDS = 8 +FAILED = 0 +PEER_HEAD = 256 +WAIT_SECONDS = 10 + + +def health_words(ctas: int) -> int: + """A rank's health words for a grid of ``ctas``: the header, a record a CTA.""" + return (ctas + 1) * RECORD_WORDS + + +def fault_inject() -> tuple[int, int, int] | None: + """(rank, epoch, microseconds) of ``VLLM_MONO_FAULT_INJECT``, or None.""" + v = envs.VLLM_MONO_FAULT_INJECT + if not v: + return None + rank, epoch, us = (int(x) for x in v.split(":")) + return rank, epoch, us diff --git a/vllm/models/common/mono/failstop.py b/vllm/models/common/mono/failstop.py new file mode 100644 index 000000000000..439952886db3 --- /dev/null +++ b/vllm/models/common/mono/failstop.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fail-stop and metrics. + +A MonoKernel never returns tokens from a step it did not finish: the runner +copies the model's error words (``abi``) next to the sampled token ids, and +``check`` raises before that step's tokens are returned. The engine then stops +like on any other engine death, and the deployment's restart policy recovers +it. ``VLLM_MONO_UNSAFE_WARN=1`` (tests and benchmarks only, no CLI flag) logs +the incident and keeps serving instead. + +Counters (``drain``, exported as ``vllm:monokernel_steps`` and +``vllm:monokernel_incidents``): steps on the kernels, steps on vLLM's path by +reason, and incidents by failed rank. +""" + +from collections import Counter + +import torch + +from vllm import envs +from vllm.logger import init_logger + +logger = init_logger(__name__) + +STEPS: Counter[tuple[str, str]] = Counter() +INCIDENTS: Counter[str] = Counter() + + +def count_step(spec: str, why: str | None) -> None: + """Counts one served step of ``spec``: on the kernels when ``why`` is None, + else on vLLM's layers for that reason.""" + STEPS[spec, why or "mono"] += 1 + + +def drain() -> dict[str, int] | None: + """The counts since the last drain: ``"steps::"`` (path "mono" + or the reason the step ran vLLM's layers) and ``"incidents:"``.""" + if not STEPS and not INCIDENTS: + return None + out = {f"steps:{spec}:{path}": n for (spec, path), n in STEPS.items()} + out.update({f"incidents:{rank}": n for rank, n in INCIDENTS.items()}) + STEPS.clear() + INCIDENTS.clear() + return out + + +def check(words: torch.Tensor) -> None: + """Fails the step whose error words these are, before its tokens return. + + Args: + words: The output rank's error words for one step, on the host (a word + a TP rank, nonzero once one of its in-kernel waits gave up). + + Raises: + RuntimeError: A rank gave up a wait, unless ``VLLM_MONO_UNSAFE_WARN``. + + """ + if not words.any(): + return + ranks = [r for r, w in enumerate(words.tolist()) if w] + for r in ranks: + INCIDENTS[str(r)] += 1 + msg = ( + f"MonoKernel step failed: TP ranks {ranks} gave up an in-kernel wait; " + "the step's tokens are discarded" + ) + if envs.VLLM_MONO_UNSAFE_WARN: + logger.error("%s (VLLM_MONO_UNSAFE_WARN=1: serving on)", msg) + return + raise RuntimeError(msg) diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py new file mode 100644 index 000000000000..0b3baf4fbbd1 --- /dev/null +++ b/vllm/models/common/mono/guards.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Refusals. ``configure`` runs once from core config post-init, before load: +it imports only the model's spec module, picks the variant this config can +serve, adds its widths to the CUDA graph capture sizes and raises with every +reason the config is refused. ``agree`` runs a later refusal stage (after load) +on every TP rank and agrees on the verdict before anything collective follows. +""" + +import importlib +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING + +from vllm import envs +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.models.common.mono.peer import tp_agree +from vllm.models.common.mono.spec import Feature, MonoSpec, MonoUnsupported + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + +_FULL_MODES = ( + CUDAGraphMode.NONE, + CUDAGraphMode.FULL, + CUDAGraphMode.FULL_DECODE_ONLY, + CUDAGraphMode.FULL_AND_PIECEWISE, +) + + +def enabled(vllm_config: "VllmConfig") -> bool: + """``--monokernel``, or ``VLLM_ROCM_MONO_DECODE=1`` (#60397's switch). + + The ``KernelConfig`` flag lands in a separate core PR, so a tree without it + reads as off rather than failing to import. + """ + flag = getattr(vllm_config.kernel_config, "monokernel", False) + return bool(flag) or envs.VLLM_ROCM_MONO_DECODE + + +def load_specs(path: str) -> tuple[MonoSpec, ...]: + """``"package.module:NAME"`` -> the spec variants, importing only that + module. ``NAME`` is a ``MonoSpec`` or a tuple of them, in the order they + are preferred.""" + module, name = path.split(":") + specs = getattr(importlib.import_module(module), name) + return (specs,) if isinstance(specs, MonoSpec) else tuple(specs) + + +def features_in_use(vllm_config: "VllmConfig") -> set[Feature]: + """The deployment's features a spec has to list in ``supports``.""" + pc = vllm_config.parallel_config + oc = vllm_config.offload_config + on = { + Feature.PIPELINE_PARALLEL: pc.pipeline_parallel_size > 1, + Feature.DATA_PARALLEL: pc.data_parallel_size > 1, + Feature.EXPERT_PARALLEL: pc.enable_expert_parallel, + Feature.EPLB: pc.enable_eplb, + Feature.DECODE_CONTEXT_PARALLEL: pc.decode_context_parallel_size > 1, + Feature.PREFILL_CONTEXT_PARALLEL: pc.prefill_context_parallel_size > 1, + Feature.LORA: vllm_config.lora_config is not None, + Feature.KV_CONNECTOR: vllm_config.kv_transfer_config is not None, + Feature.ROUTED_EXPERT_CAPTURE: ( + vllm_config.aux_output_config.enable_return_routed_experts + ), + Feature.SPECULATIVE_DECODING: vllm_config.speculative_config is not None, + Feature.SLEEP_MODE: vllm_config.model_config.enable_sleep_mode, + Feature.CPU_WEIGHT_OFFLOAD: ( + oc.uva.cpu_offload_gb > 0 or oc.prefetch.offload_group_size > 0 + ), + Feature.WEIGHT_UPDATES: vllm_config.weight_transfer_config is not None, + } + return {f for f, used in on.items() if used} + + +def refuse_config(spec: MonoSpec, vllm_config: "VllmConfig") -> list[str]: + """Every reason this config cannot run ``spec``: default-deny on features, + then the runner and the graph mode, then the spec's own.""" + reasons = [ + f"{f.value} is not supported" + for f in sorted(features_in_use(vllm_config) - spec.supports) + ] + if not vllm_config.use_v2_model_runner: + reasons.append("needs Model Runner V2") + mode = vllm_config.compilation_config.cudagraph_mode + if mode is not None and mode not in _FULL_MODES: + reasons.append(f"CUDA graph mode {mode.name}: decode needs FULL graphs") + return reasons + list(spec.refuse_config(vllm_config)) + + +def refuse_model(spec: MonoSpec, model) -> list[str]: + """Every reason the loaded model cannot run ``spec``. A MonoKernel replaces + the layer loop, which a traced graph has already inlined, so a model + carrying ``@support_torch_compile`` is refused.""" + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + + reasons = [] + if isinstance(model, TorchCompileWithNoGuardsWrapper): + reasons.append( + "model is torch.compile'd: MonoKernels replace the layer loop, " + "so the model must not carry @support_torch_compile" + ) + return reasons + list(spec.refuse_model(model)) + + +def select( + specs: Sequence[MonoSpec], vllm_config: "VllmConfig" +) -> tuple[MonoSpec | None, dict[str, list[str]]]: + """The first variant this config can serve, and every variant's refusals. + + Returns: + (spec, refusals). ``spec`` is None when every variant refused, and + ``refusals`` holds each variant's reasons, by spec name. + + """ + refusals: dict[str, list[str]] = {} + for spec in specs: + reasons = refuse_config(spec, vllm_config) + if not reasons: + return spec, refusals + refusals[spec.name] = reasons + return None, refusals + + +def refused(refusals: dict[str, list[str]]) -> str: + """Every variant's reasons as one message.""" + return "; ".join( + f"{name}: {', '.join(reasons)}" for name, reasons in refusals.items() + ) + + +def set_capture_sizes(spec: MonoSpec, vllm_config: "VllmConfig") -> None: + """The capture sizes up to the largest width become exactly the widths, so + a replay never pads a request's rows; larger sizes stay as configured (and + capture vLLM's path).""" + cc = vllm_config.compilation_config + widths = spec.widths(vllm_config) + sizes = cc.cudagraph_capture_sizes or [] + new = sorted(set(widths) | {s for s in sizes if s > max(widths)}) + if new != sorted(sizes): + logger.info("MonoKernel %s: CUDA graph capture sizes %s", spec.name, new) + cc.cudagraph_capture_sizes = new + + +def configure(vllm_config: "VllmConfig", specs_path: str) -> None: + """Picks the variant, sets the capture sizes, or raises with every reason. + + Args: + vllm_config: The engine's config, before the weights are loaded. + specs_path: ``"package.module:MONO_SPECS"``. + + Raises: + MonoUnsupported: Every variant refused this config. + + """ + spec, refusals = select(load_specs(specs_path), vllm_config) + if spec is None: + raise MonoUnsupported("MonoKernel refused: " + refused(refusals)) + set_capture_sizes(spec, vllm_config) + if vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: + logger.warning( + "MonoKernel %s under eager mode: for debugging only, serve with FULL " + "CUDA graphs", + spec.name, + ) + + +def agree(refuse: Callable[[], list[str]], group, what: str) -> None: + """``refuse``'s reasons on this rank, then every TP rank's verdict agreed: a + rank that refused alone would leave the others spinning in-kernel. + + Raises: + MonoUnsupported: This rank refused, or another rank did. + + """ + reasons = refuse() + agreed = tp_agree(not reasons, group) + if reasons: + raise MonoUnsupported(f"{what}: " + "; ".join(reasons)) + if not agreed: + raise MonoUnsupported(f"{what}: another TP rank refused") diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py new file mode 100644 index 000000000000..22fc5acd66ce --- /dev/null +++ b/vllm/models/common/mono/layer.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The hooks a model calls: ``attach`` after load, ``forward_layer`` in its +layer loop, ``health`` from ``step_health``.""" + +from typing import TYPE_CHECKING + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.models.common.mono import failstop +from vllm.models.common.mono.guards import enabled, load_specs, refused, select +from vllm.models.common.mono.routing import Mono +from vllm.models.common.mono.spec import MonoUnsupported + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + +_ATTACHED: Mono | None = None + + +def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: + """Picks the variant, packs the weights and builds this rank's layer + kernels. Collective over the TP group. + + Args: + model: The loaded model, in ``process_weights_after_loading``. + vllm_config: The engine's config. + specs_path: ``"package.module:MONO_SPECS"``. + + Returns: + This rank's ``Mono``, or None with the switch off or when no layer + takes a kernel. + + Raises: + MonoUnsupported: Every variant refused the config, or this rank (or + another) refused the loaded model. + + """ + global _ATTACHED + if not enabled(vllm_config): + return None + spec, refusals = select(load_specs(specs_path), vllm_config) + if spec is None: + raise MonoUnsupported("MonoKernel refused: " + refused(refusals)) + mono = Mono(spec, model, vllm_config) + if not mono.kernels: + logger.info("MonoKernel %s: no layer takes a kernel", spec.name) + return None + if not hasattr(model, "step_health"): + logger.warning_once( + "MonoKernel %s: the model exposes no step_health(), so the runner " + "cannot fail a step whose in-kernel wait gave up", + spec.name, + ) + _ATTACHED = mono + return mono + + +def forward_layer(mono: Mono | None, i: int, layer, positions, *args, **kwargs): + """Layer ``i`` on its kernel when this step runs on the kernels, else the + vLLM layer unchanged. The first call under a forward context decides for + every layer.""" + if mono is not None: + mono.begin_once(positions) + step = mono.step + if step is not None and (k := mono.kernels.get((i, step.width))) is not None: + return k.launch(step, positions, *args, **kwargs) + return layer(positions, *args, **kwargs) + + +def health(mono: Mono | None) -> torch.Tensor | None: + """The step's error words, read by the runner after each step: any set + word fails the step before its tokens are returned (``failstop``).""" + return None if mono is None else mono.health() + + +def account_step(num_tokens: int, runtime_mode: CUDAGraphMode) -> dict[str, int] | None: + """Called by the runner after every served step: counts it, and returns the + counters to export since the last call (``failstop.drain``).""" + if _ATTACHED is None: + return None + _ATTACHED.account(num_tokens, runtime_mode) + return failstop.drain() diff --git a/vllm/models/common/mono/peer.py b/vllm/models/common/mono/peer.py new file mode 100644 index 000000000000..d062f779ee3c --- /dev/null +++ b/vllm/models/common/mono/peer.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/runtime/peer_memory.py +# atom/mono/runtime/consensus.py +"""Peer memory for in-kernel TP collectives: one uncached buffer a rank, every +rank holding every peer's address, and the TP agreements that must precede +anything collective.""" + +import torch +import torch.distributed as dist + +from vllm.models.common.mono.spec import MonoUnsupported + + +def tp_group() -> tuple[int, int, object | None]: + """This rank's (rank, world size, CPU group) in the TP group, and + (0, 1, None) before distributed is initialized.""" + if not dist.is_available() or not dist.is_initialized(): + return 0, 1, None + from vllm.distributed import get_tp_group + + tp = get_tp_group() + return tp.rank_in_group, tp.world_size, tp.cpu_group + + +class _DeviceBytes: + """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, + for a non-owning torch view.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address + (``addresses``, int64, indexed by TP rank). What lives where, and why no + rank overwrites a peer's unread data, is the model's layout's to say. + + ``bytes`` is the whole buffer: the kernels tag every pair with the step's + epoch, so it is never zeroed between steps. + + The constructor is collective over ``group`` (the handle exchange, then an + agreement that every rank mapped every peer): every rank must reach it, so + a runner reaches it only after ``tp_agree``. A failed mapping on any rank + raises ``MonoUnsupported`` on all of them. + """ + + def __init__( + self, + nbytes: int, + group, + rank: int, + npes: int, + device: torch.device, + ): + from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + self._heap = UncachedIpcHeap + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) + self.rank = rank + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses, failure = [], None + try: + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + except RuntimeError as err: + failure = err + if not tp_agree(failure is None, group): + self.close() + raise MonoUnsupported( + "peer memory handshake failed on " + f"{'this' if failure else 'another'}" + " TP rank" + ) from failure + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def close(self) -> None: + """Closes the peer mappings and frees the local buffer; idempotent. + Only once no kernel of any rank can still touch it (the runner is being + dropped).""" + opened, self._opened = self._opened, [] + for base in opened: + self._heap.close_mem_handle(base) + if self.local: + self.bytes = None + self._heap.free_device_mem(self.local) + self.local = 0 + + +def tp_agree(ok: bool, group) -> bool: + """Agrees a verdict every TP rank must take alike. + + A MonoKernel waits in-kernel on every peer rank, so a rank that turned mono + off alone, or failed to build what the others launch, leaves the others + spinning on the GPU for ever. Every rank must call this at the same point. + + Args: + ok: This rank's local verdict. + group: The TP group's CPU (gloo) process group, or None for one rank. + + Returns: + True when every rank of ``group`` passed ``ok`` true. + + """ + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + +def bind_agreed(bind, group) -> None: + """``bind()`` on this rank (raising ``MonoUnsupported`` to refuse), then + every rank's verdict agreed: a rank that raised re-raises, the others raise + too.""" + try: + bind() + except Exception: + tp_agree(False, group) + raise + if not tp_agree(True, group): + raise MonoUnsupported("another TP rank refused mono") diff --git a/vllm/models/common/mono/routing.py b/vllm/models/common/mono/routing.py new file mode 100644 index 000000000000..4201f0c6119c --- /dev/null +++ b/vllm/models/common/mono/routing.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The per-step decision and cache re-binding. + +``Mono.begin_step`` decides once a forward, for every kernel layer, from what +every rank sees alike: the batch's width and the attention metadata, never +rank-local state and never a collective. Under a FULL CUDA graph the decision +is baked in at capture, so it is a function of the captured size: FULL graphs +serve uniform decode batches and every width is a capture size +(``guards.set_capture_sizes``), so no replay pads a request's rows. PIECEWISE +steps run vLLM's layers. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel +from vllm.models.common.mono import failstop +from vllm.models.common.mono.guards import agree, refuse_model +from vllm.models.common.mono.peer import tp_group +from vllm.models.common.mono.spec import ( + LayerKernel, + MonoContext, + MonoSpec, + MonoStep, + MonoUnsupported, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): + """A spec's layer-kernel builds as vLLM JIT warm-up keys.""" + + @dataclass(frozen=True) + class CompileKey: + spec: str + build: Any + + def __init__(self, spec: str, owners: dict[Any, LayerKernel]): + super().__init__() + self.spec = spec + self.owners = owners + + def get_warmup_keys(self) -> list["MonoWarmup.CompileKey"]: + return [self.CompileKey(self.spec, b) for b in self.owners] + + def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: + self.owners[compile_key.build].compile(compile_key.build) + self._compiled_cache[compile_key] = True + + +class Mono: + """A model's MonoKernels on this rank: built once after load.""" + + def __init__(self, spec: MonoSpec, model, vllm_config: "VllmConfig"): + rank, _, cpu_group = tp_group() + self.spec = spec + agree( + lambda: refuse_model(spec, model), + cpu_group, + f"MonoKernel {spec.name} refused for the loaded model", + ) + self.ctx = MonoContext(vllm_config, rank, next(model.parameters()).device) + self.widths = spec.widths(vllm_config) + self.kernels: dict[tuple[int, int], LayerKernel] = {} + self.packed: Any = None + self.step: MonoStep | None = None + self.baked: dict[int, str | None] = {} + self.last: str | None = "no step" + self._fc: Any = None + self._bound: tuple | None = None + + indices = self._kernel_layers(model) + if not indices: + self.cache_layers: list[str] = [] + return + self.packed = spec.pack(model, self.ctx) + for i in indices: + for w in self.widths: + self.kernels[i, w] = spec.build( + self.ctx, self.packed, model.layers[i], w + ) + self.cache_layers = sorted( + {n for k in self.kernels.values() for n in k.cache_layers} + ) + self._register_warmup() + logger.info( + "MonoKernel %s: %d layers, widths %s", + spec.name, + len(indices), + list(self.widths), + ) + + def _kernel_layers(self, model) -> list[int]: + start = getattr(model, "start_layer", 0) + end = getattr(model, "end_layer", len(model.layers)) + out = [] + for i in self.spec.layers(model.config): + if not start <= i < end: + continue + why = self.spec.skip_layer(self.ctx, model.layers[i]) + if why is None: + out.append(i) + else: + logger.info_once( + "MonoKernel %s: layer %d on vLLM's path (%s)", + self.spec.name, + i, + why, + ) + return out + + def _register_warmup(self) -> None: + owners: dict[Any, LayerKernel] = {} + for k in self.kernels.values(): + for key in k.warmup_keys(): + owners.setdefault(key, k) + if owners: + MonoWarmup(self.spec.name, owners).register_warmup() + + def repack(self, model) -> None: + """Writes ``model``'s current weights into the allocation ``pack`` + returned, in place. + + Raises: + MonoUnsupported: The spec declares no ``repack``, so its weights + cannot move after attach. + + """ + if self.spec.repack is None: + raise MonoUnsupported( + f"MonoKernel {self.spec.name}: weights cannot move after attach" + ) + self.spec.repack(model, self.ctx, self.packed) + + def _bind(self) -> None: + """Re-checks the caches whenever they moved (keyed on every tensor's + data_ptr and numel): memory profiling runs real steps on throwaway + caches before the real ones exist.""" + sfc = self.ctx.vllm_config.compilation_config.static_forward_context + caches = {n: sfc[n].kv_cache for n in self.cache_layers} + key = tuple( + (t.data_ptr(), t.numel()) + for c in caches.values() + for t in (c if isinstance(c, (list, tuple)) else (c,)) + ) + if key == self._bound: + return + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise MonoUnsupported( + f"MonoKernel {self.spec.name}: caches moved during CUDA graph capture" + ) + reasons = self.spec.bind(self.ctx, caches) + if reasons: + raise MonoUnsupported( + f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) + ) + self._bound = key + + def begin_once(self, positions: torch.Tensor) -> None: + """Decides this forward's step, on the first kernel layer to ask.""" + fc = get_forward_context() + if fc is not self._fc: + self._fc = fc + self.begin_step(positions) + + def begin_step(self, positions: torch.Tensor) -> str | None: + """None runs this step on the kernels (and sets ``step``), else the + reason it runs vLLM's layers.""" + self.step = None + why = self._decide(positions) + if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: + self.baked[positions.size(-1)] = why + self.last = why + return why + + def account(self, num_tokens: int, runtime_mode: CUDAGraphMode) -> None: + """Counts one served step, after its forward.""" + if runtime_mode == CUDAGraphMode.FULL: + why = self.baked.get(num_tokens, "not captured") + else: + why, self.last = self.last, "no step" + failstop.count_step(self.spec.name, why) + + def _decide(self, positions: torch.Tensor) -> str | None: + fc = get_forward_context() + if fc.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + return "piecewise" + n = positions.size(-1) + if n not in self.widths: + return "width" + step = MonoStep(n, n, positions, fc.attn_metadata, fc.cudagraph_runtime_mode) + why = self.spec.skip_step(self.ctx, step) + if why is None: + self._bind() + self.step = step + return why + + def health(self) -> torch.Tensor: + return next(iter(self.kernels.values())).health() diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py new file mode 100644 index 000000000000..a61bc724db72 --- /dev/null +++ b/vllm/models/common/mono/spec.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""What a model declares to run MonoKernels: its ``MonoSpec``, and the types +the spec's callbacks see.""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Any, Literal, Protocol + +import torch + +from vllm.config.compilation import CUDAGraphMode + +if TYPE_CHECKING: + from vllm.config import VllmConfig + + +class Feature(str, Enum): + """What a deployment can turn on. A feature in use that a spec's + ``supports`` does not list is refused.""" + + PIPELINE_PARALLEL = "pipeline_parallel" + DATA_PARALLEL = "data_parallel" + EXPERT_PARALLEL = "expert_parallel" + EPLB = "eplb" + DECODE_CONTEXT_PARALLEL = "decode_context_parallel" + PREFILL_CONTEXT_PARALLEL = "prefill_context_parallel" + LORA = "lora" + KV_CONNECTOR = "kv_connector" + ROUTED_EXPERT_CAPTURE = "routed_expert_capture" + SPECULATIVE_DECODING = "speculative_decoding" + SLEEP_MODE = "sleep_mode" + CPU_WEIGHT_OFFLOAD = "cpu_weight_offload" + WEIGHT_UPDATES = "weight_updates" + + +class MonoUnsupported(Exception): + """The configuration or the loaded model is outside what a spec serves.""" + + +@dataclass(frozen=True) +class CacheReq: + """A cache the kernels read: its KV cache group, dtype and layout.""" + + group: str + dtype: str + layout: str + + +@dataclass +class MonoContext: + """One rank's view of the deployment, built once at attach. + + Attributes: + vllm_config: The engine's config. + rank: This rank's index in the TP group. + device: The device the kernels launch on. + state: The spec's own per-rank state, whatever ``pack`` keeps there. + Opaque to the library, which never reads it. + + """ + + vllm_config: "VllmConfig" + rank: int + device: torch.device + state: Any = None + + +@dataclass +class MonoStep: + """One forward's batch descriptor, built once and read by every layer. + + Attributes: + num_tokens: The step's query rows. + width: The padded rows the launch serves. + positions: The step's position ids. + attn_metadata: The forward context's attention metadata, by layer name. + runtime_mode: The CUDA graph mode this step runs under. + + """ + + num_tokens: int + width: int + positions: torch.Tensor + attn_metadata: Any + runtime_mode: CUDAGraphMode + + +class LayerKernel(Protocol): + """One decoder layer's launch, at one width, on this rank.""" + + cache_layers: tuple[str, ...] + """Attention layers whose caches ``MonoSpec.bind`` checks.""" + + def launch(self, step: MonoStep, *args, **kwargs) -> Any: + """Takes and returns what the vLLM layer it replaces does.""" + ... + + def health(self) -> torch.Tensor: + """This rank's error words, in the ``abi`` layout.""" + ... + + def warmup_keys(self) -> list[Any]: + """Hashable build keys of this width's launches, compiled before CUDA + graph capture by vLLM's JIT warm-up.""" + ... + + def compile(self, key: Any) -> None: + """Compiles one of ``warmup_keys`` without launching it.""" + ... + + +@dataclass(frozen=True) +class MonoSpec: + """One (model, vendor, ISA)'s declaration: values and callbacks, no + behaviour of its own. + + A spec is a value, so a second ISA is ``dataclasses.replace`` of the first + at module scope rather than a subclass, and a model exports its candidates + as a ``MONO_SPECS`` tuple the library selects from. + + Gating runs in tiers, and the prefix is the consequence. ``refuse_*`` says + the deployment asked for something it cannot have, so it may take the + engine down; ``skip_*`` says only that this layer, or this step, takes + vLLM's path, and never raises. If the same model with the same flags could + legitimately have another layer, or a later step, that passes the check, it + is a ``skip_*``. Every tier reports all of its reasons, never the first. + + Attributes: + name: Log and metric name. + supports: Features this spec serves. Default deny: a feature in use + that is absent from it refuses the deployment. + caches: The caches the kernels read, by attention layer name. + weight_source: ``"views"`` to read vLLM's post-load layout in place, + ``"convert"`` to pack a second layout once after load. + layers: The decoder layers that get a kernel, from the model's config. + widths: The rows a launch serves, from the engine's config. These + become the CUDA graph capture sizes. + refuse_config: Deployment refusals the engine's config answers, before + the weights are loaded. + refuse_model: Deployment refusals only the loaded model answers. + skip_layer: Why one layer takes vLLM's path, None when it takes the + kernel's. + skip_step: Why one step takes vLLM's path, None when it takes the + kernel's. + pack: The weights and launch state, built once at attach. The only + place a weight copy is ever made. + repack: Writes new weights into the allocation ``pack`` returned, + in place, because the kernels and the captured graphs hold its + address. None refuses anything that moves weights after attach. + bind: Checks the KV caches the kernels will read, on every re-bind. + build: One ``LayerKernel`` for (context, packed, layer, width). + + """ + + name: str + supports: frozenset[Feature] + caches: Mapping[str, CacheReq] + weight_source: Literal["views", "convert"] + + layers: Callable[[Any], list[int]] + widths: Callable[["VllmConfig"], tuple[int, ...]] + + refuse_config: Callable[["VllmConfig"], list[str]] + refuse_model: Callable[[Any], list[str]] + skip_layer: Callable[[MonoContext, Any], str | None] + skip_step: Callable[[MonoContext, MonoStep], str | None] + + pack: Callable[[Any, MonoContext], Any] + repack: Callable[[Any, MonoContext, Any], None] | None + bind: Callable[[MonoContext, Mapping[str, Any]], list[str]] + build: Callable[[MonoContext, Any, Any, int], LayerKernel] From 53f595ae83c128967cb12519b2505e4b2c3eefd3 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 08:25:51 -0500 Subject: [PATCH 02/79] feat: port dsv41 Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 7 +- tests/models/common/mono/reference.py | 6 +- tests/models/test_dsv41_mono_decode.py | 729 +++++++++++------- tests/models/test_dsv41_mono_layers.py | 20 +- vllm/models/common/mono/__init__.py | 17 +- vllm/models/common/mono/guards.py | 3 + vllm/models/common/mono/layer.py | 32 +- vllm/models/common/mono/spec.py | 5 +- vllm/models/deepseek_v41/amd/model.py | 38 +- vllm/models/deepseek_v41/amd/mono/__init__.py | 2 +- vllm/models/deepseek_v41/amd/mono_decode.py | 435 ----------- vllm/models/deepseek_v41/amd/mono_spec.py | 517 +++++++++++++ 12 files changed, 1055 insertions(+), 756 deletions(-) delete mode 100644 vllm/models/deepseek_v41/amd/mono_decode.py create mode 100644 vllm/models/deepseek_v41/amd/mono_spec.py diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 758fe2319fa8..025473b1ff11 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -37,7 +37,7 @@ set_capture_sizes, ) from vllm.models.common.mono.layer import attach, forward_layer -from vllm.models.common.mono.spec import MonoUnsupported +from vllm.models.common.mono.spec import Feature, MonoUnsupported KERNEL_BACKENDS = ("aiter", "flydsl") """Backends a spec module must not pull in: it is imported at config time, on @@ -127,9 +127,10 @@ def test_refusals_report_every_reason(self, config): config.use_v2_model_runner = False for spec in load_specs(self.specs_path): reasons = refuse_config(spec, config) - assert sum("lora" in r for r in reasons) == 1, reasons - assert sum("speculative_decoding" in r for r in reasons) == 1, reasons assert sum("Model Runner V2" in r for r in reasons) == 1, reasons + for feature in (Feature.LORA, Feature.SPECULATIVE_DECODING): + want = 0 if feature in spec.supports else 1 + assert sum(feature.value in r for r in reasons) == want, reasons def test_refusal_names_every_variant(self, config): config.use_v2_model_runner = False diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index ef34117ba546..3de113b302b7 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -179,8 +179,10 @@ def _tp(): REFERENCE_WIDE = replace( REFERENCE, name="reference-wide", - refuse_config=lambda c: refuse_config(c) - + (["needs max_num_seqs >= 8"] if c.scheduler_config.max_num_seqs < 8 else []), + refuse_config=lambda c: ( + refuse_config(c) + + (["needs max_num_seqs >= 8"] if c.scheduler_config.max_num_seqs < 8 else []) + ), ) MONO_SPECS = (REFERENCE_WIDE, REFERENCE) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 309d278bcb42..30df8802d59e 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -1,233 +1,309 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The DeepSeek-V4.1 mono decode layer runs only where its kernels are exact: -decode-only steps of at most MAX_ROWS rows with causal SWA windows, eager or in -a FULL graph. A PIECEWISE capture is replayed for mixed batches, so it must -never record the mono launches. Layers outside the whole-layer kernels run the -FFN launch on the same steps, from wo_b's unreduced output.""" +"""DeepSeek-V4.1's MonoSpec against the conformance kit, and the choices that +are its own: which layers take the whole-layer launch and which take the FFN +launch, and which steps the kernels are exact on (decode only, at most MAX_ROWS +rows, causal SWA windows, eager or in a FULL graph). +The kernels need CDNA4 and FlyDSL; these tests stub the runner module, so they +check the spec's decisions on CPU and never launch anything. +""" + +import importlib import sys from types import ModuleType, SimpleNamespace +from typing import Any import pytest import torch -from vllm.config import CUDAGraphMode +from tests.models.common.mono.kit import MonoConformance, forward +from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -from vllm.models.deepseek_v41.amd import mono_decode as md +from vllm.models.common.mono.layer import attach, kernel_for +from vllm.models.common.mono.spec import MonoUnsupported +from vllm.models.deepseek_v41.amd import mono_spec as ms -M = 6 +SPECS_PATH = ms.SPECS_PATH +LAYERS = 6 +HIDDEN = 16 +BLOCKS = 256 -class _Runner: - def __init__(self): - self.calls = [] +class _StubRunner: + def __init__(self, world, rank, group, device): + self.world = world + self.built: list = [] + self.epoch = torch.zeros(1, dtype=torch.int32) - def forward(self, *args, **kwargs): - self.calls.append((args, kwargs)) - return ("mono outputs",) + def scratch(self, width): + pass + + def kernels(self, *args, **kwargs): + self.built.append((args, kwargs)) + + def ffn_kernel(self, width): + self.built.append(width) def first_pre_mix(self, rows): return torch.zeros(rows, 4) + def forward(self, *args, **kwargs): + return ("whole", args, kwargs) + def ffn(self, *args): - self.calls.append((args, {})) - return ("ffn outputs",) + return ("ffn", args) + + +def _stub_modules(monkeypatch) -> None: + """The FlyDSL runner package, as a spec-shaped stub.""" + runner: Any = ModuleType("vllm.models.deepseek_v41.amd.mono.runner") + runner.BLOCKS = BLOCKS + runner.MAX_TOKENS = 64 + runner.DSV41MonoLayer = _StubRunner + runner.MonoLayerWeights = _weights_stub + runner.AttnWeights = _weights_stub + plan: Any = ModuleType("vllm.models.deepseek_v41.amd.mono.attention.plan") + plan.Dims = lambda world: world + for name, module in ( + ("vllm.models.deepseek_v41.amd.mono.runner", runner), + ("vllm.models.deepseek_v41.amd.mono.attention.plan", plan), + ): + monkeypatch.setitem(sys.modules, name, module) + + +def _weights_stub(**fields) -> Any: + return SimpleNamespace(check=lambda *args: None, **fields) + + +def _bytes(*shape) -> torch.Tensor: + return torch.zeros(*shape, dtype=torch.uint8) + + +def _proj(out=HIDDEN, inp=HIDDEN) -> Any: + return SimpleNamespace( + weight=torch.nn.Parameter(torch.zeros(out, inp)), + weight_scale=_bytes(out, 4), + reduce_results=True, + ) -def _layer(comp="comp", indexer=None): - cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8) # noqa: E731 - attn = SimpleNamespace( - swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()), - compressed_cache_prefix=comp, - topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32), +def _norm() -> Any: + return SimpleNamespace(weight=torch.nn.Parameter(torch.ones(HIDDEN))) + + +def _moe(**overrides) -> Any: + experts = SimpleNamespace( + w13_weight=_bytes(4, HIDDEN), + w13_weight_scale=_bytes(4, 4), + w2_weight=_bytes(4, HIDDEN), + w2_weight_scale=_bytes(4, 4), + quant_method=SimpleNamespace(mxfp4_backend=Mxfp4MoeBackend.AITER_MXFP4_BF16), + ) + fields = dict( + experts=SimpleNamespace(routed_experts=experts), + shared_experts=SimpleNamespace(gate_up_proj=_proj(), down_proj=_proj()), + gate=SimpleNamespace( + weight=_proj().weight, e_score_correction_bias=None, tid2eid=None + ), + n_routed_experts=384, + n_activated_experts=6, + routed_scaling_factor=1.5, + swiglu_limit=10.0, + scoring_func="sqrtsoftplus", + renormalize=True, + ) + return SimpleNamespace(**(fields | overrides)) + + +def _attn(index: int, **overrides) -> Any: + fields = dict( + layer_id=index, + compressor=None, + indexer=None, compress_ratio=1, - indexer=indexer, - _compressed_kv_cache=cache, + kv_mxfp8=False, + kv_cache_dtype="fp8_ds_mla", + swa_cache_layer=SimpleNamespace(prefix=f"layers.{index}.swa", kv_cache=None), + compressed_cache_prefix=f"layers.{index}.comp", + topk_indices_buffer=torch.zeros(64, 8, dtype=torch.int32), + fused_wqa_wkv=_proj(), + q_norm=_norm(), + kv_norm=_norm(), + wq_b=_proj(), + wo_a=_proj(), + wo_b=_proj(), + attn_sink=None, + rotary_emb=SimpleNamespace(cos_sin_cache=torch.zeros(8, 8)), + _compressed_kv_cache=lambda: make_cache(), ) - return SimpleNamespace(attn=attn, engram=None, hc_mult=4) + return SimpleNamespace(**(fields | overrides)) + + +class _Layer(torch.nn.Module): + def __init__(self, index: int, **overrides): + super().__init__() + self.attn = _attn(index, **overrides.pop("attn", {})) + self.ffn = _moe(**overrides.pop("ffn", {})) + self.engram = None + self.use_sequence_parallel = False + self.fuse_seam_norm = True + self.hc_mult = 4 + self.hc_attn_fn = self.hc_ffn_fn = None + self.hc_attn_scale = self.hc_ffn_scale = 1.0 + self.hc_attn_base = self.hc_ffn_base = 0.0 + self.attn_norm = _norm() + self.ffn_norm = _norm() + self.mhc_post = lambda x, residual, post, res: residual + for key, value in overrides.items(): + setattr(self, key, value) + self.weight = torch.nn.Parameter(torch.zeros(HIDDEN)) + + +class _Model(torch.nn.Module): + def __init__(self, config): + super().__init__() + self.config = config.model_config.hf_config + self.vllm_config = config + self.layers = torch.nn.ModuleList(_Layer(i) for i in range(LAYERS)) + self.start_layer = 0 + self.end_layer = LAYERS + self.mono = None + + +def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: + return _bytes(blocks, block * ms.RECORD) + + +def make_config(**overrides) -> Any: + config = SimpleNamespace( + kernel_config=SimpleNamespace(monokernel=True), + model_config=SimpleNamespace( + enable_sleep_mode=False, + hf_config=SimpleNamespace( + num_hidden_layers=LAYERS, hidden_size=HIDDEN, index_topk=8 + ), + ), + scheduler_config=SimpleNamespace(max_num_seqs=8), + parallel_config=SimpleNamespace( + tensor_parallel_size=2, + pipeline_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=False, + enable_eplb=False, + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + offload_config=SimpleNamespace( + uva=SimpleNamespace(cpu_offload_gb=0), + prefetch=SimpleNamespace(offload_group_size=0), + ), + aux_output_config=SimpleNamespace(enable_return_routed_experts=False), + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.FULL, + cudagraph_capture_sizes=list(range(1, ms.MAX_ROWS + 1)), + static_forward_context={}, + ), + lora_config=None, + kv_transfer_config=None, + speculative_config=None, + weight_transfer_config=None, + use_v2_model_runner=True, + ) + for key, value in overrides.items(): + setattr(config, key, value) + return config -def _swa(**overrides): +def _swa(**overrides) -> Any: fields = dict( num_prefills=0, num_decodes=1, - num_decode_tokens=M, - decode_swa_width=md.SWA_WIDTH, - decode_swa_indices=torch.zeros(64, 1, md.SWA_WIDTH, dtype=torch.int32), - decode_swa_lens=torch.zeros(64, dtype=torch.int32), - token_to_req_indices=torch.zeros(64, dtype=torch.int32), - slot_mapping=torch.zeros(64, dtype=torch.int64), + num_decode_tokens=1, + decode_swa_width=ms.SWA_WIDTH, + decode_swa_indices=torch.zeros(8, 1, ms.SWA_WIDTH, dtype=torch.int32), + decode_swa_lens=torch.zeros(8, dtype=torch.int32), + token_to_req_indices=torch.zeros(8, dtype=torch.int32), + slot_mapping=torch.zeros(8, dtype=torch.int64), block_size=32, ) - fields.update(overrides) - return SimpleNamespace(**fields) - - -def _inputs(rows=M): - return dict( - x=torch.zeros(rows, 5120, dtype=torch.bfloat16), - positions=torch.zeros(rows, dtype=torch.int64), - residual=torch.zeros(rows, 4, 5120, dtype=torch.bfloat16), - post_mix=torch.zeros(rows, 4, 1), - res_mix=torch.zeros(rows, 4, 4), - pre_mix=torch.zeros(rows, 4), - ) + return SimpleNamespace(**(fields | overrides)) -@pytest.fixture -def run(monkeypatch): - """Call the mono path for one layer under a forward context built from - ``mode`` and ``metadata``; returns (result, runner calls).""" +def _metadata(model, **overrides) -> dict: + meta: dict = {} + for layer in model.layers: + meta[layer.attn.swa_cache_layer.prefix] = _swa(**overrides) + meta[layer.attn.compressed_cache_prefix] = SimpleNamespace( + block_size=32, block_table=torch.zeros(8, 4, dtype=torch.int32) + ) + return meta - def _run( - mode=CUDAGraphMode.FULL, - metadata=None, - ffn_only=False, - entry=None, - layer=None, - **inputs, - ): - if metadata is None: - metadata = { - "swa": _swa(), - "comp": SimpleNamespace(block_size=128, block_table=None), - } - fc = SimpleNamespace(cudagraph_runtime_mode=mode, attn_metadata=metadata) - runner = _Runner() - monkeypatch.setattr(md, "is_forward_context_available", lambda: True) - monkeypatch.setattr(md, "get_forward_context", lambda: fc) - monkeypatch.setattr(md, "_mono_runner", lambda device: runner) - monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) - args = {**_inputs(), **inputs} - mono = md.MonoDecodeLayer(ffn_only) - if (entry or ("ffn" if ffn_only else "layer")) == "ffn": - # x stands for wo_b's unreduced output; the first layer's attention - # has no compressed cache - part, _ = args.pop("x"), args.pop("positions") - return mono.ffn(_layer(comp=None), part, **args), runner.calls - return mono(layer or _layer(), **args), runner.calls - - return _run - - -@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) -def test_decode_step_runs_mono(run, mode): - out, calls = run(mode) - assert out == ("mono outputs",) and len(calls) == 1 - args, kwargs = calls[0] - # the caches reach the kernels as [blocks, block, record] bytes - assert args[8].shape == (4, 32, md.RECORD) and args[8].stride() == ( - 128 * md.RECORD, - md.RECORD, - 1, - ) - assert kwargs["comp_cache"].shape == (4, 128, md.RECORD) - assert kwargs["fold"] and "select" not in kwargs +class TestDsv41Spec(MonoConformance): + specs_path = SPECS_PATH -def test_first_layer_runs_mono_from_the_embedding(run): - """The first layer's residual is the embedding broadcast to the four - streams, which K1 takes as it is (no owed post to fold).""" - out, calls = run(residual=None, post_mix=None, res_mix=None, pre_mix=None) - assert out == ("mono outputs",) and len(calls) == 1 - args, kwargs = calls[0] - assert args[2].shape == (M, 4, 5120) and not kwargs["fold"] + @pytest.fixture(autouse=True) + def _stubs(self, monkeypatch): + _gpu(monkeypatch) + def make_config(self, **overrides): + return make_config(**overrides) -def test_index_source_selects_between_the_kernels(run): - """An index source hands the runner its compressor and indexer to run - between K1 and K2.""" - out, calls = run(layer=_layer(indexer=object())) - assert out == ("mono outputs",) and callable(calls[0][1]["select"]) + def make_model(self, config): + model = _Model(config) + sfc = config.compilation_config.static_forward_context + for layer in model.layers: + layer.attn.swa_cache_layer.kv_cache = make_cache() + sfc[layer.attn.swa_cache_layer.prefix] = layer.attn.swa_cache_layer + sfc[layer.attn.compressed_cache_prefix] = SimpleNamespace( + kv_cache=make_cache() + ) + return model + def decode_metadata(self, config, model): + return _metadata(model) -def test_padded_decode_step_runs_mono(run): - """A FULL graph pads the step to its capture size: the kernels take it and - skip the pad rows (their slot is -1).""" - swa = _swa(num_decode_tokens=M - 2) - comp = SimpleNamespace(block_size=128, block_table=None) - out, calls = run(metadata={"swa": swa, "comp": comp}) - assert out == ("mono outputs",) and len(calls) == 1 + def skipped_metadata(self, config, model): + return _metadata(model, num_prefills=1) + def bad_caches(self, config, model): + return {"swa": torch.zeros(4, dtype=torch.bfloat16)} -@pytest.mark.parametrize( - "case", - [ - dict(mode=CUDAGraphMode.PIECEWISE), - dict( - metadata={ - "swa": _swa(num_prefills=1), - "comp": SimpleNamespace(block_size=128), - } - ), - dict( - metadata={ - "swa": _swa(decode_swa_width=256), - "comp": SimpleNamespace(block_size=128), - } - ), - dict( - metadata={ - "swa": _swa(num_decode_tokens=M + 1), - "comp": SimpleNamespace(block_size=128), - } - ), - dict(metadata={}), # a profile / dummy run: no attention metadata - dict(**_inputs(md.MAX_ROWS + 6)), - # a later pipeline stage's first layer: its stream arrives without - # the mixes - dict(residual=None, x=torch.zeros(M, 4, 5120, dtype=torch.bfloat16)), - ], - ids=[ - "piecewise", - "prefill", - "noncausal-window", - "tokens-past-rows", - "no-metadata", - "rows", - "stage-input", - ], -) -def test_other_steps_take_the_original_path(run, case): - out, calls = run(**case) - assert out is None and not calls - - -def _decoder_layer(engram=None, **attn): - attn = ( - dict( - layer_id=21, - compressor=None, - indexer=None, - compress_ratio=1, - kv_mxfp8=False, - kv_cache_dtype="fp8_ds_mla", + def _run_layers(self, mono, model) -> None: + positions = self._positions(mono) + for i, layer in enumerate(model.layers): + kernel_for(mono, i, positions) + + +def _rocm_platform() -> ModuleType: + """The real ROCm platform module, whose arch probe needs a GPU to answer.""" + if "vllm.platforms.rocm" not in sys.modules: + probe = torch.cuda.get_device_properties + torch.cuda.get_device_properties = lambda *_, **__: SimpleNamespace( + gcnArchName="gfx950" ) - | attn - ) - backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 - quant_method = SimpleNamespace(mxfp4_backend=backend) - ffn = SimpleNamespace( - experts=SimpleNamespace( - routed_experts=SimpleNamespace(quant_method=quant_method) - ), - shared_experts=object(), - gate=SimpleNamespace(tid2eid=None), - n_routed_experts=384, - n_activated_experts=6, - routed_scaling_factor=1.5, - swiglu_limit=10.0, - scoring_func="sqrtsoftplus", - renormalize=True, - ) - return SimpleNamespace( - attn=SimpleNamespace(**attn), - ffn=ffn, - engram=engram, - use_sequence_parallel=False, - fuse_seam_norm=True, - ) + try: + importlib.import_module("vllm.platforms.rocm") + finally: + torch.cuda.get_device_properties = probe + return sys.modules["vllm.platforms.rocm"] + + +def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS) -> None: + _stub_modules(monkeypatch) + monkeypatch.setattr(ms, "_compute_units", lambda: cus) + monkeypatch.setattr(_rocm_platform(), "get_cdna_version", lambda: cdna) + + +def _attach(monkeypatch, config=None, model=None, **layer_overrides): + _gpu(monkeypatch) + config = config or make_config() + model = model if model is not None else TestDsv41Spec().make_model(config) + for layer in model.layers: + for key, value in layer_overrides.items(): + setattr(layer.attn, key, value) + return attach(model, config, SPECS_PATH), model def _indexer(): @@ -237,136 +313,223 @@ def _indexer(): @pytest.mark.parametrize( - "layer, path", + "overrides, ffn_only", [ - (_decoder_layer(), "whole"), - (_decoder_layer(compress_ratio=0), "whole"), - (_decoder_layer(engram=object()), "whole"), - (_decoder_layer(compressor=object(), indexer=_indexer()), "whole"), - (_decoder_layer(indexer=_indexer()), "whole"), - (_decoder_layer(compressor=object()), "ffn"), - (_decoder_layer(indexer=object()), "ffn"), - (_decoder_layer(kv_cache_dtype="auto"), "ffn"), - (_decoder_layer(layer_id=40), None), + (dict, False), + (lambda: dict(compress_ratio=0), False), + (lambda: dict(compressor=object(), indexer=_indexer()), False), + (lambda: dict(indexer=_indexer()), False), + (lambda: dict(compressor=object()), True), + (lambda: dict(indexer=object()), True), + (lambda: dict(kv_cache_dtype="auto"), True), + (lambda: dict(kv_mxfp8=True), True), + (lambda: dict(compress_ratio=3), True), ], ids=[ "standard", "window", - "engram", "kv-source", "index-source", "compressor-only", "other-indexer", "kv", - "draft", + "mxfp8", + "ratio", ], ) -def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): +def test_layers_take_whole_or_ffn_launch(monkeypatch, overrides, ffn_only): """The whole-layer kernels serve every backbone layer of the checkpoint (an index source's compressor and indexer run between them); a layer outside - them keeps vLLM's attention and runs its FFN half as one launch.""" - _deployment(monkeypatch) - mono = md.MonoDecodeLayer.create(layer, _config()) - got = None if mono is None else "ffn" if mono.ffn_only else "whole" - assert got == path + them keeps vLLM's attention and runs its FFN half as one launch, from + wo_b's unreduced output.""" + mono, model = _attach(monkeypatch, **overrides()) + kernels = {k.ffn_only for k in mono.kernels.values()} + assert kernels == {ffn_only} + assert all(layer.attn.wo_b.reduce_results is not ffn_only for layer in model.layers) + + +@pytest.mark.parametrize( + "skip", + [ + dict(use_sequence_parallel=True), + dict(fuse_seam_norm=False), + ], + ids=["sequence-parallel", "unfused-seam"], +) +def test_skipped_layers_keep_vllm_path(monkeypatch, skip): + config = make_config() + model = TestDsv41Spec().make_model(config) + for layer in model.layers: + for key, value in skip.items(): + setattr(layer, key, value) + mono, _ = _attach(monkeypatch, config=config, model=model) + assert mono is None -@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) -def test_decode_step_runs_ffn_launch(run, mode): - out, calls = run(mode, ffn_only=True, metadata={"swa": _swa()}) - assert out == ("ffn outputs",) and len(calls) == 1 - assert calls[0][0][1].shape == (M, 5120) # wo_b's unreduced output +def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): + config = make_config() + model = TestDsv41Spec().make_model(config) + for layer in model.layers: + layer.ffn.n_routed_experts = 256 + mono, _ = _attach(monkeypatch, config=config, model=model) + assert mono is None -def test_ffn_layer_never_runs_whole(run): - out, calls = run(ffn_only=True, entry="layer") - assert out is None and not calls +def _step(mono, model, metadata, width=1): + with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, metadata): + return mono.begin_step(torch.zeros(width, dtype=torch.long)) @pytest.mark.parametrize( - "case", + "case, why", [ - dict(mode=CUDAGraphMode.PIECEWISE), - dict(metadata={"swa": _swa(num_prefills=1)}), - dict(metadata={"swa": _swa(num_decode_tokens=M + 1)}), - dict(metadata={}), - dict(**_inputs(md.MAX_ROWS + 6)), + (dict(num_prefills=1), "not a decode-only step"), + (dict(num_decodes=0), "not a decode-only step"), + (dict(num_decode_tokens=4), "more decode tokens than rows"), + (dict(decode_swa_width=256), "SWA width"), + (dict(decode_swa_lens=None), "SWA decode metadata missing"), ], - ids=["piecewise", "prefill", "tokens-past-rows", "no-metadata", "rows"], + ids=["prefill", "no-decode", "tokens-past-rows", "noncausal-window", "no-swa"], ) -def test_other_steps_reduce_wo_b_themselves(run, case): - """None: the layer all-reduces wo_b's partial and runs the FFN as before.""" - case.setdefault("metadata", {"swa": _swa()}) - out, calls = run(ffn_only=True, **case) - assert out is None and not calls +def test_other_steps_take_vllm_path(monkeypatch, case, why): + mono, model = _attach(monkeypatch) + assert why in _step(mono, model, _metadata(model, **case)) -def _config(monokernel=True, **parallel): - parallel = ( - dict(enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1) - | parallel - ) - return SimpleNamespace( - model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), - parallel_config=SimpleNamespace(**parallel), - kernel_config=SimpleNamespace(monokernel=monokernel), - ) +def test_decode_step_takes_the_kernels(monkeypatch): + mono, model = _attach(monkeypatch) + assert _step(mono, model, _metadata(model)) is None -def _deployment(monkeypatch, cdna=4, cus=256): - """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" - import vllm.platforms.rocm as rocm +def test_profile_run_takes_vllm_path(monkeypatch): + mono, model = _attach(monkeypatch) + assert _step(mono, model, None) == "no attention metadata" - monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) - monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna) - monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2) - monkeypatch.setattr(md, "_compute_units", lambda: cus) - name = "vllm.models.deepseek_v41.amd.mono.runner" - runner = ModuleType(name) - runner.BLOCKS = 256 # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, name, runner) +def test_rows_past_the_kernels_take_vllm_path(monkeypatch): + mono, model = _attach(monkeypatch) + assert _step(mono, model, _metadata(model), width=ms.MAX_ROWS + 1) == "width" -def test_off_without_the_flag(monkeypatch): - """No --monokernel: no mono layer, whatever the deployment.""" - _deployment(monkeypatch) - assert ( - md.MonoDecodeLayer.create(_decoder_layer(), _config(monokernel=False)) is None - ) + +def test_widths_follow_the_capture_sizes(monkeypatch): + config = make_config() + config.compilation_config.cudagraph_capture_sizes = [1, 2, 8, 64, 512] + assert ms._widths(config) == (1, 2, 8) + + +def test_widths_default_to_every_row(monkeypatch): + config = make_config() + config.compilation_config.cudagraph_capture_sizes = [512] + assert ms._widths(config) == tuple(range(1, ms.MAX_ROWS + 1)) + + +@pytest.mark.parametrize("tp", [1, 3, 8]) +def test_tensor_parallel_outside_the_kernels_refuses(tp): + config = make_config() + config.parallel_config.tensor_parallel_size = tp + assert any("tensor parallel" in r for r in ms._refuse_config(config)) + + +def test_a4w4_moe_refuses(monkeypatch): + monkeypatch.setattr("vllm.envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", True) + assert any("A8W4" in r for r in ms._refuse_config(make_config())) @pytest.mark.parametrize("cdna", [3, 5]) -def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): - _deployment(monkeypatch, cdna=cdna) - with pytest.raises(ValueError, match="needs CDNA4"): - md.MonoDecodeLayer.create(_decoder_layer(), _config()) +def test_other_isa_refuses(monkeypatch, cdna): + config = make_config() + model = TestDsv41Spec().make_model(config) + _gpu(monkeypatch, cdna=cdna) + with pytest.raises(MonoUnsupported, match="CDNA4"): + attach(model, config, SPECS_PATH) -def _other_backend(): - layer = _decoder_layer() - experts = layer.ffn.experts.routed_experts - experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON - return layer +def test_too_few_compute_units_refuses(monkeypatch): + config = make_config() + model = TestDsv41Spec().make_model(config) + _gpu(monkeypatch, cus=BLOCKS // 2) + with pytest.raises(MonoUnsupported, match="compute units"): + attach(model, config, SPECS_PATH) -@pytest.mark.parametrize( - "case, match", - [ - (dict(parallel=dict(enable_expert_parallel=True)), "no expert"), - (dict(parallel=dict(data_parallel_size=2)), "no expert"), - (dict(parallel=dict(enable_eplb=True)), "no expert"), - (dict(a4w4=True), "A8W4"), - (dict(layer=_other_backend), "AITER_MXFP4_BF16"), - (dict(cus=128), "256 compute units"), - ], - ids=["expert-parallel", "data-parallel", "eplb", "a4w4", "moe-backend", "cus"], -) -def test_opt_in_outside_the_kernels_moe_raises(monkeypatch, case, match): - """The kernels read all 384 experts on each rank in AITER's A8W4 layout, with - all 256 CTAs resident: any other deployment is an error, not garbage.""" - _deployment(monkeypatch, cus=case.get("cus", 256)) - if case.get("a4w4"): - monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1") - layer = case.get("layer", _decoder_layer)() - with pytest.raises(ValueError, match=match): - md.MonoDecodeLayer.create(layer, _config(**case.get("parallel", {}))) +def test_other_moe_backend_refuses(monkeypatch): + config = make_config() + model = TestDsv41Spec().make_model(config) + for layer in model.layers: + experts = layer.ffn.experts.routed_experts + experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON + _gpu(monkeypatch) + with pytest.raises(MonoUnsupported, match="AITER_MXFP4_BF16"): + attach(model, config, SPECS_PATH) + + +def test_off_without_the_flag(monkeypatch): + config = make_config() + config.kernel_config.monokernel = False + mono, _ = _attach(monkeypatch, config=config) + assert mono is None + + +def test_whole_launch_runs_the_layer(monkeypatch): + mono, model = _attach(monkeypatch) + positions = torch.zeros(1, dtype=torch.long) + with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = kernel_for(mono, 1, positions) + out = kernel.launch( + mono.step, + torch.zeros(1, HIDDEN), + positions, + torch.zeros(1, 4, HIDDEN), + torch.zeros(1, 4, 1), + torch.zeros(1, 4, 4), + torch.zeros(1, 4), + ) + assert out[0] == "whole" + records = [a for a in out[1] if torch.is_tensor(a) and a.shape[-1:] == (ms.RECORD,)] + assert len(records) == 1 + assert out[2]["comp_cache"].shape[-1] == ms.RECORD + assert out[2]["fold"] + + +def test_first_layer_takes_the_embedding(monkeypatch): + """Layer 0's residual is the embedding broadcast to the four streams, which + K1 takes as it is (no owed post to fold).""" + mono, model = _attach(monkeypatch) + positions = torch.zeros(1, dtype=torch.long) + with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = kernel_for(mono, 0, positions) + out = kernel.launch( + mono.step, torch.zeros(1, HIDDEN), positions, None, None, None, None + ) + assert out[0] == "whole" and not out[2]["fold"] + assert out[1][2].shape == (1, 4, HIDDEN) + + +def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): + mono, model = _attach(monkeypatch, compressor=object()) + positions = torch.zeros(1, dtype=torch.long) + with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = kernel_for(mono, 1, positions) + assert kernel.ffn_only + out = kernel.launch( + mono.step, + torch.zeros(1, HIDDEN), + torch.zeros(1, 4, HIDDEN), + torch.zeros(1, 4, 1), + torch.zeros(1, 4, 4), + torch.zeros(1, 4), + ) + assert out[0] == "ffn" + + +def test_launch_declines_a_layout_it_cannot_read(monkeypatch): + """A later pipeline stage's first layer gets a stream without the mixes: + the step descriptor cannot see that, so the launch declines it.""" + mono, model = _attach(monkeypatch) + positions = torch.zeros(1, dtype=torch.long) + with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = kernel_for(mono, 1, positions) + out = kernel.launch( + mono.step, torch.zeros(1, 4, HIDDEN), positions, None, None, None, None + ) + assert out is None diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index a924cfda50c3..ddca6ba1c3dc 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -209,15 +209,21 @@ def save_capture(self, path: str) -> None: torch.save((self._captured, self._gaps), f"{path}.rank{rank}") def mono_layers(self) -> dict: - from vllm.models.deepseek_v41.amd import mono_decode - + model = self.model_runner.model # type: ignore[attr-defined] + mono = getattr(model.model, "mono", None) + kernels = {} if mono is None else mono.kernels layers = { - i: None - if m.mono is None - else ["ffn" if m.mono.ffn_only else "whole", m.mono._weights is not None] - for i, m in self._decoder_layers().items() + i: next( + ( + ["ffn" if k.ffn_only else "whole", k.weights is not None] + for (j, _), k in kernels.items() + if j == i + ), + None, + ) + for i in self._decoder_layers() } - runner = mono_decode._runner + runner = None if mono is None else mono.ctx.state.runner return { "layers": layers, "epoch": -1 if runner is None else int(runner.epoch[0]), diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index 524fe5008c19..abbeab6ebeb3 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -8,6 +8,19 @@ error-word layout (``abi``).""" from vllm.models.common.mono.guards import configure -from vllm.models.common.mono.layer import account_step, attach, forward_layer, health +from vllm.models.common.mono.layer import ( + account_step, + attach, + forward_layer, + health, + kernel_for, +) -__all__ = ["account_step", "attach", "configure", "forward_layer", "health"] +__all__ = [ + "account_step", + "attach", + "configure", + "forward_layer", + "health", + "kernel_for", +] diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index adc4ac8b694f..d1b3ee2e80e7 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -140,6 +140,7 @@ def set_capture_sizes(spec: MonoSpec, vllm_config: "VllmConfig") -> None: def configure(vllm_config: "VllmConfig", specs_path: str) -> None: """Picks the variant, sets the capture sizes, or raises with every reason. + A no-op with ``--monokernel`` off. Args: vllm_config: The engine's config, before the weights are loaded. @@ -149,6 +150,8 @@ def configure(vllm_config: "VllmConfig", specs_path: str) -> None: MonoUnsupported: Every variant refused this config. """ + if not enabled(vllm_config): + return spec, refusals = select(load_specs(specs_path), vllm_config) if spec is None: raise MonoUnsupported("MonoKernel refused: " + refused(refusals)) diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index 22fc5acd66ce..a15b6ccdfcc0 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -12,7 +12,7 @@ from vllm.models.common.mono import failstop from vllm.models.common.mono.guards import enabled, load_specs, refused, select from vllm.models.common.mono.routing import Mono -from vllm.models.common.mono.spec import MonoUnsupported +from vllm.models.common.mono.spec import LayerKernel, MonoUnsupported if TYPE_CHECKING: from vllm.config import VllmConfig @@ -60,15 +60,31 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: return mono +def kernel_for(mono: Mono | None, i: int, positions) -> LayerKernel | None: + """Layer ``i``'s kernel when this step runs on the kernels, else None. + + The first call under a forward context decides for every layer. A model + whose layer signature or dispatch site does not fit ``forward_layer`` looks + the kernel up itself and calls ``launch`` where it suits the layer. + """ + if mono is None: + return None + mono.begin_once(positions) + step = mono.step + if step is None: + return None + return mono.kernels.get((i, step.width)) + + def forward_layer(mono: Mono | None, i: int, layer, positions, *args, **kwargs): """Layer ``i`` on its kernel when this step runs on the kernels, else the - vLLM layer unchanged. The first call under a forward context decides for - every layer.""" - if mono is not None: - mono.begin_once(positions) - step = mono.step - if step is not None and (k := mono.kernels.get((i, step.width))) is not None: - return k.launch(step, positions, *args, **kwargs) + vLLM layer unchanged.""" + k = kernel_for(mono, i, positions) + step = None if mono is None else mono.step + if k is not None and step is not None: + out = k.launch(step, positions, *args, **kwargs) + if out is not None: + return out return layer(positions, *args, **kwargs) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index a61bc724db72..4b05a4954fd0 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -94,7 +94,10 @@ class LayerKernel(Protocol): """Attention layers whose caches ``MonoSpec.bind`` checks.""" def launch(self, step: MonoStep, *args, **kwargs) -> Any: - """Takes and returns what the vLLM layer it replaces does.""" + """Takes and returns what the vLLM layer it replaces does, or None to + decline this call, which sends the layer down vLLM's path. Decline only + on what the step descriptor cannot answer, such as the layout of the + tensors the layer was handed.""" ... def health(self) -> torch.Tensor: diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index a52a408c74dd..79fd3d906faa 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -50,6 +50,8 @@ make_layers, maybe_prefix, ) +from vllm.models.common.mono import attach, configure, health, kernel_for +from vllm.models.common.mono.routing import Mono from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -59,7 +61,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer +from vllm.models.deepseek_v41.amd.mono_spec import SPECS_PATH, Kernel from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers @@ -258,12 +260,6 @@ def __init__( # attn_norm / ffn_norm into its collapse, so the separate norms are # skipped for the seams it takes. self.fuse_seam_norm = HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM - # Decode steps of an eligible layer may run as the mono decode layer. - self.mono = MonoDecodeLayer.create(self, vllm_config) - if self.mono is not None and self.mono.ffn_only: - # A decode step hands wo_b's TP partial to the FFN launch; forward - # reduces it on the others. - self.attn.wo_b.reduce_results = False @staticmethod def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: @@ -281,10 +277,15 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, + mono: Mono | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - if self.mono is not None: - out = self.mono( - self, + kernel = typing.cast( + Kernel | None, kernel_for(mono, self.attn.layer_id, positions) + ) + step = None if mono is None else mono.step + if kernel is not None and step is not None and not kernel.ffn_only: + out = kernel.launch( + step, x, positions, residual, @@ -394,10 +395,11 @@ def forward( x = self.attn(positions, x, None) if self.use_sequence_parallel: x = sp_reduce_scatter(x) - elif self.mono is not None and self.mono.ffn_only: - out = self.mono.ffn(self, x, residual, post_mix, res_mix, attn_pre) - if out is not None: - return out + elif not self.attn.wo_b.reduce_results: + if kernel is not None and step is not None: + out = kernel.launch(step, x, residual, post_mix, res_mix, attn_pre) + if out is not None: + return out x = tensor_model_parallel_all_reduce(x) residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( @@ -481,6 +483,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.hc_mult = config.hc_mult self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps + self.vllm_config = vllm_config + self.mono: Mono | None = None + configure(vllm_config, SPECS_PATH) # Three aux streams: one per non-default input GEMM in # DeepseekV4Attention._run_parallel_input_projections @@ -843,6 +848,7 @@ def _run_layers( residual, engram_hashes, engram_mask, + self.mono, ) # The aux id of the first replay layer is the post of the last # full-batch layer. Capture it on the trimmed rows instead. @@ -1411,6 +1417,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() + self.model.mono = attach(self.model, self.model.vllm_config, SPECS_PATH) + + def step_health(self) -> torch.Tensor | None: + return health(self.model.mono) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py index 72a6ec2ab196..e5ba197cfb79 100644 --- a/vllm/models/deepseek_v41/amd/mono/__init__.py +++ b/vllm/models/deepseek_v41/amd/mono/__init__.py @@ -2,4 +2,4 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek-V4.1-Flash mono decode layer: one decoder layer of a decode step in two persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The -host side is ``runner``; vLLM docks it in ``..mono_decode``.""" +host side is ``runner``; vLLM docks it in ``..mono_spec``.""" diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py deleted file mode 100644 index ddeb6b26314f..000000000000 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ /dev/null @@ -1,435 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``--monokernel``). - -A decode step's backbone layer -- the attention seam, the attention, its TP -all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent -FlyDSL launches (``mono``, this directory) on the layer's loaded -weights, the TP reductions inside the kernels. Every backbone layer of the -checkpoint runs whole (compress ratio 0, 1 or 2 over ``fp8_ds_mla`` records): -the first layer and the Engram layers hand K1 a settled residual, and an index -source runs vLLM's compressor and indexer between the two launches. A layer -outside that (an MXFP8 KV cache, another indexer) keeps vLLM's attention seam -and attention and runs the rest as one launch, the FFN launch, from wo_b's -unreduced output. A step takes the kernels when it is decode only, at most -``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); -a whole layer also needs causal SWA windows. Anything else runs the layer as -before; the paths share every tensor at the layer boundary, so they interleave -freely. - -Every TP rank must take the same path at every layer (the kernels wait on each -other's pushes): the conditions read only metadata every rank holds alike. -""" - -from typing import TYPE_CHECKING, Any, cast - -import torch - -import vllm.envs as envs -from vllm.config import CUDAGraphMode -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - get_tp_group, -) -from vllm.forward_context import get_forward_context, is_forward_context_available -from vllm.logger import init_logger -from vllm.platforms import current_platform - -if TYPE_CHECKING: - from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer - from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights - from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata - from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata - -logger = init_logger(__name__) - -RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes -SWA_WIDTH = 128 # a causal decode token's window slots -MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) -# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats -CDNA_VERSIONS = (4,) - -_runner = None - - -def _mono_runner(device: torch.device): - """The process's mono runner (one a TP rank, shared by every layer). Its - peer-memory handshake is collective over the TP group: every rank builds it - at the same layer of the same step.""" - global _runner - if _runner is None: - if torch.cuda.is_current_stream_capturing(): - # its peer memory is allocated and exchanged eagerly; vLLM warms - # every graph up eagerly before capturing it - raise RuntimeError( - "--monokernel was first reached inside a CUDA graph " - "capture, before an eager decode step." - ) - from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer - - assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS - _runner = DSV41MonoLayer( - get_tensor_model_parallel_world_size(), - get_tensor_model_parallel_rank(), - get_tp_group().cpu_group, - device, - ) - return _runner - - -def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: - """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's - data rows, then its scale words: the kernels index both from it).""" - kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) - assert kv.stride(0) >= block * RECORD, (tuple(kv.shape), kv.stride(), block) - return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) - - -def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: - return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) - - -def _compute_units() -> int: - return current_platform.num_compute_units(torch.accelerator.current_device_index()) - - -def _unsupported_moe(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None: - """Why the kernels cannot run this deployment's MoE, or None. They read every - one of the 384 experts on each rank, TP-sharded, in AITER's A8W4 MXFP4 - layout.""" - from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend - - pc = vllm_config.parallel_config - if pc.enable_expert_parallel or pc.enable_eplb or pc.data_parallel_size > 1: - return "needs tensor parallelism only: no expert or data parallelism" - if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: - return "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" - experts = getattr(getattr(layer.ffn, "experts", None), "routed_experts", None) - if experts is None: - return None # not a routed MoE: the layer checks decline it - backend = getattr(experts.quant_method, "mxfp4_backend", None) - if backend != Mxfp4MoeBackend.AITER_MXFP4_BF16: - return f"needs the AITER_MXFP4_BF16 MoE backend, not {backend}" - return None - - -def _select( - layer: "DeepseekV4DecoderLayer", - positions: torch.Tensor, - normed: torch.Tensor, - qr: torch.Tensor, -) -> None: - """An index source's compressor and top-k selection, from K1's normed input - and bf16 qr: ``_forward_csa2_full`` / ``_forward_csa2_reindex`` and the - indexer op, serial.""" - attn = layer.attn - comp, indexer = attn.compressor, attn.indexer - assert indexer is not None - if comp is not None: - kv_score = torch.mm( - normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32 - ) - latent = comp(kv_score, positions) - comp.insert_cache(latent, positions, attn.rotary_emb) - indexer._produce_k(latent, positions, attn.indexer_rotary_emb) - weights, _ = indexer.weights_proj(normed) - q, q_scale, weights = indexer.forward_q( - qr, None, weights, positions, attn.indexer_rotary_emb - ) - assert weights is not None - indexer.indexer_op( - normed, (q, q_scale) if q_scale is not None else q, None, weights - ) - - -class MonoDecodeLayer: - """One decoder layer's mono path (``create``: None when the layer is not - eligible): the whole layer (``__call__``) or, with ``ffn_only``, the FFN - launch after vLLM's attention (``ffn``). Its weights are taken from the - layer at the first call, after loading.""" - - def __init__(self, ffn_only: bool = False) -> None: - self.ffn_only = ffn_only - self._weights: MonoLayerWeights | None = None - - @staticmethod - def create( - layer: "DeepseekV4DecoderLayer", vllm_config - ) -> "MonoDecodeLayer | None": - if not vllm_config.kernel_config.monokernel: - return None - from vllm.platforms.rocm import get_cdna_version - - # the deployment: an explicit opt-in that cannot run is an error - why = None - if get_cdna_version() not in CDNA_VERSIONS: - why = f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}" - elif get_tensor_model_parallel_world_size() not in (2, 4): - why = "needs tensor parallel size 2 or 4" - else: - why = _unsupported_moe(layer, vllm_config) - if why is None: - try: - from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS - except ImportError as err: - why = f"needs FlyDSL and AITER's FlyDSL helpers ({err})" - else: - # every CTA of a launch stays resident: a partitioned GPU deadlocks - cus = _compute_units() - if cus < BLOCKS: - why = f"needs {BLOCKS} compute units, the GPU has {cus}" - if why is not None: - raise ValueError(f"--monokernel {why}.") - # the layer: the kernels serve the backbone's seams and MoE - config = vllm_config.model_config.hf_config - attn, ffn = layer.attn, layer.ffn - if layer.use_sequence_parallel: - why = "sequence-parallel layer" - elif not layer.fuse_seam_norm: - why = "needs aiter's fused mHC seam" - elif attn.layer_id >= config.num_hidden_layers: - why = "draft layer" - elif ( - ffn.shared_experts is None - or ffn.gate.tid2eid is not None - or ffn.n_routed_experts != 384 - or ffn.n_activated_experts != 6 - or ffn.routed_scaling_factor != 1.5 - or ffn.swiglu_limit != 10.0 - or ffn.scoring_func != "sqrtsoftplus" - or not ffn.renormalize - ): - why = "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" - if why is not None: - logger.info_once("DSv4.1 mono decode off for some layers: %s", why) - return None - # the whole layer: the kernels' attention takes the standard layers; an - # index source's compressor and indexer run between its two launches - from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer - - if attn.compressor is not None and attn.indexer is None: - why = "a compressor without an indexer" - elif attn.indexer is not None and not isinstance( - attn.indexer, DeepseekV41RocmMxfp4Indexer - ): - why = f"a {type(attn.indexer).__name__}" - elif attn.compress_ratio not in (0, 1, 2): - why = f"compress ratio {attn.compress_ratio}" - elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": - why = f"a {attn.kv_cache_dtype} KV cache" - if why is not None: - logger.info_once( - "DSv4.1 mono decode: the FFN launch for layers with %s", why - ) - return MonoDecodeLayer(ffn_only=True) - logger.info_once("DSv4.1 mono decode: the whole standard backbone layers") - return MonoDecodeLayer() - - def weights(self, layer: "DeepseekV4DecoderLayer"): - if self._weights is not None: - return self._weights - from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights - - a, f = layer.attn, layer.ffn - e, sh = f.experts.routed_experts, f.shared_experts - assert sh is not None # ``create`` took only layers with one - self._weights = MonoLayerWeights( - attn=None if self.ffn_only else self._attn_weights(a), - hc_attn_fn=layer.hc_attn_fn, - hc_attn_scale=layer.hc_attn_scale, - hc_attn_base=layer.hc_attn_base, - attn_norm=layer.attn_norm.weight, - hc_ffn_fn=layer.hc_ffn_fn, - hc_ffn_scale=layer.hc_ffn_scale, - hc_ffn_base=layer.hc_ffn_base, - ffn_norm=layer.ffn_norm.weight, - gate_w=f.gate.weight, - bias=f.gate.e_score_correction_bias, - w13=e.w13_weight, - w13_s=e.w13_weight_scale, - w2=e.w2_weight, - w2_s=e.w2_weight_scale, - sgu=sh.gate_up_proj.weight, - sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), - sw2=sh.down_proj.weight, - sw2_s=_scale_bytes(sh.down_proj.weight_scale), - ) - self._weights.check(get_tensor_model_parallel_world_size()) - return self._weights - - @staticmethod - def _attn_weights(a): - from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims - from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights - - attn = AttnWeights( - layer_id=a.layer_id, - wqkv=a.fused_wqa_wkv.weight, - wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), - q_norm=a.q_norm.weight, - kv_norm=a.kv_norm.weight, - wq_b=a.wq_b.weight, - wq_b_scale=_scale_bytes(a.wq_b.weight_scale), - wo_a=a.wo_a.weight, - wo_a_scale=_scale_bytes(a.wo_a.weight_scale), - wo_b=a.wo_b.weight, - wo_b_scale=_scale_bytes(a.wo_b.weight_scale), - attn_sink=a.attn_sink, - cos_sin=a.rotary_emb.cos_sin_cache, - ratio=a.compress_ratio, - ) - attn.check(Dims(get_tensor_model_parallel_world_size())) - return attn - - @staticmethod - def _skip(why: str) -> None: - logger.debug_once("DSv4.1 mono decode skipped a step: %s", why) - return None - - @staticmethod - def _decode_step(layer: "DeepseekV4DecoderLayer", rows: int): - """(why this step takes the original path -- None for a decode-only step - the kernels take -- the layer's SWA metadata, its compressed cache's).""" - if not 1 <= rows <= MAX_ROWS: - return f"{rows} rows", None, None - if not is_forward_context_available(): - return "no forward context", None, None - fc = get_forward_context() - md = fc.attn_metadata - attn = layer.attn - swa, comp = None, None - if isinstance(md, dict): - swa = cast( - "DeepseekV4ROCMAiterSparseSWAMetadata | None", - md.get(attn.swa_cache_layer.prefix), - ) - if attn.compressed_cache_prefix is not None: - comp = cast( - "DeepseekV4FlashMLAMetadata | None", - md.get(attn.compressed_cache_prefix), - ) - why = None - if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): - # a PIECEWISE (breakable) capture is replayed for mixed batches too - why = f"{fc.cudagraph_runtime_mode} step" - elif not isinstance(md, dict): # a profile / dummy run, or DBO ubatches - why = "no attention metadata" - elif swa is None: - why = "layer metadata missing" - elif swa.num_prefills != 0 or swa.num_decodes == 0: - why = "not a decode-only step" - elif swa.num_decode_tokens > rows: - why = "more decode tokens than rows" - return why, swa, comp - - def __call__( - self, - layer: "DeepseekV4DecoderLayer", - x: torch.Tensor, - positions: torch.Tensor, - residual: torch.Tensor | None, - post_mix: torch.Tensor | None, - res_mix: torch.Tensor | None, - pre_mix: torch.Tensor | None, - engram_hashes: torch.Tensor | None = None, - engram_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, ...] | None: - """The layer's outputs (x, residual, post_mix, res_mix, pre_mix) by the - mono kernels, or None when this step takes the original path (always for - an ``ffn_only`` layer).""" - if self.ffn_only: - return None - first = residual is None - if first and x.dim() != 2: - return self._skip("first layer of a pipeline stage") - why, swa, comp = self._decode_step(layer, x.shape[0]) - attn = layer.attn - tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) - if why is None: - # the kernels' attention: causal SWA windows, the compressed cache - if comp is None and attn.compress_ratio: - why = "layer metadata missing" - elif swa.decode_swa_width != SWA_WIDTH: - why = f"SWA width {swa.decode_swa_width}" - elif ( - swa.decode_swa_indices is None - or swa.decode_swa_lens is None - or swa.token_to_req_indices is None - ): - why = "SWA decode metadata missing" - elif not all(t is not None and t.is_contiguous() for t in tensors): - why = "non-contiguous layer inputs" - if why is not None: - return self._skip(why) - assert swa is not None - runner = _mono_runner(x.device) - logger.info_once( - "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS - ) - # K1's seam folds the owed post into the residual, except where the - # residual arrives settled: the embedding, or after the Engram write - fold = True - if first: - residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() - pre_mix = runner.first_pre_mix(x.shape[0]) - post_mix = res_mix = pre_mix # unread without the fold - fold = False - elif layer.engram is not None and engram_hashes is not None: - residual = layer.mhc_post(x, residual, post_mix, res_mix) - hashes = engram_hashes[:, layer.engram.layer_hash_index] - residual = layer.engram(residual, hashes, engram_mask) - fold = False - kw: dict[str, Any] = {} - if attn.indexer is not None: - kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) - if attn.compress_ratio: - assert comp is not None - kw.update( - topk_indices=attn.topk_indices_buffer, - comp_cache=_record_view( - attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio - ), - comp_block_table=comp.block_table, - ) - return runner.forward( - self.weights(layer), - x, - residual, - post_mix, - res_mix, - pre_mix, - positions, - swa.slot_mapping, - _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), - swa.decode_swa_indices, - swa.decode_swa_lens, - swa.token_to_req_indices, - fold=fold, - **kw, - ) - - def ffn( - self, - layer: "DeepseekV4DecoderLayer", - part: torch.Tensor, - residual: torch.Tensor, - post_mix: torch.Tensor, - res_mix: torch.Tensor, - pre_mix: torch.Tensor, - ) -> tuple[torch.Tensor, ...] | None: - """The layer's outputs by the FFN launch from ``part``, wo_b's unreduced - output, and the attention seam's outputs; None when this step takes the - original path (the caller then reduces ``part`` itself).""" - why, _, _ = self._decode_step(layer, part.shape[0]) - tensors = (part, residual, post_mix, res_mix, pre_mix) - if why is None and not all(t.is_contiguous() for t in tensors): - why = "non-contiguous layer inputs" - if why is not None: - return self._skip(why) - runner = _mono_runner(part.device) - logger.info_once( - "DSv4.1 mono decode FFN launch active (decode steps of <= %d rows)", - MAX_ROWS, - ) - return runner.ffn(self.weights(layer), *tensors) diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py new file mode 100644 index 000000000000..9d45e01520be --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1's MonoKernels on ROCm CDNA4 (``--monokernel``, RFC #60904). + +A decode step's backbone layer (the attention seam, the attention, its TP +all-reduce, the FFN seam, the MoE and its all-reduce) runs as two persistent +FlyDSL launches (``mono``, this directory) on the layer's loaded weights, the +TP reductions inside the kernels. Every backbone layer of the checkpoint runs +whole (compress ratio 0, 1 or 2 over ``fp8_ds_mla`` records): the first layer +and the Engram layers hand K1 a settled residual, and an index source runs +vLLM's compressor and indexer between the two launches. A layer outside that +(an MXFP8 KV cache, another indexer) keeps vLLM's attention seam and attention +and runs the rest as one launch, the FFN launch, from wo_b's unreduced output. +A step takes the kernels when it is decode only, at a captured width of at most +``MAX_ROWS`` rows, with causal SWA windows. Anything else runs the layer as +before; the paths share every tensor at the layer boundary, so they interleave +freely. + +Every TP rank must take the same path at every layer (the kernels wait on each +other's pushes): the step decision reads only metadata every rank holds alike. + +This module is imported by ``mono.configure`` in the engine core process, +before the weights are loaded and with no device in hand, so it pulls in +neither FlyDSL nor AITER at import. +""" + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, cast + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.logger import init_logger +from vllm.models.common.mono.spec import ( + CacheReq, + Feature, + MonoContext, + MonoSpec, + MonoStep, +) +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + +logger = init_logger(__name__) + +SPECS_PATH = "vllm.models.deepseek_v41.amd.mono_spec:MONO_SPECS" + +RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes +SWA_WIDTH = 128 # a causal decode token's window slots +MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats +CDNA_VERSIONS = (4,) + + +def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: + """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's + data rows, then its scale words: the kernels index both from it).""" + kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) + + +def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _select( + layer: "DeepseekV4DecoderLayer", + positions: torch.Tensor, + normed: torch.Tensor, + qr: torch.Tensor, +) -> None: + """An index source's compressor and top-k selection, from K1's normed input + and bf16 qr: ``_forward_csa2_full`` / ``_forward_csa2_reindex`` and the + indexer op, serial.""" + attn = layer.attn + comp, indexer = attn.compressor, attn.indexer + assert indexer is not None + if comp is not None: + kv_score = torch.mm( + normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32 + ) + latent = comp(kv_score, positions) + comp.insert_cache(latent, positions, attn.rotary_emb) + indexer._produce_k(latent, positions, attn.indexer_rotary_emb) + weights, _ = indexer.weights_proj(normed) + q, q_scale, weights = indexer.forward_q( + qr, None, weights, positions, attn.indexer_rotary_emb + ) + assert weights is not None + indexer.indexer_op( + normed, (q, q_scale) if q_scale is not None else q, None, weights + ) + + +@dataclass +class Packed: + """This rank's launch state: the runner (one a rank, its peer memory + collective over the TP group), each kernel layer's weight views, and the + metadata keys ``skip_step`` validates for every kernel layer at once.""" + + runner: Any + first: int + words: torch.Tensor + weights: dict[int, Any] = field(default_factory=dict) + swa_keys: list[str] = field(default_factory=list) + comp_keys: list[str] = field(default_factory=list) + need_swa: bool = False + + +class Kernel: + """What both launches share: the layer, its width and its weight views.""" + + ffn_only = False + + def __init__(self, packed: Packed, layer: "DeepseekV4DecoderLayer", width: int): + self.packed = packed + self.layer = layer + self.width = width + self.weights = packed.weights[layer.attn.layer_id] + self.cache_layers: tuple[str, ...] = () + + def health(self) -> torch.Tensor: + return self.packed.words + + def warmup_keys(self) -> list[Any]: + raise NotImplementedError + + def compile(self, key: Any) -> None: + raise NotImplementedError + + def launch(self, step: MonoStep, *args, **kwargs) -> Any: + raise NotImplementedError + + +class WholeLayer(Kernel): + """The whole backbone layer as K1 and K2, from the layer's own inputs.""" + + def __init__(self, packed: Packed, layer: "DeepseekV4DecoderLayer", width: int): + super().__init__(packed, layer, width) + attn = layer.attn + self.ratio = attn.compress_ratio + self.index = attn.indexer is not None + self.swa_key = attn.swa_cache_layer.prefix + self.comp_key = attn.compressed_cache_prefix + self.cache_layers = (self.swa_key,) + ( + (self.comp_key,) if self.comp_key is not None else () + ) + # K1's seam folds the owed post into the residual, except where the + # residual arrives settled: the embedding, or after the Engram write. + settled = attn.layer_id == packed.first or layer.engram is not None + self.folds = (True, False) if settled else (True,) + + def warmup_keys(self) -> list[Any]: + return [(self.width, self.ratio, f, self.index) for f in self.folds] + + def compile(self, key: Any) -> None: + width, ratio, fold, index = key + self.packed.runner.scratch(width) + self.packed.runner.kernels(width, ratio, fold=fold, index=index) + + def launch( + self, + step: MonoStep, + x: torch.Tensor, + positions: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + pre_mix: torch.Tensor | None, + engram_hashes: torch.Tensor | None = None, + engram_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, ...] | None: + layer, attn = self.layer, self.layer.attn + runner = self.packed.runner + first = residual is None + if first and x.dim() != 2: + return None + tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) + if not all(t is not None and t.is_contiguous() for t in tensors): + return None + fold = True + if first: + residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + pre_mix = runner.first_pre_mix(x.shape[0]) + post_mix = res_mix = pre_mix # unread without the fold + fold = False + elif layer.engram is not None and engram_hashes is not None: + residual = layer.mhc_post(x, residual, post_mix, res_mix) + hashes = engram_hashes[:, layer.engram.layer_hash_index] + residual = layer.engram(residual, hashes, engram_mask) + fold = False + swa = step.attn_metadata[self.swa_key] + kw: dict[str, Any] = {} + if self.index: + kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) + if self.ratio: + comp = step.attn_metadata[self.comp_key] + kw.update( + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // self.ratio + ), + comp_block_table=comp.block_table, + ) + return runner.forward( + self.weights, + x, + residual, + post_mix, + res_mix, + pre_mix, + positions, + swa.slot_mapping, + _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), + swa.decode_swa_indices, + swa.decode_swa_lens, + swa.token_to_req_indices, + fold=fold, + **kw, + ) + + +class FfnSeam(Kernel): + """The FFN seam, MoE and both reductions as one launch, from wo_b's + unreduced output after vLLM's attention.""" + + ffn_only = True + + def warmup_keys(self) -> list[Any]: + return [(self.width, "ffn")] + + def compile(self, key: Any) -> None: + self.packed.runner.scratch(key[0]) + self.packed.runner.ffn_kernel(key[0]) + + def launch( + self, + step: MonoStep, + part: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + pre_mix: torch.Tensor, + ) -> tuple[torch.Tensor, ...] | None: + tensors = (part, residual, post_mix, res_mix, pre_mix) + if not all(t.is_contiguous() for t in tensors): + return None + return self.packed.runner.ffn(self.weights, *tensors) + + +def _compute_units() -> int: + return current_platform.num_compute_units(torch.accelerator.current_device_index()) + + +def _widths(vllm_config) -> tuple[int, ...]: + """The configured capture sizes the kernels can serve, so a deployment + that asked for fewer graphs keeps them; every size at all when it asked for + none.""" + sizes = vllm_config.compilation_config.cudagraph_capture_sizes or [] + served = tuple(s for s in sorted(sizes) if 1 <= s <= MAX_ROWS) + return served or tuple(range(1, MAX_ROWS + 1)) + + +def _refuse_config(vllm_config) -> list[str]: + """Every reason this config cannot run the kernels, before load.""" + reasons = [] + if vllm_config.parallel_config.tensor_parallel_size not in (2, 4): + reasons.append("needs tensor parallel size 2 or 4") + if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + reasons.append( + "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" + ) + try: + import vllm.models.deepseek_v41.amd.mono.runner # noqa: F401 + except ImportError as err: + reasons.append(f"needs FlyDSL and AITER's FlyDSL helpers ({err})") + return reasons + + +def _refuse_model(model) -> list[str]: + """Every reason the loaded model and this GPU cannot run the kernels. The + kernels read every one of the 384 experts on each rank, TP-sharded, in + AITER's A8W4 MXFP4 layout, from CTAs that all stay resident.""" + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS + from vllm.platforms.rocm import get_cdna_version + + reasons = [] + if get_cdna_version() not in CDNA_VERSIONS: + reasons.append(f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}") + cus = _compute_units() + if cus < BLOCKS: + reasons.append(f"needs {BLOCKS} compute units, the GPU has {cus}") + backends = set() + for layer in model.layers: + experts = getattr(getattr(layer, "ffn", None), "experts", None) + routed = getattr(experts, "routed_experts", None) + if routed is not None: + backends.add(getattr(routed.quant_method, "mxfp4_backend", None)) + wrong = backends - {Mxfp4MoeBackend.AITER_MXFP4_BF16} + if wrong: + reasons.append(f"needs the AITER_MXFP4_BF16 MoE backend, not {wrong}") + return reasons + + +def _skip_layer(ctx: MonoContext, layer) -> str | None: + """Why this layer keeps vLLM's path for every step, or None when it takes a + launch (whole or FFN only: ``_build`` picks which).""" + ffn = layer.ffn + if layer.use_sequence_parallel: + return "sequence-parallel layer" + if not layer.fuse_seam_norm: + return "needs aiter's fused mHC seam" + if ( + ffn.shared_experts is None + or ffn.gate.tid2eid is not None + or ffn.n_routed_experts != 384 + or ffn.n_activated_experts != 6 + or ffn.routed_scaling_factor != 1.5 + or ffn.swiglu_limit != 10.0 + or ffn.scoring_func != "sqrtsoftplus" + or not ffn.renormalize + ): + return "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + return None + + +def _ffn_only(layer) -> str | None: + """Why this layer's attention seam and attention stay on vLLM's path, which + leaves it the FFN launch, or None when it runs whole.""" + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + attn = layer.attn + if attn.compressor is not None and attn.indexer is None: + return "a compressor without an indexer" + if attn.indexer is not None and not isinstance( + attn.indexer, DeepseekV41RocmMxfp4Indexer + ): + return f"a {type(attn.indexer).__name__}" + if attn.compress_ratio not in (0, 1, 2): + return f"compress ratio {attn.compress_ratio}" + if attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": + return f"a {attn.kv_cache_dtype} KV cache" + return None + + +def _skip_step(ctx: MonoContext, step: MonoStep) -> str | None: + """Why this step keeps vLLM's path at every layer, or None when the kernels + take it. One verdict for every kernel layer, from the batch the whole TP + group sees alike.""" + packed: Packed = ctx.state + md = step.attn_metadata + if not 1 <= step.width <= MAX_ROWS: + return f"{step.width} rows" + if step.runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + # a PIECEWISE (breakable) capture is replayed for mixed batches too + return f"{step.runtime_mode} step" + if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + return "no attention metadata" + if any(md.get(k) is None for k in packed.comp_keys): + return "layer metadata missing" + for key in packed.swa_keys: + swa = cast("DeepseekV4ROCMAiterSparseSWAMetadata | None", md.get(key)) + if swa is None: + return "layer metadata missing" + if swa.num_prefills != 0 or swa.num_decodes == 0: + return "not a decode-only step" + if swa.num_decode_tokens > step.width: + return "more decode tokens than rows" + if not packed.need_swa: + continue + if swa.decode_swa_width != SWA_WIDTH: + return f"SWA width {swa.decode_swa_width}" + if ( + swa.decode_swa_indices is None + or swa.decode_swa_lens is None + or swa.token_to_req_indices is None + ): + return "SWA decode metadata missing" + return None + + +def _pack(model, ctx: MonoContext) -> Packed: + """This rank's runner, collective over the TP group.""" + from vllm.models.common.mono.peer import tp_group + from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer + + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + rank, world, cpu_group = tp_group() + packed = Packed( + runner=DSV41MonoLayer(world, rank, cpu_group, ctx.device), + first=getattr(model, "start_layer", 0), + words=torch.zeros(world, dtype=torch.int32, device=ctx.device), + ) + ctx.state = packed + return packed + + +def _weights(layer, ffn_only: bool, world: int): + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + + a, f = layer.attn, layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + weights = MonoLayerWeights( + attn=None if ffn_only else _attn_weights(a, world), + hc_attn_fn=layer.hc_attn_fn, + hc_attn_scale=layer.hc_attn_scale, + hc_attn_base=layer.hc_attn_base, + attn_norm=layer.attn_norm.weight, + hc_ffn_fn=layer.hc_ffn_fn, + hc_ffn_scale=layer.hc_ffn_scale, + hc_ffn_base=layer.hc_ffn_base, + ffn_norm=layer.ffn_norm.weight, + gate_w=f.gate.weight, + bias=f.gate.e_score_correction_bias, + w13=e.w13_weight, + w13_s=e.w13_weight_scale, + w2=e.w2_weight, + w2_s=e.w2_weight_scale, + sgu=sh.gate_up_proj.weight, + sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), + sw2=sh.down_proj.weight, + sw2_s=_scale_bytes(sh.down_proj.weight_scale), + ) + weights.check(world) + return weights + + +def _attn_weights(a, world: int): + from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights + + attn = AttnWeights( + layer_id=a.layer_id, + wqkv=a.fused_wqa_wkv.weight, + wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=a.wq_b.weight, + wq_b_scale=_scale_bytes(a.wq_b.weight_scale), + wo_a=a.wo_a.weight, + wo_a_scale=_scale_bytes(a.wo_a.weight_scale), + wo_b=a.wo_b.weight, + wo_b_scale=_scale_bytes(a.wo_b.weight_scale), + attn_sink=a.attn_sink, + cos_sin=a.rotary_emb.cos_sin_cache, + ratio=a.compress_ratio, + ) + attn.check(Dims(world)) + return attn + + +def _bind(ctx: MonoContext, caches) -> list[str]: + """Every reason the kernels cannot read these caches: they index each as + [blocks, block, RECORD] bytes.""" + reasons = [] + for name, cache in caches.items(): + if not isinstance(cache, torch.Tensor): + reasons.append(f"{name}: no cache tensor") + continue + if cache.element_size() != 1: + reasons.append(f"{name}: {cache.dtype} records, needs bytes") + if cache.dim() < 2: + reasons.append(f"{name}: {cache.dim()} dims, needs [blocks, block, bytes]") + elif cache.stride(0) < RECORD: + reasons.append(f"{name}: block stride {cache.stride(0)} under {RECORD}") + return reasons + + +def _build(ctx: MonoContext, packed: Packed, layer, width: int): + world = ctx.vllm_config.parallel_config.tensor_parallel_size + why = _ffn_only(layer) + if layer.attn.layer_id not in packed.weights: + packed.weights[layer.attn.layer_id] = _weights(layer, why is not None, world) + if why is not None: + logger.info_once("DSv4.1 MonoKernel: the FFN launch for layers with %s", why) + # A decode step hands wo_b's TP partial to the FFN launch; the layer + # reduces it on the steps that keep vLLM's path. + layer.attn.wo_b.reduce_results = False + return FfnSeam(packed, layer, width) + kernel = WholeLayer(packed, layer, width) + packed.need_swa = True + for key, keys in ( + (kernel.swa_key, packed.swa_keys), + (kernel.comp_key, packed.comp_keys), + ): + if key is not None and key not in keys: + keys.append(key) + return kernel + + +def _layers(config) -> list[int]: + return list(range(config.num_hidden_layers)) + + +MONO_SPEC = MonoSpec( + name="dsv41-cdna4", + supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), + caches={"swa": CacheReq("swa", "fp8_ds_mla", "blocks,block,bytes")}, + weight_source="views", + layers=_layers, + widths=_widths, + refuse_config=_refuse_config, + refuse_model=_refuse_model, + skip_layer=_skip_layer, + skip_step=_skip_step, + pack=_pack, + repack=None, + bind=_bind, + build=_build, +) + +MONO_SPECS = (MONO_SPEC,) From 88e92df4d41806dc90d0c3a8f07c6af2cb4e8c59 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 08:38:59 -0500 Subject: [PATCH 03/79] fix: remove unused envvars Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 11 ---------- vllm/envs.py | 12 ----------- vllm/models/common/mono/abi.py | 15 -------------- vllm/models/common/mono/failstop.py | 32 ++++++++--------------------- 4 files changed, 8 insertions(+), 62 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 025473b1ff11..d494ffd9d639 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -111,10 +111,8 @@ def mono(self, config, model): @pytest.fixture(autouse=True) def _counters(self): failstop.STEPS.clear() - failstop.INCIDENTS.clear() yield failstop.STEPS.clear() - failstop.INCIDENTS.clear() def spec(self, config): spec, refusals = select(load_specs(self.specs_path), config) @@ -237,18 +235,9 @@ def test_fail_stop_raises(self, mono): words[0] = 1 with pytest.raises(RuntimeError, match="MonoKernel step failed"): failstop.check(words) - assert failstop.INCIDENTS["0"] == 1 - - def test_fail_stop_warns_when_unsafe(self, mono, monkeypatch): - monkeypatch.setattr("vllm.envs.VLLM_MONO_UNSAFE_WARN", True) - words = torch.zeros_like(mono.health()) - words[0] = 1 - failstop.check(words) - assert failstop.INCIDENTS["0"] == 1 def test_healthy_step_passes(self, mono): failstop.check(torch.zeros_like(mono.health())) - assert not failstop.INCIDENTS def test_counters_follow_full_replay(self, config, model, mono): width = min(mono.widths) diff --git a/vllm/envs.py b/vllm/envs.py index b6b07524e667..bb1c03a9dabd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -156,8 +156,6 @@ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False - VLLM_MONO_UNSAFE_WARN: bool = False - VLLM_MONO_FAULT_INJECT: str = "" VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1399,16 +1397,6 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS", "False").lower() in ("true", "1") ), - # Warn instead of failing the step when a mono decode kernel reports an - # error word (an in-kernel wait gave up). For debugging only: the step's - # output is not trustworthy. By default is disabled. - "VLLM_MONO_UNSAFE_WARN": lambda: ( - os.getenv("VLLM_MONO_UNSAFE_WARN", "False").lower() in ("true", "1") - ), - # Make one mono decode launch report a fault, to test the fail-stop path: - # "rank:epoch:us" sets the rank, the step epoch and how long the injected - # stall lasts. By default is disabled. - "VLLM_MONO_FAULT_INJECT": lambda: os.getenv("VLLM_MONO_FAULT_INJECT", ""), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/common/mono/abi.py b/vllm/models/common/mono/abi.py index d684416843b6..888ce754fbf1 100644 --- a/vllm/models/common/mono/abi.py +++ b/vllm/models/common/mono/abi.py @@ -11,14 +11,8 @@ of the output rank's peer buffer. A rank whose wait gives up writes its word there, drains the store, then sets FAILED, so nothing it stores later can overtake the flag. The runner reads these words after every step. - -The fault hook (tests only, ``VLLM_MONO_FAULT_INJECT=rank:epoch:us``): a -build-time stall of one CTA of rank ``rank`` for ``us`` microseconds, once, in -the step at ``epoch``; no code when unset. """ -from vllm import envs - MONO_ABI_VERSION = 1 RECORD_WORDS = 8 @@ -30,12 +24,3 @@ def health_words(ctas: int) -> int: """A rank's health words for a grid of ``ctas``: the header, a record a CTA.""" return (ctas + 1) * RECORD_WORDS - - -def fault_inject() -> tuple[int, int, int] | None: - """(rank, epoch, microseconds) of ``VLLM_MONO_FAULT_INJECT``, or None.""" - v = envs.VLLM_MONO_FAULT_INJECT - if not v: - return None - rank, epoch, us = (int(x) for x in v.split(":")) - return rank, epoch, us diff --git a/vllm/models/common/mono/failstop.py b/vllm/models/common/mono/failstop.py index 439952886db3..a925badfbe43 100644 --- a/vllm/models/common/mono/failstop.py +++ b/vllm/models/common/mono/failstop.py @@ -6,25 +6,17 @@ copies the model's error words (``abi``) next to the sampled token ids, and ``check`` raises before that step's tokens are returned. The engine then stops like on any other engine death, and the deployment's restart policy recovers -it. ``VLLM_MONO_UNSAFE_WARN=1`` (tests and benchmarks only, no CLI flag) logs -the incident and keeps serving instead. +it. -Counters (``drain``, exported as ``vllm:monokernel_steps`` and -``vllm:monokernel_incidents``): steps on the kernels, steps on vLLM's path by -reason, and incidents by failed rank. +Counters (``drain``, exported as ``vllm:monokernel_steps``): steps on the +kernels, and steps on vLLM's path by reason. """ from collections import Counter import torch -from vllm import envs -from vllm.logger import init_logger - -logger = init_logger(__name__) - STEPS: Counter[tuple[str, str]] = Counter() -INCIDENTS: Counter[str] = Counter() def count_step(spec: str, why: str | None) -> None: @@ -34,14 +26,12 @@ def count_step(spec: str, why: str | None) -> None: def drain() -> dict[str, int] | None: - """The counts since the last drain: ``"steps::"`` (path "mono" - or the reason the step ran vLLM's layers) and ``"incidents:"``.""" - if not STEPS and not INCIDENTS: + """The counts since the last drain: ``"steps::"``, where path + is "mono" or the reason the step ran vLLM's layers.""" + if not STEPS: return None out = {f"steps:{spec}:{path}": n for (spec, path), n in STEPS.items()} - out.update({f"incidents:{rank}": n for rank, n in INCIDENTS.items()}) STEPS.clear() - INCIDENTS.clear() return out @@ -53,19 +43,13 @@ def check(words: torch.Tensor) -> None: a TP rank, nonzero once one of its in-kernel waits gave up). Raises: - RuntimeError: A rank gave up a wait, unless ``VLLM_MONO_UNSAFE_WARN``. + RuntimeError: A rank gave up a wait. """ if not words.any(): return ranks = [r for r, w in enumerate(words.tolist()) if w] - for r in ranks: - INCIDENTS[str(r)] += 1 - msg = ( + raise RuntimeError( f"MonoKernel step failed: TP ranks {ranks} gave up an in-kernel wait; " "the step's tokens are discarded" ) - if envs.VLLM_MONO_UNSAFE_WARN: - logger.error("%s (VLLM_MONO_UNSAFE_WARN=1: serving on)", msg) - return - raise RuntimeError(msg) From 7471e2a66459d483931cdcad78c9018f12852a37 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 08:54:39 -0500 Subject: [PATCH 04/79] fix: rename Packed to LaunchState Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 2 +- tests/models/common/mono/reference.py | 24 +++---- tests/models/common/mono/test_reference.py | 6 +- tests/models/test_dsv41_mono_layers.py | 2 +- vllm/models/common/mono/routing.py | 10 +-- vllm/models/common/mono/spec.py | 15 +++-- vllm/models/deepseek_v41/amd/mono_spec.py | 78 ++++++++++++---------- 7 files changed, 74 insertions(+), 63 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index d494ffd9d639..9b81a6f4d2aa 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -175,7 +175,7 @@ def test_every_rank_decides_alike(self, config, model): decisions = [] for rank in range(2): other = attach(self.make_model(config), config, self.specs_path) - other.ctx.rank = rank + other.ctx = replace(other.ctx, rank=rank) with forward( config, CUDAGraphMode.NONE, self.decode_metadata(config, model) ): diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 3de113b302b7..c9aacd796e65 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -58,7 +58,7 @@ def step_health(self) -> torch.Tensor | None: @dataclass -class Packed: +class LaunchState: """What ``pack`` builds once at attach: views of the model's weights, never a copy, and the error words every kernel of this rank shares.""" @@ -69,10 +69,10 @@ class Packed: class ReferenceKernel: """One layer's launch at one width: calls the layer it replaces.""" - def __init__(self, ctx: MonoContext, packed: Packed, layer, width: int): + def __init__(self, ctx: MonoContext, state: LaunchState, layer, width: int): self.cache_layers = (layer.attn_name,) self.ctx = ctx - self.packed = packed + self.state = state self.layer = layer self.width = width self.launched: list[int] = [] @@ -83,7 +83,7 @@ def launch(self, step: MonoStep, positions: torch.Tensor, *args, **kwargs): return self.layer(positions, *args, **kwargs) def health(self) -> torch.Tensor: - return self.packed.words + return self.state.words def warmup_keys(self) -> list[Any]: return [(self.layer.index, self.width)] @@ -93,7 +93,7 @@ def compile(self, key: Any) -> None: def fail(self, rank: int = 0) -> None: """Reports an in-kernel wait that gave up, as a launch would.""" - self.packed.words[rank] = 1 + self.state.words[rank] = 1 def layers(config) -> list[int]: @@ -122,21 +122,21 @@ def skip_layer(ctx: MonoContext, layer) -> str | None: return "dense layer" if layer.dense else None -def skip_step(ctx: MonoContext, step: MonoStep) -> str | None: +def skip_step(ctx: MonoContext, state: LaunchState, step: MonoStep) -> str | None: return "prefill" if step.attn_metadata.get("prefill") else None -def pack(model, ctx: MonoContext) -> Packed: +def pack(model, ctx: MonoContext) -> LaunchState: _, world, _ = _tp() - return Packed( + return LaunchState( weights={i: layer.weight.data for i, layer in enumerate(model.layers)}, words=torch.zeros(world, dtype=torch.int32, device=ctx.device), ) -def repack(model, ctx: MonoContext, packed: Packed) -> None: +def repack(model, ctx: MonoContext, state: LaunchState) -> None: for i, layer in enumerate(model.layers): - packed.weights[i].copy_(layer.weight.data) + state.weights[i].copy_(layer.weight.data) def bind(ctx: MonoContext, caches) -> list[str]: @@ -149,8 +149,8 @@ def bind(ctx: MonoContext, caches) -> list[str]: return reasons -def build(ctx: MonoContext, packed: Packed, layer, width: int) -> ReferenceKernel: - return ReferenceKernel(ctx, packed, layer, width) +def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceKernel: + return ReferenceKernel(ctx, state, layer, width) def _tp(): diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index c622fdbf8407..b34d39ad5d89 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -92,13 +92,13 @@ def test_repack_writes_in_place(): config = reference.make_config() model = reference.make_model(config) mono = attach(model, config, SPECS_PATH) - before = {i: w.data_ptr() for i, w in mono.packed.weights.items()} + before = {i: w.data_ptr() for i, w in mono.state.weights.items()} with torch.no_grad(): for layer in model.layers: layer.weight.copy_(torch.full_like(layer.weight, 7)) mono.repack(model) - assert {i: w.data_ptr() for i, w in mono.packed.weights.items()} == before - assert all((w == 7).all() for w in mono.packed.weights.values()) + assert {i: w.data_ptr() for i, w in mono.state.weights.items()} == before + assert all((w == 7).all() for w in mono.state.weights.values()) def test_repack_is_refused_when_the_spec_declares_none(): diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index ddca6ba1c3dc..fd151115e8d0 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -223,7 +223,7 @@ def mono_layers(self) -> dict: ) for i in self._decoder_layers() } - runner = None if mono is None else mono.ctx.state.runner + runner = None if mono is None else mono.state.runner return { "layers": layers, "epoch": -1 if runner is None else int(runner.epoch[0]), diff --git a/vllm/models/common/mono/routing.py b/vllm/models/common/mono/routing.py index 4201f0c6119c..1fe09698454d 100644 --- a/vllm/models/common/mono/routing.py +++ b/vllm/models/common/mono/routing.py @@ -72,7 +72,7 @@ def __init__(self, spec: MonoSpec, model, vllm_config: "VllmConfig"): self.ctx = MonoContext(vllm_config, rank, next(model.parameters()).device) self.widths = spec.widths(vllm_config) self.kernels: dict[tuple[int, int], LayerKernel] = {} - self.packed: Any = None + self.state: Any = None self.step: MonoStep | None = None self.baked: dict[int, str | None] = {} self.last: str | None = "no step" @@ -83,11 +83,11 @@ def __init__(self, spec: MonoSpec, model, vllm_config: "VllmConfig"): if not indices: self.cache_layers: list[str] = [] return - self.packed = spec.pack(model, self.ctx) + self.state = spec.pack(model, self.ctx) for i in indices: for w in self.widths: self.kernels[i, w] = spec.build( - self.ctx, self.packed, model.layers[i], w + self.ctx, self.state, model.layers[i], w ) self.cache_layers = sorted( {n for k in self.kernels.values() for n in k.cache_layers} @@ -140,7 +140,7 @@ def repack(self, model) -> None: raise MonoUnsupported( f"MonoKernel {self.spec.name}: weights cannot move after attach" ) - self.spec.repack(model, self.ctx, self.packed) + self.spec.repack(model, self.ctx, self.state) def _bind(self) -> None: """Re-checks the caches whenever they moved (keyed on every tensor's @@ -199,7 +199,7 @@ def _decide(self, positions: torch.Tensor) -> str | None: if n not in self.widths: return "width" step = MonoStep(n, n, positions, fc.attn_metadata, fc.cudagraph_runtime_mode) - why = self.spec.skip_step(self.ctx, step) + why = self.spec.skip_step(self.ctx, self.state, step) if why is None: self._bind() self.step = step diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 4b05a4954fd0..b7257fc8cac9 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -48,7 +48,7 @@ class CacheReq: layout: str -@dataclass +@dataclass(frozen=True) class MonoContext: """One rank's view of the deployment, built once at attach. @@ -56,15 +56,12 @@ class MonoContext: vllm_config: The engine's config. rank: This rank's index in the TP group. device: The device the kernels launch on. - state: The spec's own per-rank state, whatever ``pack`` keeps there. - Opaque to the library, which never reads it. """ vllm_config: "VllmConfig" rank: int device: torch.device - state: Any = None @dataclass @@ -148,12 +145,13 @@ class MonoSpec: skip_step: Why one step takes vLLM's path, None when it takes the kernel's. pack: The weights and launch state, built once at attach. The only - place a weight copy is ever made. + place a weight copy is ever made. Opaque to the library, which + only hands it back to the spec's own callbacks. repack: Writes new weights into the allocation ``pack`` returned, in place, because the kernels and the captured graphs hold its address. None refuses anything that moves weights after attach. bind: Checks the KV caches the kernels will read, on every re-bind. - build: One ``LayerKernel`` for (context, packed, layer, width). + build: One ``LayerKernel`` for (context, state, layer, width). """ @@ -162,14 +160,17 @@ class MonoSpec: caches: Mapping[str, CacheReq] weight_source: Literal["views", "convert"] + # Model shape layers: Callable[[Any], list[int]] widths: Callable[["VllmConfig"], tuple[int, ...]] + # Capabilities refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] skip_layer: Callable[[MonoContext, Any], str | None] - skip_step: Callable[[MonoContext, MonoStep], str | None] + skip_step: Callable[[MonoContext, Any, MonoStep], str | None] + # Callbacks pack: Callable[[Any, MonoContext], Any] repack: Callable[[Any, MonoContext, Any], None] | None bind: Callable[[MonoContext, Mapping[str, Any]], list[str]] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 9d45e01520be..fbf1fc38de70 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -97,10 +97,15 @@ def _select( @dataclass -class Packed: - """This rank's launch state: the runner (one a rank, its peer memory - collective over the TP group), each kernel layer's weight views, and the - metadata keys ``skip_step`` validates for every kernel layer at once.""" +class LaunchState: + """What this rank's launches hold across every step: the runner (one a + rank, its peer memory collective over the TP group), the error words, each + kernel layer's weight views, and the metadata keys ``skip_step`` validates + for every kernel layer at once. + + Built once at attach and never moved: the kernels and the captured graphs + hold its addresses. + """ runner: Any first: int @@ -116,15 +121,17 @@ class Kernel: ffn_only = False - def __init__(self, packed: Packed, layer: "DeepseekV4DecoderLayer", width: int): - self.packed = packed + def __init__( + self, launch_state: LaunchState, layer: "DeepseekV4DecoderLayer", width: int + ): + self.launch_state = launch_state self.layer = layer self.width = width - self.weights = packed.weights[layer.attn.layer_id] + self.weights = launch_state.weights[layer.attn.layer_id] self.cache_layers: tuple[str, ...] = () def health(self) -> torch.Tensor: - return self.packed.words + return self.launch_state.words def warmup_keys(self) -> list[Any]: raise NotImplementedError @@ -139,8 +146,10 @@ def launch(self, step: MonoStep, *args, **kwargs) -> Any: class WholeLayer(Kernel): """The whole backbone layer as K1 and K2, from the layer's own inputs.""" - def __init__(self, packed: Packed, layer: "DeepseekV4DecoderLayer", width: int): - super().__init__(packed, layer, width) + def __init__( + self, launch_state: LaunchState, layer: "DeepseekV4DecoderLayer", width: int + ): + super().__init__(launch_state, layer, width) attn = layer.attn self.ratio = attn.compress_ratio self.index = attn.indexer is not None @@ -151,7 +160,7 @@ def __init__(self, packed: Packed, layer: "DeepseekV4DecoderLayer", width: int): ) # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == packed.first or layer.engram is not None + settled = attn.layer_id == launch_state.first or layer.engram is not None self.folds = (True, False) if settled else (True,) def warmup_keys(self) -> list[Any]: @@ -159,8 +168,8 @@ def warmup_keys(self) -> list[Any]: def compile(self, key: Any) -> None: width, ratio, fold, index = key - self.packed.runner.scratch(width) - self.packed.runner.kernels(width, ratio, fold=fold, index=index) + self.launch_state.runner.scratch(width) + self.launch_state.runner.kernels(width, ratio, fold=fold, index=index) def launch( self, @@ -175,7 +184,7 @@ def launch( engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: layer, attn = self.layer, self.layer.attn - runner = self.packed.runner + runner = self.launch_state.runner first = residual is None if first and x.dim() != 2: return None @@ -234,8 +243,8 @@ def warmup_keys(self) -> list[Any]: return [(self.width, "ffn")] def compile(self, key: Any) -> None: - self.packed.runner.scratch(key[0]) - self.packed.runner.ffn_kernel(key[0]) + self.launch_state.runner.scratch(key[0]) + self.launch_state.runner.ffn_kernel(key[0]) def launch( self, @@ -249,7 +258,7 @@ def launch( tensors = (part, residual, post_mix, res_mix, pre_mix) if not all(t.is_contiguous() for t in tensors): return None - return self.packed.runner.ffn(self.weights, *tensors) + return self.launch_state.runner.ffn(self.weights, *tensors) def _compute_units() -> int: @@ -348,11 +357,12 @@ def _ffn_only(layer) -> str | None: return None -def _skip_step(ctx: MonoContext, step: MonoStep) -> str | None: +def _skip_step( + ctx: MonoContext, launch_state: LaunchState, step: MonoStep +) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP group sees alike.""" - packed: Packed = ctx.state md = step.attn_metadata if not 1 <= step.width <= MAX_ROWS: return f"{step.width} rows" @@ -361,9 +371,9 @@ def _skip_step(ctx: MonoContext, step: MonoStep) -> str | None: return f"{step.runtime_mode} step" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches return "no attention metadata" - if any(md.get(k) is None for k in packed.comp_keys): + if any(md.get(k) is None for k in launch_state.comp_keys): return "layer metadata missing" - for key in packed.swa_keys: + for key in launch_state.swa_keys: swa = cast("DeepseekV4ROCMAiterSparseSWAMetadata | None", md.get(key)) if swa is None: return "layer metadata missing" @@ -371,7 +381,7 @@ def _skip_step(ctx: MonoContext, step: MonoStep) -> str | None: return "not a decode-only step" if swa.num_decode_tokens > step.width: return "more decode tokens than rows" - if not packed.need_swa: + if not launch_state.need_swa: continue if swa.decode_swa_width != SWA_WIDTH: return f"SWA width {swa.decode_swa_width}" @@ -384,20 +394,18 @@ def _skip_step(ctx: MonoContext, step: MonoStep) -> str | None: return None -def _pack(model, ctx: MonoContext) -> Packed: +def _pack(model, ctx: MonoContext) -> LaunchState: """This rank's runner, collective over the TP group.""" from vllm.models.common.mono.peer import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS rank, world, cpu_group = tp_group() - packed = Packed( + return LaunchState( runner=DSV41MonoLayer(world, rank, cpu_group, ctx.device), first=getattr(model, "start_layer", 0), words=torch.zeros(world, dtype=torch.int32, device=ctx.device), ) - ctx.state = packed - return packed def _weights(layer, ffn_only: bool, world: int): @@ -471,22 +479,24 @@ def _bind(ctx: MonoContext, caches) -> list[str]: return reasons -def _build(ctx: MonoContext, packed: Packed, layer, width: int): +def _build(ctx: MonoContext, launch_state: LaunchState, layer, width: int): world = ctx.vllm_config.parallel_config.tensor_parallel_size why = _ffn_only(layer) - if layer.attn.layer_id not in packed.weights: - packed.weights[layer.attn.layer_id] = _weights(layer, why is not None, world) + if layer.attn.layer_id not in launch_state.weights: + launch_state.weights[layer.attn.layer_id] = _weights( + layer, why is not None, world + ) if why is not None: logger.info_once("DSv4.1 MonoKernel: the FFN launch for layers with %s", why) # A decode step hands wo_b's TP partial to the FFN launch; the layer # reduces it on the steps that keep vLLM's path. layer.attn.wo_b.reduce_results = False - return FfnSeam(packed, layer, width) - kernel = WholeLayer(packed, layer, width) - packed.need_swa = True + return FfnSeam(launch_state, layer, width) + kernel = WholeLayer(launch_state, layer, width) + launch_state.need_swa = True for key, keys in ( - (kernel.swa_key, packed.swa_keys), - (kernel.comp_key, packed.comp_keys), + (kernel.swa_key, launch_state.swa_keys), + (kernel.comp_key, launch_state.comp_keys), ): if key is not None and key not in keys: keys.append(key) From 967ddd453a8320faa8f119fb086c2cfaa6f8f789 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 09:10:03 -0500 Subject: [PATCH 05/79] fix: my pr comments Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 18 +++++----- tests/models/common/mono/test_reference.py | 8 ++--- vllm/models/common/mono/__init__.py | 3 +- vllm/models/common/mono/abi.py | 26 -------------- vllm/models/common/mono/failstop.py | 2 +- vllm/models/common/mono/guards.py | 42 +++++++++++----------- vllm/models/common/mono/layer.py | 13 ++++--- vllm/models/common/mono/routing.py | 8 ++--- vllm/models/common/mono/spec.py | 2 +- 9 files changed, 51 insertions(+), 71 deletions(-) delete mode 100644 vllm/models/common/mono/abi.py diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 9b81a6f4d2aa..a34f945791d5 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -30,11 +30,11 @@ from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.models.common.mono import failstop from vllm.models.common.mono.guards import ( + _load_specs, + _refuse_config, + _select_spec, + _set_capture_sizes, configure, - load_specs, - refuse_config, - select, - set_capture_sizes, ) from vllm.models.common.mono.layer import attach, forward_layer from vllm.models.common.mono.spec import Feature, MonoUnsupported @@ -115,7 +115,7 @@ def _counters(self): failstop.STEPS.clear() def spec(self, config): - spec, refusals = select(load_specs(self.specs_path), config) + spec, refusals = _select_spec(_load_specs(self.specs_path), config) assert spec is not None, refusals return spec @@ -123,8 +123,8 @@ def test_refusals_report_every_reason(self, config): config.lora_config = object() config.speculative_config = object() config.use_v2_model_runner = False - for spec in load_specs(self.specs_path): - reasons = refuse_config(spec, config) + for spec in _load_specs(self.specs_path): + reasons = _refuse_config(spec, config) assert sum("Model Runner V2" in r for r in reasons) == 1, reasons for feature in (Feature.LORA, Feature.SPECULATIVE_DECODING): want = 0 if feature in spec.supports else 1 @@ -134,7 +134,7 @@ def test_refusal_names_every_variant(self, config): config.use_v2_model_runner = False with pytest.raises(MonoUnsupported) as excinfo: configure(config, self.specs_path) - for spec in load_specs(self.specs_path): + for spec in _load_specs(self.specs_path): assert spec.name in str(excinfo.value) def test_spec_module_imports_no_kernel_backend(self): @@ -159,7 +159,7 @@ def test_capture_sizes_are_the_widths(self, config): config.compilation_config.cudagraph_capture_sizes = list( range(1, max(widths) * 2 + 1) ) - set_capture_sizes(spec, config) + _set_capture_sizes(spec, config) sizes = config.compilation_config.cudagraph_capture_sizes assert [s for s in sizes if s <= max(widths)] == sorted(widths) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index b34d39ad5d89..c103e1744f14 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -11,7 +11,7 @@ from tests.models.common.mono import reference from tests.models.common.mono.kit import MonoConformance, forward from vllm.config.compilation import CUDAGraphMode -from vllm.models.common.mono.guards import enabled, select +from vllm.models.common.mono.guards import _enabled, _select_spec from vllm.models.common.mono.layer import attach, forward_layer from vllm.models.common.mono.spec import MonoUnsupported @@ -42,19 +42,19 @@ def bad_caches(self, config, model): def test_switch_off_means_no_attach(): config = reference.make_config() config.kernel_config.monokernel = False - assert not enabled(config) + assert not _enabled(config) assert attach(reference.make_model(config), config, SPECS_PATH) is None def test_the_first_servable_variant_wins(): config = reference.make_config() config.scheduler_config.max_num_seqs = 4 - spec, refusals = select(reference.MONO_SPECS, config) + spec, refusals = _select_spec(reference.MONO_SPECS, config) assert spec is reference.REFERENCE assert refusals["reference-wide"] == ["needs max_num_seqs >= 8"] config.scheduler_config.max_num_seqs = 8 - spec, refusals = select(reference.MONO_SPECS, config) + spec, refusals = _select_spec(reference.MONO_SPECS, config) assert spec is reference.REFERENCE_WIDE assert refusals == {} diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index abbeab6ebeb3..039270dfaeef 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -4,8 +4,7 @@ decoder layer's decode step in a single launch. A model brings a ``MonoSpec`` (``spec``); this package owns the rest: refusals (``guards``), the per-step decision and cache re-binding (``routing``), the layer-loop hook (``layer``), -fail-stop and metrics (``failstop``), peer memory (``peer``) and the -error-word layout (``abi``).""" +fail-stop and metrics (``failstop``) and peer memory (``peer``).""" from vllm.models.common.mono.guards import configure from vllm.models.common.mono.layer import ( diff --git a/vllm/models/common/mono/abi.py b/vllm/models/common/mono/abi.py deleted file mode 100644 index 888ce754fbf1..000000000000 --- a/vllm/models/common/mono/abi.py +++ /dev/null @@ -1,26 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The layout every MonoKernel's error words and fault hook follow. - -A rank's health words: a header whose word FAILED is set once any of its -in-kernel waits gave up (sticky: a failed rank stays failed, and its later waits -exit at once), then one record a CTA of the first wait it gave up on (flag, -region id, pair index, tag, seen tag, CTA, clock, pad). - -The output rank's flags: one int32 word a TP rank in the first PEER_HEAD bytes -of the output rank's peer buffer. A rank whose wait gives up writes its word -there, drains the store, then sets FAILED, so nothing it stores later can -overtake the flag. The runner reads these words after every step. -""" - -MONO_ABI_VERSION = 1 - -RECORD_WORDS = 8 -FAILED = 0 -PEER_HEAD = 256 -WAIT_SECONDS = 10 - - -def health_words(ctas: int) -> int: - """A rank's health words for a grid of ``ctas``: the header, a record a CTA.""" - return (ctas + 1) * RECORD_WORDS diff --git a/vllm/models/common/mono/failstop.py b/vllm/models/common/mono/failstop.py index a925badfbe43..29d4dedd3c2d 100644 --- a/vllm/models/common/mono/failstop.py +++ b/vllm/models/common/mono/failstop.py @@ -3,7 +3,7 @@ """Fail-stop and metrics. A MonoKernel never returns tokens from a step it did not finish: the runner -copies the model's error words (``abi``) next to the sampled token ids, and +copies the model's error words next to the sampled token ids, and ``check`` raises before that step's tokens are returned. The engine then stops like on any other engine death, and the deployment's restart policy recovers it. diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index d1b3ee2e80e7..a03d6ae91449 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Refusals. ``configure`` runs once from core config post-init, before load: -it imports only the model's spec module, picks the variant this config can -serve, adds its widths to the CUDA graph capture sizes and raises with every -reason the config is refused. ``agree`` runs a later refusal stage (after load) -on every TP rank and agrees on the verdict before anything collective follows. +"""Refusals. ``configure``, the only name a model calls, runs once from core +config post-init, before load: it imports only the model's spec module, picks +the variant this config can serve, adds its widths to the CUDA graph capture +sizes and raises with every reason the config is refused. The rest is the +library's own and stays private to this package: ``_agree_across_ranks`` runs a +later refusal stage (after load) on every TP rank and agrees on the verdict +before anything collective follows. """ import importlib @@ -29,12 +31,12 @@ ) -def enabled(vllm_config: "VllmConfig") -> bool: +def _enabled(vllm_config: "VllmConfig") -> bool: """Whether ``--monokernel`` asked for the model's MonoKernels.""" return bool(vllm_config.kernel_config.monokernel) -def load_specs(path: str) -> tuple[MonoSpec, ...]: +def _load_specs(path: str) -> tuple[MonoSpec, ...]: """``"package.module:NAME"`` -> the spec variants, importing only that module. ``NAME`` is a ``MonoSpec`` or a tuple of them, in the order they are preferred.""" @@ -43,7 +45,7 @@ def load_specs(path: str) -> tuple[MonoSpec, ...]: return (specs,) if isinstance(specs, MonoSpec) else tuple(specs) -def features_in_use(vllm_config: "VllmConfig") -> set[Feature]: +def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: """The deployment's features a spec has to list in ``supports``.""" pc = vllm_config.parallel_config oc = vllm_config.offload_config @@ -69,12 +71,12 @@ def features_in_use(vllm_config: "VllmConfig") -> set[Feature]: return {f for f, used in on.items() if used} -def refuse_config(spec: MonoSpec, vllm_config: "VllmConfig") -> list[str]: +def _refuse_config(spec: MonoSpec, vllm_config: "VllmConfig") -> list[str]: """Every reason this config cannot run ``spec``: default-deny on features, then the runner and the graph mode, then the spec's own.""" reasons = [ f"{f.value} is not supported" - for f in sorted(features_in_use(vllm_config) - spec.supports) + for f in sorted(_features_in_use(vllm_config) - spec.supports) ] if not vllm_config.use_v2_model_runner: reasons.append("needs Model Runner V2") @@ -84,7 +86,7 @@ def refuse_config(spec: MonoSpec, vllm_config: "VllmConfig") -> list[str]: return reasons + list(spec.refuse_config(vllm_config)) -def refuse_model(spec: MonoSpec, model) -> list[str]: +def _refuse_model(spec: MonoSpec, model) -> list[str]: """Every reason the loaded model cannot run ``spec``. A MonoKernel replaces the layer loop, which a traced graph has already inlined, so a model carrying ``@support_torch_compile`` is refused.""" @@ -99,7 +101,7 @@ def refuse_model(spec: MonoSpec, model) -> list[str]: return reasons + list(spec.refuse_model(model)) -def select( +def _select_spec( specs: Sequence[MonoSpec], vllm_config: "VllmConfig" ) -> tuple[MonoSpec | None, dict[str, list[str]]]: """The first variant this config can serve, and every variant's refusals. @@ -111,21 +113,21 @@ def select( """ refusals: dict[str, list[str]] = {} for spec in specs: - reasons = refuse_config(spec, vllm_config) + reasons = _refuse_config(spec, vllm_config) if not reasons: return spec, refusals refusals[spec.name] = reasons return None, refusals -def refused(refusals: dict[str, list[str]]) -> str: +def _refusal_message(refusals: dict[str, list[str]]) -> str: """Every variant's reasons as one message.""" return "; ".join( f"{name}: {', '.join(reasons)}" for name, reasons in refusals.items() ) -def set_capture_sizes(spec: MonoSpec, vllm_config: "VllmConfig") -> None: +def _set_capture_sizes(spec: MonoSpec, vllm_config: "VllmConfig") -> None: """The capture sizes up to the largest width become exactly the widths, so a replay never pads a request's rows; larger sizes stay as configured (and capture vLLM's path).""" @@ -150,12 +152,12 @@ def configure(vllm_config: "VllmConfig", specs_path: str) -> None: MonoUnsupported: Every variant refused this config. """ - if not enabled(vllm_config): + if not _enabled(vllm_config): return - spec, refusals = select(load_specs(specs_path), vllm_config) + spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) if spec is None: - raise MonoUnsupported("MonoKernel refused: " + refused(refusals)) - set_capture_sizes(spec, vllm_config) + raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) + _set_capture_sizes(spec, vllm_config) if vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: logger.warning( "MonoKernel %s under eager mode: for debugging only, serve with FULL " @@ -164,7 +166,7 @@ def configure(vllm_config: "VllmConfig", specs_path: str) -> None: ) -def agree(refuse: Callable[[], list[str]], group, what: str) -> None: +def _agree_across_ranks(refuse: Callable[[], list[str]], group, what: str) -> None: """``refuse``'s reasons on this rank, then every TP rank's verdict agreed: a rank that refused alone would leave the others spinning in-kernel. diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index a15b6ccdfcc0..3fd532778430 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -10,7 +10,12 @@ from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger from vllm.models.common.mono import failstop -from vllm.models.common.mono.guards import enabled, load_specs, refused, select +from vllm.models.common.mono.guards import ( + _enabled, + _load_specs, + _refusal_message, + _select_spec, +) from vllm.models.common.mono.routing import Mono from vllm.models.common.mono.spec import LayerKernel, MonoUnsupported @@ -41,11 +46,11 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: """ global _ATTACHED - if not enabled(vllm_config): + if not _enabled(vllm_config): return None - spec, refusals = select(load_specs(specs_path), vllm_config) + spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) if spec is None: - raise MonoUnsupported("MonoKernel refused: " + refused(refusals)) + raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) mono = Mono(spec, model, vllm_config) if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) diff --git a/vllm/models/common/mono/routing.py b/vllm/models/common/mono/routing.py index 1fe09698454d..05b6f58cee85 100644 --- a/vllm/models/common/mono/routing.py +++ b/vllm/models/common/mono/routing.py @@ -7,7 +7,7 @@ rank-local state and never a collective. Under a FULL CUDA graph the decision is baked in at capture, so it is a function of the captured size: FULL graphs serve uniform decode batches and every width is a capture size -(``guards.set_capture_sizes``), so no replay pads a request's rows. PIECEWISE +(``guards``), so no replay pads a request's rows. PIECEWISE steps run vLLM's layers. """ @@ -21,7 +21,7 @@ from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel from vllm.models.common.mono import failstop -from vllm.models.common.mono.guards import agree, refuse_model +from vllm.models.common.mono.guards import _agree_across_ranks, _refuse_model from vllm.models.common.mono.peer import tp_group from vllm.models.common.mono.spec import ( LayerKernel, @@ -64,8 +64,8 @@ class Mono: def __init__(self, spec: MonoSpec, model, vllm_config: "VllmConfig"): rank, _, cpu_group = tp_group() self.spec = spec - agree( - lambda: refuse_model(spec, model), + _agree_across_ranks( + lambda: _refuse_model(spec, model), cpu_group, f"MonoKernel {spec.name} refused for the loaded model", ) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index b7257fc8cac9..009f32e9f845 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -98,7 +98,7 @@ def launch(self, step: MonoStep, *args, **kwargs) -> Any: ... def health(self) -> torch.Tensor: - """This rank's error words, in the ``abi`` layout.""" + """This rank's error words: any set word fails the step.""" ... def warmup_keys(self) -> list[Any]: From 549aa945ac921dee7ecd6c0f9b1efff36a9ddd88 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 10:17:57 -0500 Subject: [PATCH 06/79] fix: rename pack->launch_steate and repack->reload_weights Signed-off-by: simondanielsson --- tests/models/common/mono/reference.py | 13 +++++++------ tests/models/common/mono/test_reference.py | 10 +++++----- vllm/models/common/mono/layer.py | 2 +- vllm/models/common/mono/routing.py | 16 ++++++++-------- vllm/models/common/mono/spec.py | 17 +++++++++-------- vllm/models/deepseek_v41/amd/mono_spec.py | 6 +++--- 6 files changed, 33 insertions(+), 31 deletions(-) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index c9aacd796e65..6201a4aa8445 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -59,8 +59,9 @@ def step_health(self) -> torch.Tensor | None: @dataclass class LaunchState: - """What ``pack`` builds once at attach: views of the model's weights, never - a copy, and the error words every kernel of this rank shares.""" + """What ``launch_state`` builds once at attach: views of the model's + weights, never a copy, and the error words every kernel of this rank + shares.""" weights: dict[int, torch.Tensor] words: torch.Tensor @@ -126,7 +127,7 @@ def skip_step(ctx: MonoContext, state: LaunchState, step: MonoStep) -> str | Non return "prefill" if step.attn_metadata.get("prefill") else None -def pack(model, ctx: MonoContext) -> LaunchState: +def launch_state(model, ctx: MonoContext) -> LaunchState: _, world, _ = _tp() return LaunchState( weights={i: layer.weight.data for i, layer in enumerate(model.layers)}, @@ -134,7 +135,7 @@ def pack(model, ctx: MonoContext) -> LaunchState: ) -def repack(model, ctx: MonoContext, state: LaunchState) -> None: +def reload_weights(model, ctx: MonoContext, state: LaunchState) -> None: for i, layer in enumerate(model.layers): state.weights[i].copy_(layer.weight.data) @@ -170,8 +171,8 @@ def _tp(): refuse_model=refuse_model, skip_layer=skip_layer, skip_step=skip_step, - pack=pack, - repack=repack, + launch_state=launch_state, + reload_weights=reload_weights, bind=bind, build=build, ) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index c103e1744f14..e0799a8f8ca1 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -88,7 +88,7 @@ def test_kernels_run_the_layers_they_replace(): assert [k.launched for k in mono.kernels.values() if k.launched] == [[width]] * 3 -def test_repack_writes_in_place(): +def test_reload_weights_writes_in_place(): config = reference.make_config() model = reference.make_model(config) mono = attach(model, config, SPECS_PATH) @@ -96,18 +96,18 @@ def test_repack_writes_in_place(): with torch.no_grad(): for layer in model.layers: layer.weight.copy_(torch.full_like(layer.weight, 7)) - mono.repack(model) + mono.reload_weights(model) assert {i: w.data_ptr() for i, w in mono.state.weights.items()} == before assert all((w == 7).all() for w in mono.state.weights.values()) -def test_repack_is_refused_when_the_spec_declares_none(): +def test_reload_weights_is_refused_when_the_spec_declares_none(): config = reference.make_config() model = reference.make_model(config) mono = attach(model, config, SPECS_PATH) - mono.spec = replace(mono.spec, repack=None) + mono.spec = replace(mono.spec, reload_weights=None) with pytest.raises(MonoUnsupported, match="weights cannot move"): - mono.repack(model) + mono.reload_weights(model) def test_model_refusal_is_raised(): diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index 3fd532778430..21ab9d2932b0 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -28,7 +28,7 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: - """Picks the variant, packs the weights and builds this rank's layer + """Picks the variant, builds the launch state and this rank's layer kernels. Collective over the TP group. Args: diff --git a/vllm/models/common/mono/routing.py b/vllm/models/common/mono/routing.py index 05b6f58cee85..eabb146ed8d8 100644 --- a/vllm/models/common/mono/routing.py +++ b/vllm/models/common/mono/routing.py @@ -83,7 +83,7 @@ def __init__(self, spec: MonoSpec, model, vllm_config: "VllmConfig"): if not indices: self.cache_layers: list[str] = [] return - self.state = spec.pack(model, self.ctx) + self.state = spec.launch_state(model, self.ctx) for i in indices: for w in self.widths: self.kernels[i, w] = spec.build( @@ -127,20 +127,20 @@ def _register_warmup(self) -> None: if owners: MonoWarmup(self.spec.name, owners).register_warmup() - def repack(self, model) -> None: - """Writes ``model``'s current weights into the allocation ``pack`` - returned, in place. + def reload_weights(self, model) -> None: + """Writes ``model``'s current weights into the allocation + ``launch_state`` returned, in place. Raises: - MonoUnsupported: The spec declares no ``repack``, so its weights - cannot move after attach. + MonoUnsupported: The spec declares no ``reload_weights``, so its + weights cannot move after attach. """ - if self.spec.repack is None: + if self.spec.reload_weights is None: raise MonoUnsupported( f"MonoKernel {self.spec.name}: weights cannot move after attach" ) - self.spec.repack(model, self.ctx, self.state) + self.spec.reload_weights(model, self.ctx, self.state) def _bind(self) -> None: """Re-checks the caches whenever they moved (keyed on every tensor's diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 009f32e9f845..0f49ea7fd34e 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -133,7 +133,7 @@ class MonoSpec: that is absent from it refuses the deployment. caches: The caches the kernels read, by attention layer name. weight_source: ``"views"`` to read vLLM's post-load layout in place, - ``"convert"`` to pack a second layout once after load. + ``"convert"`` to build a second layout once after load. layers: The decoder layers that get a kernel, from the model's config. widths: The rows a launch serves, from the engine's config. These become the CUDA graph capture sizes. @@ -144,12 +144,13 @@ class MonoSpec: kernel's. skip_step: Why one step takes vLLM's path, None when it takes the kernel's. - pack: The weights and launch state, built once at attach. The only - place a weight copy is ever made. Opaque to the library, which + launch_state: The weights and launch state, built once at attach. The + only place a weight copy is ever made. Opaque to the library, which only hands it back to the spec's own callbacks. - repack: Writes new weights into the allocation ``pack`` returned, - in place, because the kernels and the captured graphs hold its - address. None refuses anything that moves weights after attach. + reload_weights: Writes new weights into the allocation + ``launch_state`` returned, in place, because the kernels and the + captured graphs hold its address. None refuses anything that moves + weights after attach. bind: Checks the KV caches the kernels will read, on every re-bind. build: One ``LayerKernel`` for (context, state, layer, width). @@ -171,7 +172,7 @@ class MonoSpec: skip_step: Callable[[MonoContext, Any, MonoStep], str | None] # Callbacks - pack: Callable[[Any, MonoContext], Any] - repack: Callable[[Any, MonoContext, Any], None] | None + launch_state: Callable[[Any, MonoContext], Any] + reload_weights: Callable[[Any, MonoContext, Any], None] | None bind: Callable[[MonoContext, Mapping[str, Any]], list[str]] build: Callable[[MonoContext, Any, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index fbf1fc38de70..2dd5a43fcaf2 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -394,7 +394,7 @@ def _skip_step( return None -def _pack(model, ctx: MonoContext) -> LaunchState: +def _launch_state(model, ctx: MonoContext) -> LaunchState: """This rank's runner, collective over the TP group.""" from vllm.models.common.mono.peer import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer @@ -518,8 +518,8 @@ def _layers(config) -> list[int]: refuse_model=_refuse_model, skip_layer=_skip_layer, skip_step=_skip_step, - pack=_pack, - repack=None, + launch_state=_launch_state, + reload_weights=None, bind=_bind, build=_build, ) From b728e0ae4b80b87f7c8151aa0024cf92c4e31060 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 10:39:57 -0500 Subject: [PATCH 07/79] fix: improve type annotations Signed-off-by: simondanielsson --- tests/models/common/mono/reference.py | 2 +- vllm/models/common/mono/guards.py | 14 +++++++------- vllm/models/common/mono/routing.py | 2 +- vllm/models/common/mono/spec.py | 22 ++++++++++++++++------ vllm/models/deepseek_v41/amd/mono_spec.py | 2 +- 5 files changed, 26 insertions(+), 16 deletions(-) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 6201a4aa8445..d29818b3b0c7 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -160,7 +160,7 @@ def _tp(): return tp_group() -REFERENCE = MonoSpec( +REFERENCE = MonoSpec[LaunchState]( name="reference", supports=frozenset(), caches={"attn": CACHE}, diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index a03d6ae91449..24e66fecd833 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -11,7 +11,7 @@ import importlib from collections.abc import Callable, Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger @@ -36,7 +36,7 @@ def _enabled(vllm_config: "VllmConfig") -> bool: return bool(vllm_config.kernel_config.monokernel) -def _load_specs(path: str) -> tuple[MonoSpec, ...]: +def _load_specs(path: str) -> tuple[MonoSpec[Any], ...]: """``"package.module:NAME"`` -> the spec variants, importing only that module. ``NAME`` is a ``MonoSpec`` or a tuple of them, in the order they are preferred.""" @@ -71,7 +71,7 @@ def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: return {f for f, used in on.items() if used} -def _refuse_config(spec: MonoSpec, vllm_config: "VllmConfig") -> list[str]: +def _refuse_config(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> list[str]: """Every reason this config cannot run ``spec``: default-deny on features, then the runner and the graph mode, then the spec's own.""" reasons = [ @@ -86,7 +86,7 @@ def _refuse_config(spec: MonoSpec, vllm_config: "VllmConfig") -> list[str]: return reasons + list(spec.refuse_config(vllm_config)) -def _refuse_model(spec: MonoSpec, model) -> list[str]: +def _refuse_model(spec: MonoSpec[Any], model) -> list[str]: """Every reason the loaded model cannot run ``spec``. A MonoKernel replaces the layer loop, which a traced graph has already inlined, so a model carrying ``@support_torch_compile`` is refused.""" @@ -102,8 +102,8 @@ def _refuse_model(spec: MonoSpec, model) -> list[str]: def _select_spec( - specs: Sequence[MonoSpec], vllm_config: "VllmConfig" -) -> tuple[MonoSpec | None, dict[str, list[str]]]: + specs: Sequence[MonoSpec[Any]], vllm_config: "VllmConfig" +) -> tuple[MonoSpec[Any] | None, dict[str, list[str]]]: """The first variant this config can serve, and every variant's refusals. Returns: @@ -127,7 +127,7 @@ def _refusal_message(refusals: dict[str, list[str]]) -> str: ) -def _set_capture_sizes(spec: MonoSpec, vllm_config: "VllmConfig") -> None: +def _set_capture_sizes(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> None: """The capture sizes up to the largest width become exactly the widths, so a replay never pads a request's rows; larger sizes stay as configured (and capture vLLM's path).""" diff --git a/vllm/models/common/mono/routing.py b/vllm/models/common/mono/routing.py index eabb146ed8d8..e81b83e5e2fa 100644 --- a/vllm/models/common/mono/routing.py +++ b/vllm/models/common/mono/routing.py @@ -61,7 +61,7 @@ def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: class Mono: """A model's MonoKernels on this rank: built once after load.""" - def __init__(self, spec: MonoSpec, model, vllm_config: "VllmConfig"): + def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): rank, _, cpu_group = tp_group() self.spec = spec _agree_across_ranks( diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 0f49ea7fd34e..ee6897f4cd30 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Literal, Protocol +from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, TypeVar import torch @@ -35,6 +35,11 @@ class Feature(str, Enum): WEIGHT_UPDATES = "weight_updates" +LaunchStateT = TypeVar("LaunchStateT") +"""Whatever one spec's ``launch_state`` builds. The library never looks inside +it, so a spec chooses the type and the callbacks that receive it agree.""" + + class MonoUnsupported(Exception): """The configuration or the loaded model is outside what a spec serves.""" @@ -112,7 +117,7 @@ def compile(self, key: Any) -> None: @dataclass(frozen=True) -class MonoSpec: +class MonoSpec(Generic[LaunchStateT]): """One (model, vendor, ISA)'s declaration: values and callbacks, no behaviour of its own. @@ -120,6 +125,11 @@ class MonoSpec: at module scope rather than a subclass, and a model exports its candidates as a ``MONO_SPECS`` tuple the library selects from. + The parameter is the spec's own launch state: ``launch_state`` builds it, + and ``reload_weights``, ``skip_step`` and ``build`` are handed back the + same object. Annotate the spec with it (``MonoSpec[MyLaunchState](...)``) + and those four callbacks are checked against each other. + Gating runs in tiers, and the prefix is the consequence. ``refuse_*`` says the deployment asked for something it cannot have, so it may take the engine down; ``skip_*`` says only that this layer, or this step, takes @@ -169,10 +179,10 @@ class MonoSpec: refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] skip_layer: Callable[[MonoContext, Any], str | None] - skip_step: Callable[[MonoContext, Any, MonoStep], str | None] + skip_step: Callable[[MonoContext, LaunchStateT, MonoStep], str | None] # Callbacks - launch_state: Callable[[Any, MonoContext], Any] - reload_weights: Callable[[Any, MonoContext, Any], None] | None + launch_state: Callable[[Any, MonoContext], LaunchStateT] + reload_weights: Callable[[Any, MonoContext, LaunchStateT], None] | None bind: Callable[[MonoContext, Mapping[str, Any]], list[str]] - build: Callable[[MonoContext, Any, Any, int], LayerKernel] + build: Callable[[MonoContext, LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 2dd5a43fcaf2..e045b87c04a7 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -507,7 +507,7 @@ def _layers(config) -> list[int]: return list(range(config.num_hidden_layers)) -MONO_SPEC = MonoSpec( +MONO_SPEC = MonoSpec[LaunchState]( name="dsv41-cdna4", supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), caches={"swa": CacheReq("swa", "fp8_ds_mla", "blocks,block,bytes")}, From 78554c8293662051838316f3ed5022eda8a25472 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 11:20:40 -0500 Subject: [PATCH 08/79] fix: rename Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/model.py | 4 ++-- vllm/models/deepseek_v41/amd/mono_spec.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 79fd3d906faa..dbcbe6a6b963 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -61,7 +61,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v41.amd.mono_spec import SPECS_PATH, Kernel +from vllm.models.deepseek_v41.amd.mono_spec import SPECS_PATH, LayerKernelBase from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers @@ -280,7 +280,7 @@ def forward( mono: Mono | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: kernel = typing.cast( - Kernel | None, kernel_for(mono, self.attn.layer_id, positions) + LayerKernelBase | None, kernel_for(mono, self.attn.layer_id, positions) ) step = None if mono is None else mono.step if kernel is not None and step is not None and not kernel.ffn_only: diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index e045b87c04a7..8911738dc948 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -116,8 +116,9 @@ class LaunchState: need_swa: bool = False -class Kernel: - """What both launches share: the layer, its width and its weight views.""" +class LayerKernelBase: + """The ``LayerKernel`` both launches implement: the layer, its width and + its weight views, with the launch shape left to the subclass.""" ffn_only = False @@ -143,7 +144,7 @@ def launch(self, step: MonoStep, *args, **kwargs) -> Any: raise NotImplementedError -class WholeLayer(Kernel): +class WholeLayer(LayerKernelBase): """The whole backbone layer as K1 and K2, from the layer's own inputs.""" def __init__( @@ -233,7 +234,7 @@ def launch( ) -class FfnSeam(Kernel): +class FfnSeam(LayerKernelBase): """The FFN seam, MoE and both reductions as one launch, from wo_b's unreduced output after vLLM's attention.""" @@ -479,7 +480,9 @@ def _bind(ctx: MonoContext, caches) -> list[str]: return reasons -def _build(ctx: MonoContext, launch_state: LaunchState, layer, width: int): +def _build( + ctx: MonoContext, launch_state: LaunchState, layer, width: int +) -> LayerKernelBase: world = ctx.vllm_config.parallel_config.tensor_parallel_size why = _ffn_only(layer) if layer.attn.layer_id not in launch_state.weights: From 8f13e97b8608f750660cde34ed3adb6bdf3cd6d3 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 11:52:48 -0500 Subject: [PATCH 09/79] fix: rename to runtime Signed-off-by: simondanielsson --- vllm/models/common/mono/__init__.py | 2 +- vllm/models/common/mono/layer.py | 16 ++++++++-------- .../common/mono/{routing.py => runtime.py} | 10 ++++++---- vllm/models/deepseek_v41/amd/model.py | 6 +++--- 4 files changed, 18 insertions(+), 16 deletions(-) rename vllm/models/common/mono/{routing.py => runtime.py} (95%) diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index 039270dfaeef..b433157d90ed 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -3,7 +3,7 @@ """MonoKernels (RFC #60904): persistent kernels that run most or all of one decoder layer's decode step in a single launch. A model brings a ``MonoSpec`` (``spec``); this package owns the rest: refusals (``guards``), the per-step -decision and cache re-binding (``routing``), the layer-loop hook (``layer``), +decision and cache re-binding (``runtime``), the layer-loop hook (``layer``), fail-stop and metrics (``failstop``) and peer memory (``peer``).""" from vllm.models.common.mono.guards import configure diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index 21ab9d2932b0..b8f4139fa9bb 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -16,7 +16,7 @@ _refusal_message, _select_spec, ) -from vllm.models.common.mono.routing import Mono +from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.mono.spec import LayerKernel, MonoUnsupported if TYPE_CHECKING: @@ -24,10 +24,10 @@ logger = init_logger(__name__) -_ATTACHED: Mono | None = None +_ATTACHED: MonoRuntime | None = None -def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: +def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | None: """Picks the variant, builds the launch state and this rank's layer kernels. Collective over the TP group. @@ -37,7 +37,7 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: specs_path: ``"package.module:MONO_SPECS"``. Returns: - This rank's ``Mono``, or None with the switch off or when no layer + This rank's ``MonoRuntime``, or None with the switch off or when no layer takes a kernel. Raises: @@ -51,7 +51,7 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) if spec is None: raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) - mono = Mono(spec, model, vllm_config) + mono = MonoRuntime(spec, model, vllm_config) if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) return None @@ -65,7 +65,7 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> Mono | None: return mono -def kernel_for(mono: Mono | None, i: int, positions) -> LayerKernel | None: +def kernel_for(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | None: """Layer ``i``'s kernel when this step runs on the kernels, else None. The first call under a forward context decides for every layer. A model @@ -81,7 +81,7 @@ def kernel_for(mono: Mono | None, i: int, positions) -> LayerKernel | None: return mono.kernels.get((i, step.width)) -def forward_layer(mono: Mono | None, i: int, layer, positions, *args, **kwargs): +def forward_layer(mono: MonoRuntime | None, i: int, layer, positions, *args, **kwargs): """Layer ``i`` on its kernel when this step runs on the kernels, else the vLLM layer unchanged.""" k = kernel_for(mono, i, positions) @@ -93,7 +93,7 @@ def forward_layer(mono: Mono | None, i: int, layer, positions, *args, **kwargs): return layer(positions, *args, **kwargs) -def health(mono: Mono | None) -> torch.Tensor | None: +def health(mono: MonoRuntime | None) -> torch.Tensor | None: """The step's error words, read by the runner after each step: any set word fails the step before its tokens are returned (``failstop``).""" return None if mono is None else mono.health() diff --git a/vllm/models/common/mono/routing.py b/vllm/models/common/mono/runtime.py similarity index 95% rename from vllm/models/common/mono/routing.py rename to vllm/models/common/mono/runtime.py index e81b83e5e2fa..057a0c4fbbb1 100644 --- a/vllm/models/common/mono/routing.py +++ b/vllm/models/common/mono/runtime.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The per-step decision and cache re-binding. +"""The live side of a spec: one ``MonoRuntime`` a rank, built at attach, owning +this rank's layer kernels and making the per-step decision and cache re-binding. +It launches nothing itself. -``Mono.begin_step`` decides once a forward, for every kernel layer, from what -every rank sees alike: the batch's width and the attention metadata, never +``MonoRuntime.begin_step`` decides once a forward, for every kernel layer, from +what every rank sees alike: the batch's width and the attention metadata, never rank-local state and never a collective. Under a FULL CUDA graph the decision is baked in at capture, so it is a function of the captured size: FULL graphs serve uniform decode batches and every width is a capture size @@ -58,7 +60,7 @@ def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: self._compiled_cache[compile_key] = True -class Mono: +class MonoRuntime: """A model's MonoKernels on this rank: built once after load.""" def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index dbcbe6a6b963..7125d693e314 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -51,7 +51,7 @@ maybe_prefix, ) from vllm.models.common.mono import attach, configure, health, kernel_for -from vllm.models.common.mono.routing import Mono +from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -277,7 +277,7 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - mono: Mono | None = None, + mono: MonoRuntime | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: kernel = typing.cast( LayerKernelBase | None, kernel_for(mono, self.attn.layer_id, positions) @@ -484,7 +484,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps self.vllm_config = vllm_config - self.mono: Mono | None = None + self.mono: MonoRuntime | None = None configure(vllm_config, SPECS_PATH) # Three aux streams: one per non-default input GEMM in From cbf4a5666544a416cd68057f0f2c5a1fc948a026 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 15:00:06 -0500 Subject: [PATCH 10/79] fix: simplify Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 23 +++------- tests/models/common/mono/reference.py | 26 +---------- vllm/models/common/mono/__init__.py | 10 +++-- vllm/models/common/mono/failstop.py | 55 ----------------------- vllm/models/common/mono/layer.py | 23 ++-------- vllm/models/common/mono/runtime.py | 24 +++++++--- vllm/models/common/mono/spec.py | 4 -- vllm/models/deepseek_v41/amd/model.py | 8 +++- vllm/models/deepseek_v41/amd/mono_spec.py | 3 -- 9 files changed, 44 insertions(+), 132 deletions(-) delete mode 100644 vllm/models/common/mono/failstop.py diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index a34f945791d5..01cb2827b05e 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -7,8 +7,8 @@ are complete, the spec module is importable without a kernel backend, the capture sizes are the spec's widths, one forward decides once and alike on every rank, caches re-bind when they move, attach does not move the model's -weights, cache checks report every mismatch, a failed step fails, served steps -are counted, and every width is warmed up before capture. +weights, cache checks report every mismatch, served steps are counted, and +every width is warmed up before capture. The kit never launches a kernel: it drives the library through the spec's own callbacks, so a spec whose kernels need a GPU still conforms on CPU. @@ -28,7 +28,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry -from vllm.models.common.mono import failstop +from vllm.models.common.mono import runtime from vllm.models.common.mono.guards import ( _load_specs, _refuse_config, @@ -110,9 +110,9 @@ def mono(self, config, model): @pytest.fixture(autouse=True) def _counters(self): - failstop.STEPS.clear() + runtime.STEPS.clear() yield - failstop.STEPS.clear() + runtime.STEPS.clear() def spec(self, config): spec, refusals = _select_spec(_load_specs(self.specs_path), config) @@ -230,29 +230,20 @@ def test_attach_does_not_move_weights(self, config): after = {n: p.data_ptr() for n, p in model.named_parameters()} assert before == after - def test_fail_stop_raises(self, mono): - words = torch.zeros_like(mono.health()) - words[0] = 1 - with pytest.raises(RuntimeError, match="MonoKernel step failed"): - failstop.check(words) - - def test_healthy_step_passes(self, mono): - failstop.check(torch.zeros_like(mono.health())) - def test_counters_follow_full_replay(self, config, model, mono): width = min(mono.widths) with forward(config, CUDAGraphMode.FULL, self.decode_metadata(config, model)): mono.begin_step(torch.zeros(width, dtype=torch.long)) for _ in range(3): mono.account(width, CUDAGraphMode.FULL) - assert failstop.drain() == {f"steps:{mono.spec.name}:mono": 3} + assert runtime.drain() == {f"steps:{mono.spec.name}:mono": 3} def test_counters_name_the_skipped_path(self, config, model, mono): width = min(mono.widths) with forward(config, CUDAGraphMode.FULL, self.skipped_metadata(config, model)): why = mono.begin_step(torch.zeros(width, dtype=torch.long)) mono.account(width, CUDAGraphMode.FULL) - assert failstop.drain() == {f"steps:{mono.spec.name}:{why}": 1} + assert runtime.drain() == {f"steps:{mono.spec.name}:{why}": 1} def test_warmup_registered_for_every_width(self, config): registry = JitWarmupRegistry(config) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index d29818b3b0c7..4e28686a26e0 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -51,20 +51,13 @@ def __init__(self, layers: int = 4, hidden: int = 8): self.end_layer = layers self.mono = None - def step_health(self) -> torch.Tensor | None: - from vllm.models.common.mono import health - - return health(self.mono) - @dataclass class LaunchState: """What ``launch_state`` builds once at attach: views of the model's - weights, never a copy, and the error words every kernel of this rank - shares.""" + weights, never a copy.""" weights: dict[int, torch.Tensor] - words: torch.Tensor class ReferenceKernel: @@ -83,19 +76,12 @@ def launch(self, step: MonoStep, positions: torch.Tensor, *args, **kwargs): self.launched.append(step.num_tokens) return self.layer(positions, *args, **kwargs) - def health(self) -> torch.Tensor: - return self.state.words - def warmup_keys(self) -> list[Any]: return [(self.layer.index, self.width)] def compile(self, key: Any) -> None: self.compiled.append(key) - def fail(self, rank: int = 0) -> None: - """Reports an in-kernel wait that gave up, as a launch would.""" - self.state.words[rank] = 1 - def layers(config) -> list[int]: return list(range(config.num_hidden_layers)) @@ -128,10 +114,8 @@ def skip_step(ctx: MonoContext, state: LaunchState, step: MonoStep) -> str | Non def launch_state(model, ctx: MonoContext) -> LaunchState: - _, world, _ = _tp() return LaunchState( - weights={i: layer.weight.data for i, layer in enumerate(model.layers)}, - words=torch.zeros(world, dtype=torch.int32, device=ctx.device), + weights={i: layer.weight.data for i, layer in enumerate(model.layers)} ) @@ -154,12 +138,6 @@ def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceK return ReferenceKernel(ctx, state, layer, width) -def _tp(): - from vllm.models.common.mono.peer import tp_group - - return tp_group() - - REFERENCE = MonoSpec[LaunchState]( name="reference", supports=frozenset(), diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index b433157d90ed..dfca170d9045 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -3,15 +3,18 @@ """MonoKernels (RFC #60904): persistent kernels that run most or all of one decoder layer's decode step in a single launch. A model brings a ``MonoSpec`` (``spec``); this package owns the rest: refusals (``guards``), the per-step -decision and cache re-binding (``runtime``), the layer-loop hook (``layer``), -fail-stop and metrics (``failstop``) and peer memory (``peer``).""" +decision, cache re-binding and step counters (``runtime``), the layer-loop hook +(``layer``) and peer memory (``peer``). + +Detecting and reporting an in-kernel failure is the model's own: only its +kernels know what they record, so a model that needs fail-stop reads it from +the launch state it built.""" from vllm.models.common.mono.guards import configure from vllm.models.common.mono.layer import ( account_step, attach, forward_layer, - health, kernel_for, ) @@ -20,6 +23,5 @@ "attach", "configure", "forward_layer", - "health", "kernel_for", ] diff --git a/vllm/models/common/mono/failstop.py b/vllm/models/common/mono/failstop.py deleted file mode 100644 index 29d4dedd3c2d..000000000000 --- a/vllm/models/common/mono/failstop.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fail-stop and metrics. - -A MonoKernel never returns tokens from a step it did not finish: the runner -copies the model's error words next to the sampled token ids, and -``check`` raises before that step's tokens are returned. The engine then stops -like on any other engine death, and the deployment's restart policy recovers -it. - -Counters (``drain``, exported as ``vllm:monokernel_steps``): steps on the -kernels, and steps on vLLM's path by reason. -""" - -from collections import Counter - -import torch - -STEPS: Counter[tuple[str, str]] = Counter() - - -def count_step(spec: str, why: str | None) -> None: - """Counts one served step of ``spec``: on the kernels when ``why`` is None, - else on vLLM's layers for that reason.""" - STEPS[spec, why or "mono"] += 1 - - -def drain() -> dict[str, int] | None: - """The counts since the last drain: ``"steps::"``, where path - is "mono" or the reason the step ran vLLM's layers.""" - if not STEPS: - return None - out = {f"steps:{spec}:{path}": n for (spec, path), n in STEPS.items()} - STEPS.clear() - return out - - -def check(words: torch.Tensor) -> None: - """Fails the step whose error words these are, before its tokens return. - - Args: - words: The output rank's error words for one step, on the host (a word - a TP rank, nonzero once one of its in-kernel waits gave up). - - Raises: - RuntimeError: A rank gave up a wait. - - """ - if not words.any(): - return - ranks = [r for r, w in enumerate(words.tolist()) if w] - raise RuntimeError( - f"MonoKernel step failed: TP ranks {ranks} gave up an in-kernel wait; " - "the step's tokens are discarded" - ) diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index b8f4139fa9bb..487187c769d9 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -1,22 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The hooks a model calls: ``attach`` after load, ``forward_layer`` in its -layer loop, ``health`` from ``step_health``.""" +layer loop.""" from typing import TYPE_CHECKING -import torch - from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger -from vllm.models.common.mono import failstop from vllm.models.common.mono.guards import ( _enabled, _load_specs, _refusal_message, _select_spec, ) -from vllm.models.common.mono.runtime import MonoRuntime +from vllm.models.common.mono.runtime import MonoRuntime, drain from vllm.models.common.mono.spec import LayerKernel, MonoUnsupported if TYPE_CHECKING: @@ -55,12 +52,6 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | N if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) return None - if not hasattr(model, "step_health"): - logger.warning_once( - "MonoKernel %s: the model exposes no step_health(), so the runner " - "cannot fail a step whose in-kernel wait gave up", - spec.name, - ) _ATTACHED = mono return mono @@ -93,16 +84,10 @@ def forward_layer(mono: MonoRuntime | None, i: int, layer, positions, *args, **k return layer(positions, *args, **kwargs) -def health(mono: MonoRuntime | None) -> torch.Tensor | None: - """The step's error words, read by the runner after each step: any set - word fails the step before its tokens are returned (``failstop``).""" - return None if mono is None else mono.health() - - def account_step(num_tokens: int, runtime_mode: CUDAGraphMode) -> dict[str, int] | None: """Called by the runner after every served step: counts it, and returns the - counters to export since the last call (``failstop.drain``).""" + counters to export since the last call (``runtime.drain``).""" if _ATTACHED is None: return None _ATTACHED.account(num_tokens, runtime_mode) - return failstop.drain() + return drain() diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 057a0c4fbbb1..bb809b4a04d7 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -11,8 +11,13 @@ serve uniform decode batches and every width is a capture size (``guards``), so no replay pads a request's rows. PIECEWISE steps run vLLM's layers. + +It also counts the steps it decided (``STEPS``, drained by ``drain`` and +exported as ``vllm:monokernel_steps``): steps on the kernels, and steps on +vLLM's path by reason. """ +from collections import Counter from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -22,7 +27,6 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel -from vllm.models.common.mono import failstop from vllm.models.common.mono.guards import _agree_across_ranks, _refuse_model from vllm.models.common.mono.peer import tp_group from vllm.models.common.mono.spec import ( @@ -38,6 +42,19 @@ logger = init_logger(__name__) +STEPS: Counter[tuple[str, str]] = Counter() +"""Served steps by spec and path, where path is "mono" or the reason the step +ran vLLM's layers.""" + + +def drain() -> dict[str, int] | None: + """The counts since the last drain, keyed ``"steps::"``.""" + if not STEPS: + return None + out = {f"steps:{spec}:{path}": n for (spec, path), n in STEPS.items()} + STEPS.clear() + return out + class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): """A spec's layer-kernel builds as vLLM JIT warm-up keys.""" @@ -191,7 +208,7 @@ def account(self, num_tokens: int, runtime_mode: CUDAGraphMode) -> None: why = self.baked.get(num_tokens, "not captured") else: why, self.last = self.last, "no step" - failstop.count_step(self.spec.name, why) + STEPS[self.spec.name, why or "mono"] += 1 def _decide(self, positions: torch.Tensor) -> str | None: fc = get_forward_context() @@ -206,6 +223,3 @@ def _decide(self, positions: torch.Tensor) -> str | None: self._bind() self.step = step return why - - def health(self) -> torch.Tensor: - return next(iter(self.kernels.values())).health() diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index ee6897f4cd30..3c8e82936c87 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -102,10 +102,6 @@ def launch(self, step: MonoStep, *args, **kwargs) -> Any: tensors the layer was handed.""" ... - def health(self) -> torch.Tensor: - """This rank's error words: any set word fails the step.""" - ... - def warmup_keys(self) -> list[Any]: """Hashable build keys of this width's launches, compiled before CUDA graph capture by vLLM's JIT warm-up.""" diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 7125d693e314..fa06e09e136f 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -50,7 +50,7 @@ make_layers, maybe_prefix, ) -from vllm.models.common.mono import attach, configure, health, kernel_for +from vllm.models.common.mono import attach, configure, kernel_for from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, @@ -1420,7 +1420,11 @@ def process_weights_after_loading(self) -> None: self.model.mono = attach(self.model, self.model.vllm_config, SPECS_PATH) def step_health(self) -> torch.Tensor | None: - return health(self.model.mono) + """This rank's error words, one a TP rank, nonzero once one of its + in-kernel waits gave up: any set word fails the step before its tokens + are returned.""" + mono = self.model.mono + return None if mono is None else mono.state.words def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 8911738dc948..cf87f495bd93 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -131,9 +131,6 @@ def __init__( self.weights = launch_state.weights[layer.attn.layer_id] self.cache_layers: tuple[str, ...] = () - def health(self) -> torch.Tensor: - return self.launch_state.words - def warmup_keys(self) -> list[Any]: raise NotImplementedError From 01dd055dd54460162068781a098e868abd7cbdd7 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 15:11:14 -0500 Subject: [PATCH 11/79] fix: trim morre Signed-off-by: simondanielsson --- vllm/models/common/mono/__init__.py | 12 +- vllm/models/common/mono/guards.py | 42 ++++++- vllm/models/common/mono/peer.py | 140 ---------------------- vllm/models/common/mono/runtime.py | 7 +- vllm/models/deepseek_v41/amd/mono_spec.py | 2 +- 5 files changed, 53 insertions(+), 150 deletions(-) delete mode 100644 vllm/models/common/mono/peer.py diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index dfca170d9045..10c0349ba44b 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -2,13 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MonoKernels (RFC #60904): persistent kernels that run most or all of one decoder layer's decode step in a single launch. A model brings a ``MonoSpec`` -(``spec``); this package owns the rest: refusals (``guards``), the per-step -decision, cache re-binding and step counters (``runtime``), the layer-loop hook -(``layer``) and peer memory (``peer``). +(``spec``); this package owns the rest: refusals and TP agreement (``guards``), +the per-step decision, cache re-binding and step counters (``runtime``), and the +layer-loop hook (``layer``). -Detecting and reporting an in-kernel failure is the model's own: only its -kernels know what they record, so a model that needs fail-stop reads it from -the launch state it built.""" +Peer memory and in-kernel failure reporting are the model's own: only its +kernels know what they put where and what they record, so a model that needs +either builds it in the launch state it brings.""" from vllm.models.common.mono.guards import configure from vllm.models.common.mono.layer import ( diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index 24e66fecd833..389e17525767 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -7,15 +7,22 @@ library's own and stays private to this package: ``_agree_across_ranks`` runs a later refusal stage (after load) on every TP rank and agrees on the verdict before anything collective follows. + +Peer memory is the model's own: only its kernels know what they put where, and +``tp_group`` is where it reads this rank's place in the TP group. What this +package keeps is the agreement every stage goes through (``tp_agree``), because +a rank that refuses alone leaves every other rank waiting in-kernel for ever. """ import importlib from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any +import torch +import torch.distributed as dist + from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger -from vllm.models.common.mono.peer import tp_agree from vllm.models.common.mono.spec import Feature, MonoSpec, MonoUnsupported if TYPE_CHECKING: @@ -31,6 +38,39 @@ ) +def tp_group() -> tuple[int, int, object | None]: + """This rank's (rank, world size, CPU group) in the TP group, and + (0, 1, None) before distributed is initialized.""" + if not dist.is_available() or not dist.is_initialized(): + return 0, 1, None + from vllm.distributed import get_tp_group + + tp = get_tp_group() + return tp.rank_in_group, tp.world_size, tp.cpu_group + + +def tp_agree(ok: bool, group) -> bool: + """Agrees a verdict every TP rank must take alike. + + A MonoKernel waits in-kernel on every peer rank, so a rank that turned mono + off alone, or failed to build what the others launch, leaves the others + spinning on the GPU for ever. Every rank must call this at the same point. + + Args: + ok: This rank's local verdict. + group: The TP group's CPU (gloo) process group, or None for one rank. + + Returns: + True when every rank of ``group`` passed ``ok`` true. + + """ + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + def _enabled(vllm_config: "VllmConfig") -> bool: """Whether ``--monokernel`` asked for the model's MonoKernels.""" return bool(vllm_config.kernel_config.monokernel) diff --git a/vllm/models/common/mono/peer.py b/vllm/models/common/mono/peer.py deleted file mode 100644 index d062f779ee3c..000000000000 --- a/vllm/models/common/mono/peer.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -# Copyright (C) 2026, Advanced Micro Devices, Inc.: -# atom/mono/runtime/peer_memory.py -# atom/mono/runtime/consensus.py -"""Peer memory for in-kernel TP collectives: one uncached buffer a rank, every -rank holding every peer's address, and the TP agreements that must precede -anything collective.""" - -import torch -import torch.distributed as dist - -from vllm.models.common.mono.spec import MonoUnsupported - - -def tp_group() -> tuple[int, int, object | None]: - """This rank's (rank, world size, CPU group) in the TP group, and - (0, 1, None) before distributed is initialized.""" - if not dist.is_available() or not dist.is_initialized(): - return 0, 1, None - from vllm.distributed import get_tp_group - - tp = get_tp_group() - return tp.rank_in_group, tp.world_size, tp.cpu_group - - -class _DeviceBytes: - """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, - for a non-owning torch view.""" - - def __init__(self, ptr: int, nbytes: int): - self.__cuda_array_interface__ = { - "shape": (nbytes,), - "typestr": "|u1", - "data": (ptr, False), - "version": 3, - } - - -class PeerBuffer: - """One uncached buffer per rank, every rank holding every peer's address - (``addresses``, int64, indexed by TP rank). What lives where, and why no - rank overwrites a peer's unread data, is the model's layout's to say. - - ``bytes`` is the whole buffer: the kernels tag every pair with the step's - epoch, so it is never zeroed between steps. - - The constructor is collective over ``group`` (the handle exchange, then an - agreement that every rank mapped every peer): every rank must reach it, so - a runner reaches it only after ``tp_agree``. A failed mapping on any rank - raises ``MonoUnsupported`` on all of them. - """ - - def __init__( - self, - nbytes: int, - group, - rank: int, - npes: int, - device: torch.device, - ): - from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap - - self._heap = UncachedIpcHeap - self.local = UncachedIpcHeap.alloc_uncached(nbytes) - self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) - self.rank = rank - self._opened: list[int] = [] - addresses = [self.local] - if npes > 1: - handles = UncachedIpcHeap.gather_object_list_via_broadcast( - group, UncachedIpcHeap.get_mem_handle_bytes(self.local) - ) - addresses, failure = [], None - try: - for peer, handle in enumerate(handles): - if peer == rank: - addresses.append(self.local) - else: - base = UncachedIpcHeap.open_mem_handle(handle) - self._opened.append(base) - addresses.append(base) - except RuntimeError as err: - failure = err - if not tp_agree(failure is None, group): - self.close() - raise MonoUnsupported( - "peer memory handshake failed on " - f"{'this' if failure else 'another'}" - " TP rank" - ) from failure - self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) - - def close(self) -> None: - """Closes the peer mappings and frees the local buffer; idempotent. - Only once no kernel of any rank can still touch it (the runner is being - dropped).""" - opened, self._opened = self._opened, [] - for base in opened: - self._heap.close_mem_handle(base) - if self.local: - self.bytes = None - self._heap.free_device_mem(self.local) - self.local = 0 - - -def tp_agree(ok: bool, group) -> bool: - """Agrees a verdict every TP rank must take alike. - - A MonoKernel waits in-kernel on every peer rank, so a rank that turned mono - off alone, or failed to build what the others launch, leaves the others - spinning on the GPU for ever. Every rank must call this at the same point. - - Args: - ok: This rank's local verdict. - group: The TP group's CPU (gloo) process group, or None for one rank. - - Returns: - True when every rank of ``group`` passed ``ok`` true. - - """ - if group is None: - return ok - flag = torch.tensor([int(ok)], dtype=torch.int32) - dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) - return bool(flag.item()) - - -def bind_agreed(bind, group) -> None: - """``bind()`` on this rank (raising ``MonoUnsupported`` to refuse), then - every rank's verdict agreed: a rank that raised re-raises, the others raise - too.""" - try: - bind() - except Exception: - tp_agree(False, group) - raise - if not tp_agree(True, group): - raise MonoUnsupported("another TP rank refused mono") diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index bb809b4a04d7..1e479bf99c63 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -27,8 +27,11 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel -from vllm.models.common.mono.guards import _agree_across_ranks, _refuse_model -from vllm.models.common.mono.peer import tp_group +from vllm.models.common.mono.guards import ( + _agree_across_ranks, + _refuse_model, + tp_group, +) from vllm.models.common.mono.spec import ( LayerKernel, MonoContext, diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index cf87f495bd93..363ca9b76122 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -394,7 +394,7 @@ def _skip_step( def _launch_state(model, ctx: MonoContext) -> LaunchState: """This rank's runner, collective over the TP group.""" - from vllm.models.common.mono.peer import tp_group + from vllm.models.common.mono.guards import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS From d311c5e392bd228eca02b9bdb23ee2cf5ffaf023 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 15:14:43 -0500 Subject: [PATCH 12/79] chore: rename the shared layer kernel Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/model.py | 4 ++-- vllm/models/deepseek_v41/amd/mono_spec.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index fa06e09e136f..7feeef132cf5 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -61,7 +61,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v41.amd.mono_spec import SPECS_PATH, LayerKernelBase +from vllm.models.deepseek_v41.amd.mono_spec import SPECS_PATH, DSV41LayerKernel from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers @@ -280,7 +280,7 @@ def forward( mono: MonoRuntime | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: kernel = typing.cast( - LayerKernelBase | None, kernel_for(mono, self.attn.layer_id, positions) + DSV41LayerKernel | None, kernel_for(mono, self.attn.layer_id, positions) ) step = None if mono is None else mono.step if kernel is not None and step is not None and not kernel.ffn_only: diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 363ca9b76122..583128e726d3 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -116,7 +116,7 @@ class LaunchState: need_swa: bool = False -class LayerKernelBase: +class DSV41LayerKernel: """The ``LayerKernel`` both launches implement: the layer, its width and its weight views, with the launch shape left to the subclass.""" @@ -141,7 +141,7 @@ def launch(self, step: MonoStep, *args, **kwargs) -> Any: raise NotImplementedError -class WholeLayer(LayerKernelBase): +class WholeLayer(DSV41LayerKernel): """The whole backbone layer as K1 and K2, from the layer's own inputs.""" def __init__( @@ -231,7 +231,7 @@ def launch( ) -class FfnSeam(LayerKernelBase): +class FfnSeam(DSV41LayerKernel): """The FFN seam, MoE and both reductions as one launch, from wo_b's unreduced output after vLLM's attention.""" @@ -479,7 +479,7 @@ def _bind(ctx: MonoContext, caches) -> list[str]: def _build( ctx: MonoContext, launch_state: LaunchState, layer, width: int -) -> LayerKernelBase: +) -> DSV41LayerKernel: world = ctx.vllm_config.parallel_config.tensor_parallel_size why = _ffn_only(layer) if layer.attn.layer_id not in launch_state.weights: From 4f9b3c407a0d882366d91ec5403a4bba03ed198c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 15:23:35 -0500 Subject: [PATCH 13/79] fix: cleanup non-crucial code Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 14 +++++++- tests/models/common/mono/test_reference.py | 4 +-- tests/models/test_dsv41_mono_decode.py | 12 +++---- vllm/models/common/mono/__init__.py | 11 ++----- vllm/models/common/mono/layer.py | 38 ++++------------------ vllm/models/deepseek_v41/amd/model.py | 5 +-- 6 files changed, 32 insertions(+), 52 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 01cb2827b05e..de10449a51eb 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -36,7 +36,7 @@ _set_capture_sizes, configure, ) -from vllm.models.common.mono.layer import attach, forward_layer +from vllm.models.common.mono.layer import attach, kernel_for_layer from vllm.models.common.mono.spec import Feature, MonoUnsupported KERNEL_BACKENDS = ("aiter", "flydsl") @@ -57,6 +57,18 @@ def forward(config, mode: CUDAGraphMode, attn_metadata: Any) -> Iterator[None]: yield +def forward_layer(mono, i: int, layer, positions, *args, **kwargs): + """Layer ``i`` on its kernel when this step runs on the kernels, else the + vLLM layer unchanged: the layer loop a model writes for itself.""" + k = kernel_for_layer(mono, i, positions) + step = None if mono is None else mono.step + if k is not None and step is not None: + out = k.launch(step, positions, *args, **kwargs) + if out is not None: + return out + return layer(positions, *args, **kwargs) + + def _spy(spec, field: str, calls: list): """``spec`` with one callback recording its calls.""" inner = getattr(spec, field) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index e0799a8f8ca1..5d3e6001dba0 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -9,10 +9,10 @@ import torch from tests.models.common.mono import reference -from tests.models.common.mono.kit import MonoConformance, forward +from tests.models.common.mono.kit import MonoConformance, forward, forward_layer from vllm.config.compilation import CUDAGraphMode from vllm.models.common.mono.guards import _enabled, _select_spec -from vllm.models.common.mono.layer import attach, forward_layer +from vllm.models.common.mono.layer import attach from vllm.models.common.mono.spec import MonoUnsupported SPECS_PATH = "tests.models.common.mono.reference:MONO_SPECS" diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 30df8802d59e..82c58a565f8b 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -20,7 +20,7 @@ from tests.models.common.mono.kit import MonoConformance, forward from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -from vllm.models.common.mono.layer import attach, kernel_for +from vllm.models.common.mono.layer import attach, kernel_for_layer from vllm.models.common.mono.spec import MonoUnsupported from vllm.models.deepseek_v41.amd import mono_spec as ms @@ -273,7 +273,7 @@ def bad_caches(self, config, model): def _run_layers(self, mono, model) -> None: positions = self._positions(mono) for i, layer in enumerate(model.layers): - kernel_for(mono, i, positions) + kernel_for_layer(mono, i, positions) def _rocm_platform() -> ModuleType: @@ -474,7 +474,7 @@ def test_whole_launch_runs_the_layer(monkeypatch): mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for(mono, 1, positions) + kernel = kernel_for_layer(mono, 1, positions) out = kernel.launch( mono.step, torch.zeros(1, HIDDEN), @@ -497,7 +497,7 @@ def test_first_layer_takes_the_embedding(monkeypatch): mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for(mono, 0, positions) + kernel = kernel_for_layer(mono, 0, positions) out = kernel.launch( mono.step, torch.zeros(1, HIDDEN), positions, None, None, None, None ) @@ -509,7 +509,7 @@ def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): mono, model = _attach(monkeypatch, compressor=object()) positions = torch.zeros(1, dtype=torch.long) with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for(mono, 1, positions) + kernel = kernel_for_layer(mono, 1, positions) assert kernel.ffn_only out = kernel.launch( mono.step, @@ -528,7 +528,7 @@ def test_launch_declines_a_layout_it_cannot_read(monkeypatch): mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for(mono, 1, positions) + kernel = kernel_for_layer(mono, 1, positions) out = kernel.launch( mono.step, torch.zeros(1, 4, HIDDEN), positions, None, None, None, None ) diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index 10c0349ba44b..74b02af89f94 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -11,17 +11,10 @@ either builds it in the launch state it brings.""" from vllm.models.common.mono.guards import configure -from vllm.models.common.mono.layer import ( - account_step, - attach, - forward_layer, - kernel_for, -) +from vllm.models.common.mono.layer import attach, kernel_for_layer __all__ = [ - "account_step", "attach", "configure", - "forward_layer", - "kernel_for", + "kernel_for_layer", ] diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index 487187c769d9..57e3f6b0bb9a 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The hooks a model calls: ``attach`` after load, ``forward_layer`` in its +"""The hooks a model calls: ``attach`` after load, ``kernel_for_layer`` in its layer loop.""" from typing import TYPE_CHECKING -from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger from vllm.models.common.mono.guards import ( _enabled, @@ -13,7 +12,7 @@ _refusal_message, _select_spec, ) -from vllm.models.common.mono.runtime import MonoRuntime, drain +from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.mono.spec import LayerKernel, MonoUnsupported if TYPE_CHECKING: @@ -21,8 +20,6 @@ logger = init_logger(__name__) -_ATTACHED: MonoRuntime | None = None - def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | None: """Picks the variant, builds the launch state and this rank's layer @@ -42,7 +39,6 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | N another) refused the loaded model. """ - global _ATTACHED if not _enabled(vllm_config): return None spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) @@ -52,16 +48,15 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | N if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) return None - _ATTACHED = mono return mono -def kernel_for(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | None: +def kernel_for_layer(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | None: """Layer ``i``'s kernel when this step runs on the kernels, else None. - The first call under a forward context decides for every layer. A model - whose layer signature or dispatch site does not fit ``forward_layer`` looks - the kernel up itself and calls ``launch`` where it suits the layer. + The first call under a forward context decides for every layer. The model + calls ``launch`` where it suits the layer, and runs the vLLM layer + unchanged when this returns None. """ if mono is None: return None @@ -70,24 +65,3 @@ def kernel_for(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | Non if step is None: return None return mono.kernels.get((i, step.width)) - - -def forward_layer(mono: MonoRuntime | None, i: int, layer, positions, *args, **kwargs): - """Layer ``i`` on its kernel when this step runs on the kernels, else the - vLLM layer unchanged.""" - k = kernel_for(mono, i, positions) - step = None if mono is None else mono.step - if k is not None and step is not None: - out = k.launch(step, positions, *args, **kwargs) - if out is not None: - return out - return layer(positions, *args, **kwargs) - - -def account_step(num_tokens: int, runtime_mode: CUDAGraphMode) -> dict[str, int] | None: - """Called by the runner after every served step: counts it, and returns the - counters to export since the last call (``runtime.drain``).""" - if _ATTACHED is None: - return None - _ATTACHED.account(num_tokens, runtime_mode) - return drain() diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 7feeef132cf5..f232937377c8 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -50,7 +50,7 @@ make_layers, maybe_prefix, ) -from vllm.models.common.mono import attach, configure, kernel_for +from vllm.models.common.mono import attach, configure, kernel_for_layer from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, @@ -280,7 +280,8 @@ def forward( mono: MonoRuntime | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: kernel = typing.cast( - DSV41LayerKernel | None, kernel_for(mono, self.attn.layer_id, positions) + DSV41LayerKernel | None, + kernel_for_layer(mono, self.attn.layer_id, positions), ) step = None if mono is None else mono.step if kernel is not None and step is not None and not kernel.ffn_only: From fdecc85fa962ce9cbb5c8441fd3f2d8f1daf484c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 15:42:14 -0500 Subject: [PATCH 14/79] fix: remove step_health Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/model.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index f232937377c8..276b77472a53 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -1420,12 +1420,5 @@ def process_weights_after_loading(self) -> None: self.model.finalize_mhc_broadcast_weights() self.model.mono = attach(self.model, self.model.vllm_config, SPECS_PATH) - def step_health(self) -> torch.Tensor | None: - """This rank's error words, one a TP rank, nonzero once one of its - in-kernel waits gave up: any set word fails the step before its tokens - are returned.""" - mono = self.model.mono - return None if mono is None else mono.state.words - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() From 4f3683932de08fce14df411b0ec6430815227e05 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 15:49:07 -0500 Subject: [PATCH 15/79] fix: remove account/drain Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 26 ++----------------------- vllm/models/common/mono/runtime.py | 31 +----------------------------- 2 files changed, 3 insertions(+), 54 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index de10449a51eb..7fc5829e6e4a 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -7,8 +7,8 @@ are complete, the spec module is importable without a kernel backend, the capture sizes are the spec's widths, one forward decides once and alike on every rank, caches re-bind when they move, attach does not move the model's -weights, cache checks report every mismatch, served steps are counted, and -every width is warmed up before capture. +weights, cache checks report every mismatch, and every width is warmed up +before capture. The kit never launches a kernel: it drives the library through the spec's own callbacks, so a spec whose kernels need a GPU still conforms on CPU. @@ -28,7 +28,6 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry -from vllm.models.common.mono import runtime from vllm.models.common.mono.guards import ( _load_specs, _refuse_config, @@ -120,12 +119,6 @@ def model(self, config): def mono(self, config, model): return attach(model, config, self.specs_path) - @pytest.fixture(autouse=True) - def _counters(self): - runtime.STEPS.clear() - yield - runtime.STEPS.clear() - def spec(self, config): spec, refusals = _select_spec(_load_specs(self.specs_path), config) assert spec is not None, refusals @@ -242,21 +235,6 @@ def test_attach_does_not_move_weights(self, config): after = {n: p.data_ptr() for n, p in model.named_parameters()} assert before == after - def test_counters_follow_full_replay(self, config, model, mono): - width = min(mono.widths) - with forward(config, CUDAGraphMode.FULL, self.decode_metadata(config, model)): - mono.begin_step(torch.zeros(width, dtype=torch.long)) - for _ in range(3): - mono.account(width, CUDAGraphMode.FULL) - assert runtime.drain() == {f"steps:{mono.spec.name}:mono": 3} - - def test_counters_name_the_skipped_path(self, config, model, mono): - width = min(mono.widths) - with forward(config, CUDAGraphMode.FULL, self.skipped_metadata(config, model)): - why = mono.begin_step(torch.zeros(width, dtype=torch.long)) - mono.account(width, CUDAGraphMode.FULL) - assert runtime.drain() == {f"steps:{mono.spec.name}:{why}": 1} - def test_warmup_registered_for_every_width(self, config): registry = JitWarmupRegistry(config) with registry.activate(): diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 1e479bf99c63..f88cd81cc877 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -11,13 +11,8 @@ serve uniform decode batches and every width is a capture size (``guards``), so no replay pads a request's rows. PIECEWISE steps run vLLM's layers. - -It also counts the steps it decided (``STEPS``, drained by ``drain`` and -exported as ``vllm:monokernel_steps``): steps on the kernels, and steps on -vLLM's path by reason. """ -from collections import Counter from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -45,19 +40,6 @@ logger = init_logger(__name__) -STEPS: Counter[tuple[str, str]] = Counter() -"""Served steps by spec and path, where path is "mono" or the reason the step -ran vLLM's layers.""" - - -def drain() -> dict[str, int] | None: - """The counts since the last drain, keyed ``"steps::"``.""" - if not STEPS: - return None - out = {f"steps:{spec}:{path}": n for (spec, path), n in STEPS.items()} - STEPS.clear() - return out - class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): """A spec's layer-kernel builds as vLLM JIT warm-up keys.""" @@ -96,8 +78,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.kernels: dict[tuple[int, int], LayerKernel] = {} self.state: Any = None self.step: MonoStep | None = None - self.baked: dict[int, str | None] = {} - self.last: str | None = "no step" + self.last: str | None = None self._fc: Any = None self._bound: tuple | None = None @@ -200,19 +181,9 @@ def begin_step(self, positions: torch.Tensor) -> str | None: reason it runs vLLM's layers.""" self.step = None why = self._decide(positions) - if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: - self.baked[positions.size(-1)] = why self.last = why return why - def account(self, num_tokens: int, runtime_mode: CUDAGraphMode) -> None: - """Counts one served step, after its forward.""" - if runtime_mode == CUDAGraphMode.FULL: - why = self.baked.get(num_tokens, "not captured") - else: - why, self.last = self.last, "no step" - STEPS[self.spec.name, why or "mono"] += 1 - def _decide(self, positions: torch.Tensor) -> str | None: fc = get_forward_context() if fc.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: From 524a46ec5e89dfa77fffae0b374addce77385534 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 16:01:10 -0500 Subject: [PATCH 16/79] chore: remove reload_weights Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 1 - tests/models/common/mono/reference.py | 6 ------ tests/models/common/mono/test_reference.py | 24 ---------------------- vllm/models/common/mono/runtime.py | 20 +----------------- vllm/models/common/mono/spec.py | 11 +++------- vllm/models/deepseek_v41/amd/mono_spec.py | 1 - 6 files changed, 4 insertions(+), 59 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 7fc5829e6e4a..3bbbea0d9392 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -191,7 +191,6 @@ def test_skipped_step_runs_vllm_layers(self, config, model, mono): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): mono.begin_once(self._positions(mono)) assert mono.step is None - assert mono.last is not None def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): off = max(mono.widths) + 1 diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 4e28686a26e0..b974d10b66ce 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -119,11 +119,6 @@ def launch_state(model, ctx: MonoContext) -> LaunchState: ) -def reload_weights(model, ctx: MonoContext, state: LaunchState) -> None: - for i, layer in enumerate(model.layers): - state.weights[i].copy_(layer.weight.data) - - def bind(ctx: MonoContext, caches) -> list[str]: reasons = [] for name, cache in caches.items(): @@ -150,7 +145,6 @@ def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceK skip_layer=skip_layer, skip_step=skip_step, launch_state=launch_state, - reload_weights=reload_weights, bind=bind, build=build, ) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index 5d3e6001dba0..0b8ff02b466b 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -3,8 +3,6 @@ """The conformance kit against the reference spec: what the kit asserts is the library's behaviour, so this is also the library's own test.""" -from dataclasses import replace - import pytest import torch @@ -88,28 +86,6 @@ def test_kernels_run_the_layers_they_replace(): assert [k.launched for k in mono.kernels.values() if k.launched] == [[width]] * 3 -def test_reload_weights_writes_in_place(): - config = reference.make_config() - model = reference.make_model(config) - mono = attach(model, config, SPECS_PATH) - before = {i: w.data_ptr() for i, w in mono.state.weights.items()} - with torch.no_grad(): - for layer in model.layers: - layer.weight.copy_(torch.full_like(layer.weight, 7)) - mono.reload_weights(model) - assert {i: w.data_ptr() for i, w in mono.state.weights.items()} == before - assert all((w == 7).all() for w in mono.state.weights.values()) - - -def test_reload_weights_is_refused_when_the_spec_declares_none(): - config = reference.make_config() - model = reference.make_model(config) - mono = attach(model, config, SPECS_PATH) - mono.spec = replace(mono.spec, reload_weights=None) - with pytest.raises(MonoUnsupported, match="weights cannot move"): - mono.reload_weights(model) - - def test_model_refusal_is_raised(): config = reference.make_config() with pytest.raises(MonoUnsupported, match="not a ReferenceModel"): diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index f88cd81cc877..04d090c3ccb3 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -78,7 +78,6 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.kernels: dict[tuple[int, int], LayerKernel] = {} self.state: Any = None self.step: MonoStep | None = None - self.last: str | None = None self._fc: Any = None self._bound: tuple | None = None @@ -130,21 +129,6 @@ def _register_warmup(self) -> None: if owners: MonoWarmup(self.spec.name, owners).register_warmup() - def reload_weights(self, model) -> None: - """Writes ``model``'s current weights into the allocation - ``launch_state`` returned, in place. - - Raises: - MonoUnsupported: The spec declares no ``reload_weights``, so its - weights cannot move after attach. - - """ - if self.spec.reload_weights is None: - raise MonoUnsupported( - f"MonoKernel {self.spec.name}: weights cannot move after attach" - ) - self.spec.reload_weights(model, self.ctx, self.state) - def _bind(self) -> None: """Re-checks the caches whenever they moved (keyed on every tensor's data_ptr and numel): memory profiling runs real steps on throwaway @@ -180,9 +164,7 @@ def begin_step(self, positions: torch.Tensor) -> str | None: """None runs this step on the kernels (and sets ``step``), else the reason it runs vLLM's layers.""" self.step = None - why = self._decide(positions) - self.last = why - return why + return self._decide(positions) def _decide(self, positions: torch.Tensor) -> str | None: fc = get_forward_context() diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 3c8e82936c87..25072fdfdd26 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -122,9 +122,9 @@ class MonoSpec(Generic[LaunchStateT]): as a ``MONO_SPECS`` tuple the library selects from. The parameter is the spec's own launch state: ``launch_state`` builds it, - and ``reload_weights``, ``skip_step`` and ``build`` are handed back the - same object. Annotate the spec with it (``MonoSpec[MyLaunchState](...)``) - and those four callbacks are checked against each other. + and ``skip_step`` and ``build`` are handed back the same object. Annotate + the spec with it (``MonoSpec[MyLaunchState](...)``) and those three + callbacks are checked against each other. Gating runs in tiers, and the prefix is the consequence. ``refuse_*`` says the deployment asked for something it cannot have, so it may take the @@ -153,10 +153,6 @@ class MonoSpec(Generic[LaunchStateT]): launch_state: The weights and launch state, built once at attach. The only place a weight copy is ever made. Opaque to the library, which only hands it back to the spec's own callbacks. - reload_weights: Writes new weights into the allocation - ``launch_state`` returned, in place, because the kernels and the - captured graphs hold its address. None refuses anything that moves - weights after attach. bind: Checks the KV caches the kernels will read, on every re-bind. build: One ``LayerKernel`` for (context, state, layer, width). @@ -179,6 +175,5 @@ class MonoSpec(Generic[LaunchStateT]): # Callbacks launch_state: Callable[[Any, MonoContext], LaunchStateT] - reload_weights: Callable[[Any, MonoContext, LaunchStateT], None] | None bind: Callable[[MonoContext, Mapping[str, Any]], list[str]] build: Callable[[MonoContext, LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 583128e726d3..73ac6f9e4458 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -519,7 +519,6 @@ def _layers(config) -> list[int]: skip_layer=_skip_layer, skip_step=_skip_step, launch_state=_launch_state, - reload_weights=None, bind=_bind, build=_build, ) From aafcbd87f85e54010f97dd07098a699a32b471f9 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 16:12:29 -0500 Subject: [PATCH 17/79] cleanup: remove monostep Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 7 +++-- tests/models/common/mono/reference.py | 12 +++++---- tests/models/test_dsv41_mono_decode.py | 10 +++---- vllm/models/common/mono/layer.py | 5 ++-- vllm/models/common/mono/runtime.py | 15 +++++------ vllm/models/common/mono/spec.py | 33 +++++------------------ vllm/models/deepseek_v41/amd/model.py | 8 +++--- vllm/models/deepseek_v41/amd/mono_spec.py | 28 +++++++++---------- 8 files changed, 43 insertions(+), 75 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 3bbbea0d9392..204918cd23d6 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -60,9 +60,8 @@ def forward_layer(mono, i: int, layer, positions, *args, **kwargs): """Layer ``i`` on its kernel when this step runs on the kernels, else the vLLM layer unchanged: the layer loop a model writes for itself.""" k = kernel_for_layer(mono, i, positions) - step = None if mono is None else mono.step - if k is not None and step is not None: - out = k.launch(step, positions, *args, **kwargs) + if k is not None: + out = k.launch(positions, *args, **kwargs) if out is not None: return out return layer(positions, *args, **kwargs) @@ -190,7 +189,7 @@ def test_every_rank_decides_alike(self, config, model): def test_skipped_step_runs_vllm_layers(self, config, model, mono): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): mono.begin_once(self._positions(mono)) - assert mono.step is None + assert mono.width is None def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): off = max(mono.widths) + 1 diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index b974d10b66ce..eac7de6a0749 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -16,7 +16,8 @@ import torch from vllm.config.compilation import CUDAGraphMode -from vllm.models.common.mono.spec import CacheReq, MonoContext, MonoSpec, MonoStep +from vllm.forward_context import get_forward_context +from vllm.models.common.mono.spec import CacheReq, MonoContext, MonoSpec CACHE = CacheReq(group="kv", dtype="bfloat16", layout="2nhd") @@ -72,8 +73,8 @@ def __init__(self, ctx: MonoContext, state: LaunchState, layer, width: int): self.launched: list[int] = [] self.compiled: list[Any] = [] - def launch(self, step: MonoStep, positions: torch.Tensor, *args, **kwargs): - self.launched.append(step.num_tokens) + def launch(self, positions: torch.Tensor, *args, **kwargs): + self.launched.append(self.width) return self.layer(positions, *args, **kwargs) def warmup_keys(self) -> list[Any]: @@ -109,8 +110,9 @@ def skip_layer(ctx: MonoContext, layer) -> str | None: return "dense layer" if layer.dense else None -def skip_step(ctx: MonoContext, state: LaunchState, step: MonoStep) -> str | None: - return "prefill" if step.attn_metadata.get("prefill") else None +def skip_step(ctx: MonoContext, state: LaunchState, width: int) -> str | None: + md = get_forward_context().attn_metadata + return "prefill" if md.get("prefill") else None def launch_state(model, ctx: MonoContext) -> LaunchState: diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 82c58a565f8b..60c924a7acda 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -476,7 +476,6 @@ def test_whole_launch_runs_the_layer(monkeypatch): with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 1, positions) out = kernel.launch( - mono.step, torch.zeros(1, HIDDEN), positions, torch.zeros(1, 4, HIDDEN), @@ -498,9 +497,7 @@ def test_first_layer_takes_the_embedding(monkeypatch): positions = torch.zeros(1, dtype=torch.long) with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 0, positions) - out = kernel.launch( - mono.step, torch.zeros(1, HIDDEN), positions, None, None, None, None - ) + out = kernel.launch(torch.zeros(1, HIDDEN), positions, None, None, None, None) assert out[0] == "whole" and not out[2]["fold"] assert out[1][2].shape == (1, 4, HIDDEN) @@ -512,7 +509,6 @@ def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): kernel = kernel_for_layer(mono, 1, positions) assert kernel.ffn_only out = kernel.launch( - mono.step, torch.zeros(1, HIDDEN), torch.zeros(1, 4, HIDDEN), torch.zeros(1, 4, 1), @@ -524,12 +520,12 @@ def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): def test_launch_declines_a_layout_it_cannot_read(monkeypatch): """A later pipeline stage's first layer gets a stream without the mixes: - the step descriptor cannot see that, so the launch declines it.""" + ``skip_step`` cannot see that, so the launch declines it.""" mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 1, positions) out = kernel.launch( - mono.step, torch.zeros(1, 4, HIDDEN), positions, None, None, None, None + torch.zeros(1, 4, HIDDEN), positions, None, None, None, None ) assert out is None diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index 57e3f6b0bb9a..34d25b0a01cd 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -61,7 +61,6 @@ def kernel_for_layer(mono: MonoRuntime | None, i: int, positions) -> LayerKernel if mono is None: return None mono.begin_once(positions) - step = mono.step - if step is None: + if mono.width is None: return None - return mono.kernels.get((i, step.width)) + return mono.kernels.get((i, mono.width)) diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 04d090c3ccb3..9d478e6e8b46 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -31,7 +31,6 @@ LayerKernel, MonoContext, MonoSpec, - MonoStep, MonoUnsupported, ) @@ -77,7 +76,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.widths = spec.widths(vllm_config) self.kernels: dict[tuple[int, int], LayerKernel] = {} self.state: Any = None - self.step: MonoStep | None = None + self.width: int | None = None self._fc: Any = None self._bound: tuple | None = None @@ -161,21 +160,19 @@ def begin_once(self, positions: torch.Tensor) -> None: self.begin_step(positions) def begin_step(self, positions: torch.Tensor) -> str | None: - """None runs this step on the kernels (and sets ``step``), else the + """None runs this step on the kernels (and sets ``width``), else the reason it runs vLLM's layers.""" - self.step = None + self.width = None return self._decide(positions) def _decide(self, positions: torch.Tensor) -> str | None: - fc = get_forward_context() - if fc.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: return "piecewise" n = positions.size(-1) if n not in self.widths: return "width" - step = MonoStep(n, n, positions, fc.attn_metadata, fc.cudagraph_runtime_mode) - why = self.spec.skip_step(self.ctx, self.state, step) + why = self.spec.skip_step(self.ctx, self.state, n) if why is None: self._bind() - self.step = step + self.width = n return why diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 25072fdfdd26..7bc7fedef3cd 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -10,8 +10,6 @@ import torch -from vllm.config.compilation import CUDAGraphMode - if TYPE_CHECKING: from vllm.config import VllmConfig @@ -69,37 +67,17 @@ class MonoContext: device: torch.device -@dataclass -class MonoStep: - """One forward's batch descriptor, built once and read by every layer. - - Attributes: - num_tokens: The step's query rows. - width: The padded rows the launch serves. - positions: The step's position ids. - attn_metadata: The forward context's attention metadata, by layer name. - runtime_mode: The CUDA graph mode this step runs under. - - """ - - num_tokens: int - width: int - positions: torch.Tensor - attn_metadata: Any - runtime_mode: CUDAGraphMode - - class LayerKernel(Protocol): """One decoder layer's launch, at one width, on this rank.""" cache_layers: tuple[str, ...] """Attention layers whose caches ``MonoSpec.bind`` checks.""" - def launch(self, step: MonoStep, *args, **kwargs) -> Any: + def launch(self, *args, **kwargs) -> Any: """Takes and returns what the vLLM layer it replaces does, or None to decline this call, which sends the layer down vLLM's path. Decline only - on what the step descriptor cannot answer, such as the layout of the - tensors the layer was handed.""" + on what ``skip_step`` cannot answer, such as the layout of the tensors + the layer was handed.""" ... def warmup_keys(self) -> list[Any]: @@ -149,7 +127,8 @@ class MonoSpec(Generic[LaunchStateT]): skip_layer: Why one layer takes vLLM's path, None when it takes the kernel's. skip_step: Why one step takes vLLM's path, None when it takes the - kernel's. + kernel's. Takes the step's width; anything else it needs it reads + from the forward context, which every rank sees alike. launch_state: The weights and launch state, built once at attach. The only place a weight copy is ever made. Opaque to the library, which only hands it back to the spec's own callbacks. @@ -171,7 +150,7 @@ class MonoSpec(Generic[LaunchStateT]): refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] skip_layer: Callable[[MonoContext, Any], str | None] - skip_step: Callable[[MonoContext, LaunchStateT, MonoStep], str | None] + skip_step: Callable[[MonoContext, LaunchStateT, int], str | None] # Callbacks launch_state: Callable[[Any, MonoContext], LaunchStateT] diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 276b77472a53..a364e6cf9eec 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -283,10 +283,8 @@ def forward( DSV41LayerKernel | None, kernel_for_layer(mono, self.attn.layer_id, positions), ) - step = None if mono is None else mono.step - if kernel is not None and step is not None and not kernel.ffn_only: + if kernel is not None and not kernel.ffn_only: out = kernel.launch( - step, x, positions, residual, @@ -397,8 +395,8 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) elif not self.attn.wo_b.reduce_results: - if kernel is not None and step is not None: - out = kernel.launch(step, x, residual, post_mix, res_mix, attn_pre) + if kernel is not None: + out = kernel.launch(x, residual, post_mix, res_mix, attn_pre) if out is not None: return out x = tensor_model_parallel_all_reduce(x) diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 73ac6f9e4458..e9f36dee2288 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -31,13 +31,13 @@ import vllm.envs as envs from vllm.config import CUDAGraphMode +from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.models.common.mono.spec import ( CacheReq, Feature, MonoContext, MonoSpec, - MonoStep, ) from vllm.platforms import current_platform @@ -137,7 +137,7 @@ def warmup_keys(self) -> list[Any]: def compile(self, key: Any) -> None: raise NotImplementedError - def launch(self, step: MonoStep, *args, **kwargs) -> Any: + def launch(self, *args, **kwargs) -> Any: raise NotImplementedError @@ -171,7 +171,6 @@ def compile(self, key: Any) -> None: def launch( self, - step: MonoStep, x: torch.Tensor, positions: torch.Tensor, residual: torch.Tensor | None, @@ -200,12 +199,13 @@ def launch( hashes = engram_hashes[:, layer.engram.layer_hash_index] residual = layer.engram(residual, hashes, engram_mask) fold = False - swa = step.attn_metadata[self.swa_key] + md = cast("dict[str, Any]", get_forward_context().attn_metadata) + swa = md[self.swa_key] kw: dict[str, Any] = {} if self.index: kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) if self.ratio: - comp = step.attn_metadata[self.comp_key] + comp = md[cast("str", self.comp_key)] kw.update( topk_indices=attn.topk_indices_buffer, comp_cache=_record_view( @@ -246,7 +246,6 @@ def compile(self, key: Any) -> None: def launch( self, - step: MonoStep, part: torch.Tensor, residual: torch.Tensor, post_mix: torch.Tensor, @@ -355,18 +354,17 @@ def _ffn_only(layer) -> str | None: return None -def _skip_step( - ctx: MonoContext, launch_state: LaunchState, step: MonoStep -) -> str | None: +def _skip_step(ctx: MonoContext, launch_state: LaunchState, width: int) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP group sees alike.""" - md = step.attn_metadata - if not 1 <= step.width <= MAX_ROWS: - return f"{step.width} rows" - if step.runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + fc = get_forward_context() + md = fc.attn_metadata + if not 1 <= width <= MAX_ROWS: + return f"{width} rows" + if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): # a PIECEWISE (breakable) capture is replayed for mixed batches too - return f"{step.runtime_mode} step" + return f"{fc.cudagraph_runtime_mode} step" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches return "no attention metadata" if any(md.get(k) is None for k in launch_state.comp_keys): @@ -377,7 +375,7 @@ def _skip_step( return "layer metadata missing" if swa.num_prefills != 0 or swa.num_decodes == 0: return "not a decode-only step" - if swa.num_decode_tokens > step.width: + if swa.num_decode_tokens > width: return "more decode tokens than rows" if not launch_state.need_swa: continue From 2b1b190f644ff0dc0490ba3a696cb74c76a81501 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 16:25:37 -0500 Subject: [PATCH 18/79] chore: cleanup docstrings Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 14 +---------- tests/models/common/mono/reference.py | 9 +------ tests/models/common/mono/test_reference.py | 3 +-- tests/models/test_dsv41_mono_decode.py | 9 +------ tests/models/test_dsv41_mono_layers.py | 19 +-------------- vllm/models/common/mono/__init__.py | 10 +------- vllm/models/common/mono/guards.py | 14 +---------- vllm/models/common/mono/layer.py | 3 +-- vllm/models/common/mono/runtime.py | 13 +--------- vllm/models/common/mono/spec.py | 3 +-- vllm/models/deepseek_v41/amd/mono/__init__.py | 4 +--- vllm/models/deepseek_v41/amd/mono_spec.py | 24 +------------------ 12 files changed, 12 insertions(+), 113 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 204918cd23d6..43e7e96a5697 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -1,18 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The MonoKernel conformance kit (RFC #60904). - -Every model that ships a ``MonoSpec`` subclasses ``MonoConformance`` and fills -in the five hooks, which gives it the contract the library relies on: refusals -are complete, the spec module is importable without a kernel backend, the -capture sizes are the spec's widths, one forward decides once and alike on -every rank, caches re-bind when they move, attach does not move the model's -weights, cache checks report every mismatch, and every width is warmed up -before capture. - -The kit never launches a kernel: it drives the library through the spec's own -callbacks, so a spec whose kernels need a GPU still conforms on CPU. -""" +"""The MonoKernel conformance kit (RFC #60904): what a ``MonoSpec`` must satisfy.""" import subprocess import sys diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index eac7de6a0749..ba0a237b2325 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -1,13 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The reference MonoKernel: a ``LayerKernel`` that calls the vLLM layer it -replaces, and the ``MONO_SPECS`` that declare it. - -It computes nothing of its own and launches nothing on a device, so it runs on -CPU in the conformance kit (``kit``). What it is for is the contract: a model -porting to MonoKernels can read one complete spec, and the library's own -behaviour can be tested without a GPU or a kernel backend. -""" +"""The reference MonoKernel: a ``LayerKernel`` over the layer it replaces.""" from dataclasses import dataclass, replace from types import SimpleNamespace diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index 0b8ff02b466b..312d935b659b 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The conformance kit against the reference spec: what the kit asserts is the -library's behaviour, so this is also the library's own test.""" +"""The conformance kit against the reference spec: the library's own test.""" import pytest import torch diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 60c924a7acda..246263945230 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -1,13 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's MonoSpec against the conformance kit, and the choices that -are its own: which layers take the whole-layer launch and which take the FFN -launch, and which steps the kernels are exact on (decode only, at most MAX_ROWS -rows, causal SWA windows, eager or in a FULL graph). - -The kernels need CDNA4 and FlyDSL; these tests stub the runner module, so they -check the spec's decisions on CPU and never launch anything. -""" +"""DeepSeek-V4.1's MonoSpec against the conformance kit, plus its own choices.""" import importlib import sys diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index fd151115e8d0..7008346ffe3a 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -1,23 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The DeepSeek-V4.1 mono decode path through the real model: a five-layer -DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``--monokernel`` on -and off on the same random weights: every decoder layer's outputs at every -decode step must match. Unlike ``test_dsv41_mono_numerics.py``, which -checks the kernels against vLLM's ops one stage at a time, this runs -``DeepseekV4DecoderLayer.forward`` as the engine does -- the model's hooks, -the per-step path choice, the weights taken off the layer, the metadata and -the indexer's top-k -- so a change to the decoder layer that the kernels do -not follow fails here. - -Every layer takes the whole-layer path (K1 + K2); the five hold each variant: -- 0: the first layer, sliding window only (K1 takes the embedding settled) -- 1: ratio-2 KV and index source (vLLM's compressor and indexer between K1 - and K2) -- 2: ratio-2 consumer with an Engram (K1 takes the residual settled) -- 3: ratio-1 KV, index and candidate source -- 4: ratio-1 index source over 3's candidates and KV -""" +"""DeepSeek-V4.1 mono decode against vLLM's path, through the real model.""" import zlib diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index 74b02af89f94..26c561f1cbbf 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -1,14 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MonoKernels (RFC #60904): persistent kernels that run most or all of one -decoder layer's decode step in a single launch. A model brings a ``MonoSpec`` -(``spec``); this package owns the rest: refusals and TP agreement (``guards``), -the per-step decision, cache re-binding and step counters (``runtime``), and the -layer-loop hook (``layer``). - -Peer memory and in-kernel failure reporting are the model's own: only its -kernels know what they put where and what they record, so a model that needs -either builds it in the launch state it brings.""" +"""MonoKernels (RFC #60904): a decoder layer's decode step in one launch.""" from vllm.models.common.mono.guards import configure from vllm.models.common.mono.layer import attach, kernel_for_layer diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index 389e17525767..01faa37fcc6f 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -1,18 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Refusals. ``configure``, the only name a model calls, runs once from core -config post-init, before load: it imports only the model's spec module, picks -the variant this config can serve, adds its widths to the CUDA graph capture -sizes and raises with every reason the config is refused. The rest is the -library's own and stays private to this package: ``_agree_across_ranks`` runs a -later refusal stage (after load) on every TP rank and agrees on the verdict -before anything collective follows. - -Peer memory is the model's own: only its kernels know what they put where, and -``tp_group`` is where it reads this rank's place in the TP group. What this -package keeps is the agreement every stage goes through (``tp_agree``), because -a rank that refuses alone leaves every other rank waiting in-kernel for ever. -""" +"""Refusals: the variant a config can serve, and TP agreement on every stage.""" import importlib from collections.abc import Callable, Sequence diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py index 34d25b0a01cd..6112e662fa21 100644 --- a/vllm/models/common/mono/layer.py +++ b/vllm/models/common/mono/layer.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The hooks a model calls: ``attach`` after load, ``kernel_for_layer`` in its -layer loop.""" +"""The hooks a model calls: ``attach``, then ``kernel_for_layer`` per layer.""" from typing import TYPE_CHECKING diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 9d478e6e8b46..731de233847d 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -1,17 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The live side of a spec: one ``MonoRuntime`` a rank, built at attach, owning -this rank's layer kernels and making the per-step decision and cache re-binding. -It launches nothing itself. - -``MonoRuntime.begin_step`` decides once a forward, for every kernel layer, from -what every rank sees alike: the batch's width and the attention metadata, never -rank-local state and never a collective. Under a FULL CUDA graph the decision -is baked in at capture, so it is a function of the captured size: FULL graphs -serve uniform decode batches and every width is a capture size -(``guards``), so no replay pads a request's rows. PIECEWISE -steps run vLLM's layers. -""" +"""One ``MonoRuntime`` a rank: the per-step decision and cache re-binding.""" from dataclasses import dataclass from typing import TYPE_CHECKING, Any diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 7bc7fedef3cd..685907a8afdb 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""What a model declares to run MonoKernels: its ``MonoSpec``, and the types -the spec's callbacks see.""" +"""What a model declares to run MonoKernels: ``MonoSpec`` and its callbacks.""" from collections.abc import Callable, Mapping from dataclasses import dataclass diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py index e5ba197cfb79..d8de17ee19bc 100644 --- a/vllm/models/deepseek_v41/amd/mono/__init__.py +++ b/vllm/models/deepseek_v41/amd/mono/__init__.py @@ -1,5 +1,3 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1-Flash mono decode layer: one decoder layer of a decode step in -two persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The -host side is ``runner``; vLLM docks it in ``..mono_spec``.""" +"""DeepSeek-V4.1-Flash mono decode layer, in two FlyDSL launches (README.md).""" diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index e9f36dee2288..52e08dcb9cb6 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -1,28 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's MonoKernels on ROCm CDNA4 (``--monokernel``, RFC #60904). - -A decode step's backbone layer (the attention seam, the attention, its TP -all-reduce, the FFN seam, the MoE and its all-reduce) runs as two persistent -FlyDSL launches (``mono``, this directory) on the layer's loaded weights, the -TP reductions inside the kernels. Every backbone layer of the checkpoint runs -whole (compress ratio 0, 1 or 2 over ``fp8_ds_mla`` records): the first layer -and the Engram layers hand K1 a settled residual, and an index source runs -vLLM's compressor and indexer between the two launches. A layer outside that -(an MXFP8 KV cache, another indexer) keeps vLLM's attention seam and attention -and runs the rest as one launch, the FFN launch, from wo_b's unreduced output. -A step takes the kernels when it is decode only, at a captured width of at most -``MAX_ROWS`` rows, with causal SWA windows. Anything else runs the layer as -before; the paths share every tensor at the layer boundary, so they interleave -freely. - -Every TP rank must take the same path at every layer (the kernels wait on each -other's pushes): the step decision reads only metadata every rank holds alike. - -This module is imported by ``mono.configure`` in the engine core process, -before the weights are loaded and with no device in hand, so it pulls in -neither FlyDSL nor AITER at import. -""" +"""DeepSeek-V4.1's MonoKernels on ROCm CDNA4 (``--monokernel``, RFC #60904).""" from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast From 0f474eb3ff9c6e7f3d77b3790ca90b0d3eb5bbb3 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 16:48:54 -0500 Subject: [PATCH 19/79] fix: rename bind->refuse-caches Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 18 +++++++++--------- tests/models/common/mono/reference.py | 4 ++-- vllm/models/common/mono/runtime.py | 17 +++++++---------- vllm/models/common/mono/spec.py | 7 ++++--- vllm/models/deepseek_v41/amd/mono_spec.py | 4 ++-- 5 files changed, 24 insertions(+), 26 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 43e7e96a5697..b918ea686b03 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -91,7 +91,7 @@ def skipped_metadata(self, config, model) -> Any: raise NotImplementedError def bad_caches(self, config, model) -> dict[str, Any]: - """Caches that break at least two of ``bind``'s checks at once.""" + """Caches that break at least two of ``refuse_caches``'s checks.""" raise NotImplementedError @pytest.fixture @@ -189,21 +189,21 @@ def test_piecewise_runs_vllm_layers(self, config, model, mono): with forward(config, CUDAGraphMode.PIECEWISE, meta): assert mono.begin_step(self._positions(mono)) == "piecewise" - def test_caches_rebind_when_they_move(self, config, model, mono): - binds: list = [] - mono.spec = _spy(mono.spec, "bind", binds) + def test_caches_rechecked_when_they_move(self, config, model, mono): + checks: list = [] + mono.spec = _spy(mono.spec, "refuse_caches", checks) self._decode(config, model, mono) - assert len(binds) == 1 + assert len(checks) == 1 self._decode(config, model, mono) - assert len(binds) == 1 + assert len(checks) == 1 sfc = config.compilation_config.static_forward_context for name in mono.cache_layers: sfc[name].kv_cache = sfc[name].kv_cache.clone() self._decode(config, model, mono) - assert len(binds) == 2 + assert len(checks) == 2 - def test_bind_reports_every_mismatch(self, config, model, mono): - reasons = mono.spec.bind(mono.ctx, self.bad_caches(config, model)) + def test_refuse_caches_reports_every_reason(self, config, model, mono): + reasons = mono.spec.refuse_caches(mono.ctx, self.bad_caches(config, model)) assert len(reasons) >= 2, reasons def test_mismatched_caches_refuse_the_step(self, config, model, mono): diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index ba0a237b2325..e390d9f0d2fe 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -114,7 +114,7 @@ def launch_state(model, ctx: MonoContext) -> LaunchState: ) -def bind(ctx: MonoContext, caches) -> list[str]: +def refuse_caches(ctx: MonoContext, caches) -> list[str]: reasons = [] for name, cache in caches.items(): if str(cache.dtype).removeprefix("torch.") != CACHE.dtype: @@ -140,7 +140,7 @@ def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceK skip_layer=skip_layer, skip_step=skip_step, launch_state=launch_state, - bind=bind, + refuse_caches=refuse_caches, build=build, ) diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 731de233847d..558bcb256511 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""One ``MonoRuntime`` a rank: the per-step decision and cache re-binding.""" +"""One ``MonoRuntime`` a rank: the per-step decision and the cache re-checks.""" from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -67,7 +67,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.state: Any = None self.width: int | None = None self._fc: Any = None - self._bound: tuple | None = None + self._checked: tuple | None = None indices = self._kernel_layers(model) if not indices: @@ -117,7 +117,7 @@ def _register_warmup(self) -> None: if owners: MonoWarmup(self.spec.name, owners).register_warmup() - def _bind(self) -> None: + def _refuse_caches(self) -> None: """Re-checks the caches whenever they moved (keyed on every tensor's data_ptr and numel): memory profiling runs real steps on throwaway caches before the real ones exist.""" @@ -128,18 +128,18 @@ def _bind(self) -> None: for c in caches.values() for t in (c if isinstance(c, (list, tuple)) else (c,)) ) - if key == self._bound: + if key == self._checked: return if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): raise MonoUnsupported( f"MonoKernel {self.spec.name}: caches moved during CUDA graph capture" ) - reasons = self.spec.bind(self.ctx, caches) + reasons = self.spec.refuse_caches(self.ctx, caches) if reasons: raise MonoUnsupported( f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) ) - self._bound = key + self._checked = key def begin_once(self, positions: torch.Tensor) -> None: """Decides this forward's step, on the first kernel layer to ask.""" @@ -152,9 +152,6 @@ def begin_step(self, positions: torch.Tensor) -> str | None: """None runs this step on the kernels (and sets ``width``), else the reason it runs vLLM's layers.""" self.width = None - return self._decide(positions) - - def _decide(self, positions: torch.Tensor) -> str | None: if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: return "piecewise" n = positions.size(-1) @@ -162,6 +159,6 @@ def _decide(self, positions: torch.Tensor) -> str | None: return "width" why = self.spec.skip_step(self.ctx, self.state, n) if why is None: - self._bind() + self._refuse_caches() self.width = n return why diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 685907a8afdb..a32a83730182 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -70,7 +70,7 @@ class LayerKernel(Protocol): """One decoder layer's launch, at one width, on this rank.""" cache_layers: tuple[str, ...] - """Attention layers whose caches ``MonoSpec.bind`` checks.""" + """Attention layers whose caches ``MonoSpec.refuse_caches`` checks.""" def launch(self, *args, **kwargs) -> Any: """Takes and returns what the vLLM layer it replaces does, or None to @@ -131,7 +131,8 @@ class MonoSpec(Generic[LaunchStateT]): launch_state: The weights and launch state, built once at attach. The only place a weight copy is ever made. Opaque to the library, which only hands it back to the spec's own callbacks. - bind: Checks the KV caches the kernels will read, on every re-bind. + refuse_caches: Deployment refusals the KV caches the kernels will read + answer, re-run whenever those caches move. build: One ``LayerKernel`` for (context, state, layer, width). """ @@ -153,5 +154,5 @@ class MonoSpec(Generic[LaunchStateT]): # Callbacks launch_state: Callable[[Any, MonoContext], LaunchStateT] - bind: Callable[[MonoContext, Mapping[str, Any]], list[str]] + refuse_caches: Callable[[MonoContext, Mapping[str, Any]], list[str]] build: Callable[[MonoContext, LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 52e08dcb9cb6..a5317f516638 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -436,7 +436,7 @@ def _attn_weights(a, world: int): return attn -def _bind(ctx: MonoContext, caches) -> list[str]: +def _refuse_caches(ctx: MonoContext, caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as [blocks, block, RECORD] bytes.""" reasons = [] @@ -495,7 +495,7 @@ def _layers(config) -> list[int]: skip_layer=_skip_layer, skip_step=_skip_step, launch_state=_launch_state, - bind=_bind, + refuse_caches=_refuse_caches, build=_build, ) From 0681c31034d3b206e426b79c4555e058b9d63517 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 16:52:50 -0500 Subject: [PATCH 20/79] chore: move refuse_caches Signed-off-by: simondanielsson --- tests/models/common/mono/reference.py | 2 +- vllm/models/common/mono/spec.py | 6 +++--- vllm/models/deepseek_v41/amd/mono_spec.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index e390d9f0d2fe..49f9fb115ca2 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -137,10 +137,10 @@ def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceK widths=widths, refuse_config=refuse_config, refuse_model=refuse_model, + refuse_caches=refuse_caches, skip_layer=skip_layer, skip_step=skip_step, launch_state=launch_state, - refuse_caches=refuse_caches, build=build, ) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index a32a83730182..3b955eef9fe6 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -123,6 +123,8 @@ class MonoSpec(Generic[LaunchStateT]): refuse_config: Deployment refusals the engine's config answers, before the weights are loaded. refuse_model: Deployment refusals only the loaded model answers. + refuse_caches: Deployment refusals the KV caches the kernels will read + answer, re-run whenever those caches move. skip_layer: Why one layer takes vLLM's path, None when it takes the kernel's. skip_step: Why one step takes vLLM's path, None when it takes the @@ -131,8 +133,6 @@ class MonoSpec(Generic[LaunchStateT]): launch_state: The weights and launch state, built once at attach. The only place a weight copy is ever made. Opaque to the library, which only hands it back to the spec's own callbacks. - refuse_caches: Deployment refusals the KV caches the kernels will read - answer, re-run whenever those caches move. build: One ``LayerKernel`` for (context, state, layer, width). """ @@ -149,10 +149,10 @@ class MonoSpec(Generic[LaunchStateT]): # Capabilities refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] + refuse_caches: Callable[[MonoContext, Mapping[str, Any]], list[str]] skip_layer: Callable[[MonoContext, Any], str | None] skip_step: Callable[[MonoContext, LaunchStateT, int], str | None] # Callbacks launch_state: Callable[[Any, MonoContext], LaunchStateT] - refuse_caches: Callable[[MonoContext, Mapping[str, Any]], list[str]] build: Callable[[MonoContext, LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index a5317f516638..6179bd17feac 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -492,10 +492,10 @@ def _layers(config) -> list[int]: widths=_widths, refuse_config=_refuse_config, refuse_model=_refuse_model, + refuse_caches=_refuse_caches, skip_layer=_skip_layer, skip_step=_skip_step, launch_state=_launch_state, - refuse_caches=_refuse_caches, build=_build, ) From b97d5b6623d190344f2a2363f2e9b6da1ec21ec0 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 17:03:07 -0500 Subject: [PATCH 21/79] chore: fix Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 2 +- tests/models/common/mono/reference.py | 14 ++--- tests/models/common/mono/test_reference.py | 2 +- tests/models/test_dsv41_mono_decode.py | 2 +- vllm/models/common/mono/__init__.py | 2 +- vllm/models/common/mono/layer.py | 65 ---------------------- vllm/models/common/mono/runtime.py | 53 +++++++++++++++++- vllm/models/common/mono/spec.py | 11 ---- vllm/models/deepseek_v41/amd/mono_spec.py | 2 - 9 files changed, 63 insertions(+), 90 deletions(-) delete mode 100644 vllm/models/common/mono/layer.py diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index b918ea686b03..98d558cf638b 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -23,7 +23,7 @@ _set_capture_sizes, configure, ) -from vllm.models.common.mono.layer import attach, kernel_for_layer +from vllm.models.common.mono.runtime import attach, kernel_for_layer from vllm.models.common.mono.spec import Feature, MonoUnsupported KERNEL_BACKENDS = ("aiter", "flydsl") diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 49f9fb115ca2..c82444a20f17 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -10,9 +10,10 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context -from vllm.models.common.mono.spec import CacheReq, MonoContext, MonoSpec +from vllm.models.common.mono.spec import MonoContext, MonoSpec -CACHE = CacheReq(group="kv", dtype="bfloat16", layout="2nhd") +CACHE_DTYPE = "bfloat16" +CACHE_LAYOUT = "2nhd" _LAYOUT_DIMS = {"2nhd": 4} @@ -117,10 +118,10 @@ def launch_state(model, ctx: MonoContext) -> LaunchState: def refuse_caches(ctx: MonoContext, caches) -> list[str]: reasons = [] for name, cache in caches.items(): - if str(cache.dtype).removeprefix("torch.") != CACHE.dtype: - reasons.append(f"{name}: dtype {cache.dtype}, needs {CACHE.dtype}") - if cache.ndim != _LAYOUT_DIMS[CACHE.layout]: - reasons.append(f"{name}: {cache.ndim}d, needs layout {CACHE.layout}") + if str(cache.dtype).removeprefix("torch.") != CACHE_DTYPE: + reasons.append(f"{name}: dtype {cache.dtype}, needs {CACHE_DTYPE}") + if cache.ndim != _LAYOUT_DIMS[CACHE_LAYOUT]: + reasons.append(f"{name}: {cache.ndim}d, needs layout {CACHE_LAYOUT}") return reasons @@ -131,7 +132,6 @@ def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceK REFERENCE = MonoSpec[LaunchState]( name="reference", supports=frozenset(), - caches={"attn": CACHE}, weight_source="views", layers=layers, widths=widths, diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index 312d935b659b..8b62d1f59672 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -9,7 +9,7 @@ from tests.models.common.mono.kit import MonoConformance, forward, forward_layer from vllm.config.compilation import CUDAGraphMode from vllm.models.common.mono.guards import _enabled, _select_spec -from vllm.models.common.mono.layer import attach +from vllm.models.common.mono.runtime import attach from vllm.models.common.mono.spec import MonoUnsupported SPECS_PATH = "tests.models.common.mono.reference:MONO_SPECS" diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 246263945230..8914ba7a6119 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -13,7 +13,7 @@ from tests.models.common.mono.kit import MonoConformance, forward from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -from vllm.models.common.mono.layer import attach, kernel_for_layer +from vllm.models.common.mono.runtime import attach, kernel_for_layer from vllm.models.common.mono.spec import MonoUnsupported from vllm.models.deepseek_v41.amd import mono_spec as ms diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index 26c561f1cbbf..510326126c91 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -3,7 +3,7 @@ """MonoKernels (RFC #60904): a decoder layer's decode step in one launch.""" from vllm.models.common.mono.guards import configure -from vllm.models.common.mono.layer import attach, kernel_for_layer +from vllm.models.common.mono.runtime import attach, kernel_for_layer __all__ = [ "attach", diff --git a/vllm/models/common/mono/layer.py b/vllm/models/common/mono/layer.py deleted file mode 100644 index 6112e662fa21..000000000000 --- a/vllm/models/common/mono/layer.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The hooks a model calls: ``attach``, then ``kernel_for_layer`` per layer.""" - -from typing import TYPE_CHECKING - -from vllm.logger import init_logger -from vllm.models.common.mono.guards import ( - _enabled, - _load_specs, - _refusal_message, - _select_spec, -) -from vllm.models.common.mono.runtime import MonoRuntime -from vllm.models.common.mono.spec import LayerKernel, MonoUnsupported - -if TYPE_CHECKING: - from vllm.config import VllmConfig - -logger = init_logger(__name__) - - -def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | None: - """Picks the variant, builds the launch state and this rank's layer - kernels. Collective over the TP group. - - Args: - model: The loaded model, in ``process_weights_after_loading``. - vllm_config: The engine's config. - specs_path: ``"package.module:MONO_SPECS"``. - - Returns: - This rank's ``MonoRuntime``, or None with the switch off or when no layer - takes a kernel. - - Raises: - MonoUnsupported: Every variant refused the config, or this rank (or - another) refused the loaded model. - - """ - if not _enabled(vllm_config): - return None - spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) - if spec is None: - raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) - mono = MonoRuntime(spec, model, vllm_config) - if not mono.kernels: - logger.info("MonoKernel %s: no layer takes a kernel", spec.name) - return None - return mono - - -def kernel_for_layer(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | None: - """Layer ``i``'s kernel when this step runs on the kernels, else None. - - The first call under a forward context decides for every layer. The model - calls ``launch`` where it suits the layer, and runs the vLLM layer - unchanged when this returns None. - """ - if mono is None: - return None - mono.begin_once(positions) - if mono.width is None: - return None - return mono.kernels.get((i, mono.width)) diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 558bcb256511..c6c879049149 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""One ``MonoRuntime`` a rank: the per-step decision and the cache re-checks.""" +"""One ``MonoRuntime`` a rank: the per-step decision and the cache re-checks, +behind the hooks a model calls (``attach``, then ``kernel_for_layer`` per layer). +""" from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -13,7 +15,11 @@ from vllm.model_executor.warmup.jit_warmup import VllmJitKernel from vllm.models.common.mono.guards import ( _agree_across_ranks, + _enabled, + _load_specs, + _refusal_message, _refuse_model, + _select_spec, tp_group, ) from vllm.models.common.mono.spec import ( @@ -162,3 +168,48 @@ def begin_step(self, positions: torch.Tensor) -> str | None: self._refuse_caches() self.width = n return why + + +def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | None: + """Picks the variant, builds the launch state and this rank's layer + kernels. Collective over the TP group. + + Args: + model: The loaded model, in ``process_weights_after_loading``. + vllm_config: The engine's config. + specs_path: ``"package.module:MONO_SPECS"``. + + Returns: + This rank's ``MonoRuntime``, or None with the switch off or when no layer + takes a kernel. + + Raises: + MonoUnsupported: Every variant refused the config, or this rank (or + another) refused the loaded model. + + """ + if not _enabled(vllm_config): + return None + spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) + if spec is None: + raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) + mono = MonoRuntime(spec, model, vllm_config) + if not mono.kernels: + logger.info("MonoKernel %s: no layer takes a kernel", spec.name) + return None + return mono + + +def kernel_for_layer(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | None: + """Layer ``i``'s kernel when this step runs on the kernels, else None. + + The first call under a forward context decides for every layer. The model + calls ``launch`` where it suits the layer, and runs the vLLM layer + unchanged when this returns None. + """ + if mono is None: + return None + mono.begin_once(positions) + if mono.width is None: + return None + return mono.kernels.get((i, mono.width)) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 3b955eef9fe6..eefbf240d499 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -41,15 +41,6 @@ class MonoUnsupported(Exception): """The configuration or the loaded model is outside what a spec serves.""" -@dataclass(frozen=True) -class CacheReq: - """A cache the kernels read: its KV cache group, dtype and layout.""" - - group: str - dtype: str - layout: str - - @dataclass(frozen=True) class MonoContext: """One rank's view of the deployment, built once at attach. @@ -114,7 +105,6 @@ class MonoSpec(Generic[LaunchStateT]): name: Log and metric name. supports: Features this spec serves. Default deny: a feature in use that is absent from it refuses the deployment. - caches: The caches the kernels read, by attention layer name. weight_source: ``"views"`` to read vLLM's post-load layout in place, ``"convert"`` to build a second layout once after load. layers: The decoder layers that get a kernel, from the model's config. @@ -139,7 +129,6 @@ class MonoSpec(Generic[LaunchStateT]): name: str supports: frozenset[Feature] - caches: Mapping[str, CacheReq] weight_source: Literal["views", "convert"] # Model shape diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 6179bd17feac..4425d46f8872 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -12,7 +12,6 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.models.common.mono.spec import ( - CacheReq, Feature, MonoContext, MonoSpec, @@ -486,7 +485,6 @@ def _layers(config) -> list[int]: MONO_SPEC = MonoSpec[LaunchState]( name="dsv41-cdna4", supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), - caches={"swa": CacheReq("swa", "fp8_ds_mla", "blocks,block,bytes")}, weight_source="views", layers=_layers, widths=_widths, From 7cf9683d8d979174ef0d8e9f5e722ce9e03c4f10 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 18:40:10 -0500 Subject: [PATCH 22/79] fix: remove all duplicate functionality Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 5 +- tests/models/common/mono/reference.py | 8 +-- vllm/models/common/mono/guards.py | 18 +----- vllm/models/common/mono/runtime.py | 62 +++++++++---------- vllm/models/common/mono/spec.py | 20 ++---- .../amd/mono/common/peer_memory.py | 23 +------ vllm/models/deepseek_v41/amd/mono_spec.py | 8 +-- 7 files changed, 42 insertions(+), 102 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 98d558cf638b..091c3f6f3d8c 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -165,9 +165,8 @@ def test_decided_once_a_forward(self, config, model, mono): def test_every_rank_decides_alike(self, config, model): decisions = [] - for rank in range(2): + for _ in range(2): other = attach(self.make_model(config), config, self.specs_path) - other.ctx = replace(other.ctx, rank=rank) with forward( config, CUDAGraphMode.NONE, self.decode_metadata(config, model) ): @@ -176,7 +175,7 @@ def test_every_rank_decides_alike(self, config, model): def test_skipped_step_runs_vllm_layers(self, config, model, mono): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): - mono.begin_once(self._positions(mono)) + mono.begin_step(self._positions(mono)) assert mono.width is None def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index c82444a20f17..9b31a0c40202 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -78,10 +78,6 @@ def compile(self, key: Any) -> None: self.compiled.append(key) -def layers(config) -> list[int]: - return list(range(config.num_hidden_layers)) - - def widths(vllm_config) -> tuple[int, ...]: most = vllm_config.scheduler_config.max_num_seqs return tuple(w for w in (1, 2, 4, 8) if w <= most) @@ -100,7 +96,7 @@ def refuse_model(model) -> list[str]: return [] if isinstance(model, ReferenceModel) else ["not a ReferenceModel"] -def skip_layer(ctx: MonoContext, layer) -> str | None: +def skip_layer(layer) -> str | None: return "dense layer" if layer.dense else None @@ -132,8 +128,6 @@ def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceK REFERENCE = MonoSpec[LaunchState]( name="reference", supports=frozenset(), - weight_source="views", - layers=layers, widths=widths, refuse_config=refuse_config, refuse_model=refuse_model, diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index 01faa37fcc6f..28a431398cda 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -3,7 +3,7 @@ """Refusals: the variant a config can serve, and TP agreement on every stage.""" import importlib -from collections.abc import Callable, Sequence +from collections.abc import Sequence from typing import TYPE_CHECKING, Any import torch @@ -192,19 +192,3 @@ def configure(vllm_config: "VllmConfig", specs_path: str) -> None: "CUDA graphs", spec.name, ) - - -def _agree_across_ranks(refuse: Callable[[], list[str]], group, what: str) -> None: - """``refuse``'s reasons on this rank, then every TP rank's verdict agreed: a - rank that refused alone would leave the others spinning in-kernel. - - Raises: - MonoUnsupported: This rank refused, or another rank did. - - """ - reasons = refuse() - agreed = tp_agree(not reasons, group) - if reasons: - raise MonoUnsupported(f"{what}: " + "; ".join(reasons)) - if not agreed: - raise MonoUnsupported(f"{what}: another TP rank refused") diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index c6c879049149..76d912472735 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -14,12 +14,12 @@ from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel from vllm.models.common.mono.guards import ( - _agree_across_ranks, _enabled, _load_specs, _refusal_message, _refuse_model, _select_spec, + tp_agree, tp_group, ) from vllm.models.common.mono.spec import ( @@ -60,18 +60,20 @@ class MonoRuntime: """A model's MonoKernels on this rank: built once after load.""" def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): - rank, _, cpu_group = tp_group() self.spec = spec - _agree_across_ranks( - lambda: _refuse_model(spec, model), - cpu_group, - f"MonoKernel {spec.name} refused for the loaded model", - ) - self.ctx = MonoContext(vllm_config, rank, next(model.parameters()).device) + reasons = _refuse_model(spec, model) + agreed = tp_agree(not reasons, tp_group()[2]) + refused = f"MonoKernel {spec.name} refused for the loaded model" + if reasons: + raise MonoUnsupported(f"{refused}: " + "; ".join(reasons)) + if not agreed: + raise MonoUnsupported(f"{refused}: another TP rank refused") + self.ctx = MonoContext(vllm_config, next(model.parameters()).device) self.widths = spec.widths(vllm_config) self.kernels: dict[tuple[int, int], LayerKernel] = {} self.state: Any = None self.width: int | None = None + self.why: str | None = None self._fc: Any = None self._checked: tuple | None = None @@ -100,10 +102,8 @@ def _kernel_layers(self, model) -> list[int]: start = getattr(model, "start_layer", 0) end = getattr(model, "end_layer", len(model.layers)) out = [] - for i in self.spec.layers(model.config): - if not start <= i < end: - continue - why = self.spec.skip_layer(self.ctx, model.layers[i]) + for i in range(start, end): + why = self.spec.skip_layer(model.layers[i]) if why is None: out.append(i) else: @@ -147,27 +147,25 @@ def _refuse_caches(self) -> None: ) self._checked = key - def begin_once(self, positions: torch.Tensor) -> None: - """Decides this forward's step, on the first kernel layer to ask.""" - fc = get_forward_context() - if fc is not self._fc: - self._fc = fc - self.begin_step(positions) - def begin_step(self, positions: torch.Tensor) -> str | None: - """None runs this step on the kernels (and sets ``width``), else the - reason it runs vLLM's layers.""" - self.width = None - if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: - return "piecewise" + """This forward's decision, taken on the first kernel layer to ask and + reused by the rest: None runs the step on the kernels (and sets + ``width``), else the reason it runs vLLM's layers.""" + fc = get_forward_context() + if fc is self._fc: + return self.why n = positions.size(-1) - if n not in self.widths: - return "width" - why = self.spec.skip_step(self.ctx, self.state, n) - if why is None: - self._refuse_caches() - self.width = n - return why + self._fc, self.width = fc, None + if fc.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + self.why = "piecewise" + elif n not in self.widths: + self.why = "width" + else: + self.why = self.spec.skip_step(self.ctx, self.state, n) + if self.why is None: + self._refuse_caches() + self.width = n + return self.why def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | None: @@ -209,7 +207,7 @@ def kernel_for_layer(mono: MonoRuntime | None, i: int, positions) -> LayerKernel """ if mono is None: return None - mono.begin_once(positions) + mono.begin_step(positions) if mono.width is None: return None return mono.kernels.get((i, mono.width)) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index eefbf240d499..0956da7443da 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -5,7 +5,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar import torch @@ -47,13 +47,11 @@ class MonoContext: Attributes: vllm_config: The engine's config. - rank: This rank's index in the TP group. device: The device the kernels launch on. """ vllm_config: "VllmConfig" - rank: int device: torch.device @@ -105,9 +103,6 @@ class MonoSpec(Generic[LaunchStateT]): name: Log and metric name. supports: Features this spec serves. Default deny: a feature in use that is absent from it refuses the deployment. - weight_source: ``"views"`` to read vLLM's post-load layout in place, - ``"convert"`` to build a second layout once after load. - layers: The decoder layers that get a kernel, from the model's config. widths: The rows a launch serves, from the engine's config. These become the CUDA graph capture sizes. refuse_config: Deployment refusals the engine's config answers, before @@ -116,7 +111,8 @@ class MonoSpec(Generic[LaunchStateT]): refuse_caches: Deployment refusals the KV caches the kernels will read answer, re-run whenever those caches move. skip_layer: Why one layer takes vLLM's path, None when it takes the - kernel's. + kernel's. Takes only the layer, so what it reads is what every rank + of the same model sees alike. skip_step: Why one step takes vLLM's path, None when it takes the kernel's. Takes the step's width; anything else it needs it reads from the forward context, which every rank sees alike. @@ -129,19 +125,11 @@ class MonoSpec(Generic[LaunchStateT]): name: str supports: frozenset[Feature] - weight_source: Literal["views", "convert"] - - # Model shape - layers: Callable[[Any], list[int]] widths: Callable[["VllmConfig"], tuple[int, ...]] - - # Capabilities refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] refuse_caches: Callable[[MonoContext, Mapping[str, Any]], list[str]] - skip_layer: Callable[[MonoContext, Any], str | None] + skip_layer: Callable[[Any], str | None] skip_step: Callable[[MonoContext, LaunchStateT, int], str | None] - - # Callbacks launch_state: Callable[[Any, MonoContext], LaunchStateT] build: Callable[[MonoContext, LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py index 58c95a567e07..afab63471e1c 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py +++ b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py @@ -4,32 +4,15 @@ # Copyright (C) 2026, Advanced Micro Devices, Inc.: # atom/mono/runtime/consensus.py # atom/mono/runtime/peer_memory.py -"""Symmetric peer memory for the in-kernel TP all-reduces, and the agreement every -rank must reach before anything collective (a rank that turned the mono layer -off alone would leave the others spinning on the GPU).""" +"""Symmetric peer memory for the in-kernel TP all-reduces.""" import torch -import torch.distributed as dist from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap +from vllm.models.common.mono.guards import tp_agree +from vllm.models.common.mono.spec import MonoUnsupported -# ---------------------------------------------------------------- TP agreement -class MonoUnsupported(Exception): - """The loaded model or runtime configuration is outside what mono serves.""" - -def tp_agree(ok: bool, group) -> bool: - """True when every rank of ``group`` passed ``ok`` true. Every rank must call - it at the same point; ``group`` is a CPU (gloo) process group, or None for a - single rank.""" - if group is None: - return ok - flag = torch.tensor([int(ok)], dtype=torch.int32) - dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) - return bool(flag.item()) - - -# ---------------------------------------------------------------- peer buffers class _DeviceBytes: """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a non-owning torch view.""" diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 4425d46f8872..826fd4db7088 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -290,7 +290,7 @@ def _refuse_model(model) -> list[str]: return reasons -def _skip_layer(ctx: MonoContext, layer) -> str | None: +def _skip_layer(layer) -> str | None: """Why this layer keeps vLLM's path for every step, or None when it takes a launch (whole or FFN only: ``_build`` picks which).""" ffn = layer.ffn @@ -478,15 +478,9 @@ def _build( return kernel -def _layers(config) -> list[int]: - return list(range(config.num_hidden_layers)) - - MONO_SPEC = MonoSpec[LaunchState]( name="dsv41-cdna4", supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), - weight_source="views", - layers=_layers, widths=_widths, refuse_config=_refuse_config, refuse_model=_refuse_model, From 14fdc8fae1e3f5b2baf48e5efc9ac7f394cabcba Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 19:27:53 -0500 Subject: [PATCH 23/79] fix: reference kernel Signed-off-by: simondanielsson --- tests/models/common/mono/reference.py | 30 ++++++++++++--------------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 9b31a0c40202..8a59d357b54f 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -12,10 +12,8 @@ from vllm.forward_context import get_forward_context from vllm.models.common.mono.spec import MonoContext, MonoSpec -CACHE_DTYPE = "bfloat16" -CACHE_LAYOUT = "2nhd" - -_LAYOUT_DIMS = {"2nhd": 4} +CACHE_DTYPE = torch.bfloat16 +CACHE_NDIM = 4 class ReferenceLayer(torch.nn.Module): @@ -44,7 +42,6 @@ def __init__(self, layers: int = 4, hidden: int = 8): ) self.start_layer = 0 self.end_layer = layers - self.mono = None @dataclass @@ -56,26 +53,25 @@ class LaunchState: class ReferenceKernel: - """One layer's launch at one width: calls the layer it replaces.""" + """One layer's launch at one width: the work of the layer it replaces, off + the weight view its launch state holds.""" def __init__(self, ctx: MonoContext, state: LaunchState, layer, width: int): self.cache_layers = (layer.attn_name,) - self.ctx = ctx - self.state = state - self.layer = layer + self.index = layer.index + self.weight = state.weights[layer.index] self.width = width self.launched: list[int] = [] - self.compiled: list[Any] = [] - def launch(self, positions: torch.Tensor, *args, **kwargs): + def launch(self, positions: torch.Tensor, hidden: torch.Tensor): self.launched.append(self.width) - return self.layer(positions, *args, **kwargs) + return hidden + self.weight def warmup_keys(self) -> list[Any]: - return [(self.layer.index, self.width)] + return [(self.index, self.width)] def compile(self, key: Any) -> None: - self.compiled.append(key) + pass def widths(vllm_config) -> tuple[int, ...]: @@ -114,10 +110,10 @@ def launch_state(model, ctx: MonoContext) -> LaunchState: def refuse_caches(ctx: MonoContext, caches) -> list[str]: reasons = [] for name, cache in caches.items(): - if str(cache.dtype).removeprefix("torch.") != CACHE_DTYPE: + if cache.dtype is not CACHE_DTYPE: reasons.append(f"{name}: dtype {cache.dtype}, needs {CACHE_DTYPE}") - if cache.ndim != _LAYOUT_DIMS[CACHE_LAYOUT]: - reasons.append(f"{name}: {cache.ndim}d, needs layout {CACHE_LAYOUT}") + if cache.ndim != CACHE_NDIM: + reasons.append(f"{name}: {cache.ndim}d, needs {CACHE_NDIM}d") return reasons From 5c54b691bc18743bbb833e31582c76ef175ce3d8 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 19:43:58 -0500 Subject: [PATCH 24/79] fix: delete monocontext Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 4 ++- tests/models/common/mono/reference.py | 14 +++++----- tests/models/test_dsv41_mono_decode.py | 10 +++---- vllm/models/common/mono/runtime.py | 19 +++++-------- vllm/models/common/mono/spec.py | 33 +++++++---------------- vllm/models/deepseek_v41/amd/mono_spec.py | 21 +++++++-------- 6 files changed, 40 insertions(+), 61 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 091c3f6f3d8c..d2c374bc302e 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -202,7 +202,9 @@ def test_caches_rechecked_when_they_move(self, config, model, mono): assert len(checks) == 2 def test_refuse_caches_reports_every_reason(self, config, model, mono): - reasons = mono.spec.refuse_caches(mono.ctx, self.bad_caches(config, model)) + reasons = mono.spec.refuse_caches( + mono.vllm_config, self.bad_caches(config, model) + ) assert len(reasons) >= 2, reasons def test_mismatched_caches_refuse_the_step(self, config, model, mono): diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 8a59d357b54f..8656a389253e 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -10,7 +10,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context -from vllm.models.common.mono.spec import MonoContext, MonoSpec +from vllm.models.common.mono.spec import MonoSpec CACHE_DTYPE = torch.bfloat16 CACHE_NDIM = 4 @@ -56,7 +56,7 @@ class ReferenceKernel: """One layer's launch at one width: the work of the layer it replaces, off the weight view its launch state holds.""" - def __init__(self, ctx: MonoContext, state: LaunchState, layer, width: int): + def __init__(self, vllm_config, state: LaunchState, layer, width: int): self.cache_layers = (layer.attn_name,) self.index = layer.index self.weight = state.weights[layer.index] @@ -96,18 +96,18 @@ def skip_layer(layer) -> str | None: return "dense layer" if layer.dense else None -def skip_step(ctx: MonoContext, state: LaunchState, width: int) -> str | None: +def skip_step(vllm_config, state: LaunchState, width: int) -> str | None: md = get_forward_context().attn_metadata return "prefill" if md.get("prefill") else None -def launch_state(model, ctx: MonoContext) -> LaunchState: +def launch_state(model, vllm_config) -> LaunchState: return LaunchState( weights={i: layer.weight.data for i, layer in enumerate(model.layers)} ) -def refuse_caches(ctx: MonoContext, caches) -> list[str]: +def refuse_caches(vllm_config, caches) -> list[str]: reasons = [] for name, cache in caches.items(): if cache.dtype is not CACHE_DTYPE: @@ -117,8 +117,8 @@ def refuse_caches(ctx: MonoContext, caches) -> list[str]: return reasons -def build(ctx: MonoContext, state: LaunchState, layer, width: int) -> ReferenceKernel: - return ReferenceKernel(ctx, state, layer, width) +def build(vllm_config, state: LaunchState, layer, width: int) -> ReferenceKernel: + return ReferenceKernel(vllm_config, state, layer, width) REFERENCE = MonoSpec[LaunchState]( diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 8914ba7a6119..7a814ae8fcd7 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -369,7 +369,7 @@ def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): def _step(mono, model, metadata, width=1): - with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, metadata): + with forward(mono.vllm_config, CUDAGraphMode.FULL, metadata): return mono.begin_step(torch.zeros(width, dtype=torch.long)) @@ -466,7 +466,7 @@ def test_off_without_the_flag(monkeypatch): def test_whole_launch_runs_the_layer(monkeypatch): mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 1, positions) out = kernel.launch( torch.zeros(1, HIDDEN), @@ -488,7 +488,7 @@ def test_first_layer_takes_the_embedding(monkeypatch): K1 takes as it is (no owed post to fold).""" mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 0, positions) out = kernel.launch(torch.zeros(1, HIDDEN), positions, None, None, None, None) assert out[0] == "whole" and not out[2]["fold"] @@ -498,7 +498,7 @@ def test_first_layer_takes_the_embedding(monkeypatch): def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): mono, model = _attach(monkeypatch, compressor=object()) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 1, positions) assert kernel.ffn_only out = kernel.launch( @@ -516,7 +516,7 @@ def test_launch_declines_a_layout_it_cannot_read(monkeypatch): ``skip_step`` cannot see that, so the launch declines it.""" mono, model = _attach(monkeypatch) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.ctx.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): kernel = kernel_for_layer(mono, 1, positions) out = kernel.launch( torch.zeros(1, 4, HIDDEN), positions, None, None, None, None diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 76d912472735..ca0262ade25d 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -22,12 +22,7 @@ tp_agree, tp_group, ) -from vllm.models.common.mono.spec import ( - LayerKernel, - MonoContext, - MonoSpec, - MonoUnsupported, -) +from vllm.models.common.mono.spec import LayerKernel, MonoSpec, MonoUnsupported if TYPE_CHECKING: from vllm.config import VllmConfig @@ -68,7 +63,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): raise MonoUnsupported(f"{refused}: " + "; ".join(reasons)) if not agreed: raise MonoUnsupported(f"{refused}: another TP rank refused") - self.ctx = MonoContext(vllm_config, next(model.parameters()).device) + self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) self.kernels: dict[tuple[int, int], LayerKernel] = {} self.state: Any = None @@ -81,11 +76,11 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): if not indices: self.cache_layers: list[str] = [] return - self.state = spec.launch_state(model, self.ctx) + self.state = spec.launch_state(model, vllm_config) for i in indices: for w in self.widths: self.kernels[i, w] = spec.build( - self.ctx, self.state, model.layers[i], w + vllm_config, self.state, model.layers[i], w ) self.cache_layers = sorted( {n for k in self.kernels.values() for n in k.cache_layers} @@ -127,7 +122,7 @@ def _refuse_caches(self) -> None: """Re-checks the caches whenever they moved (keyed on every tensor's data_ptr and numel): memory profiling runs real steps on throwaway caches before the real ones exist.""" - sfc = self.ctx.vllm_config.compilation_config.static_forward_context + sfc = self.vllm_config.compilation_config.static_forward_context caches = {n: sfc[n].kv_cache for n in self.cache_layers} key = tuple( (t.data_ptr(), t.numel()) @@ -140,7 +135,7 @@ def _refuse_caches(self) -> None: raise MonoUnsupported( f"MonoKernel {self.spec.name}: caches moved during CUDA graph capture" ) - reasons = self.spec.refuse_caches(self.ctx, caches) + reasons = self.spec.refuse_caches(self.vllm_config, caches) if reasons: raise MonoUnsupported( f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) @@ -161,7 +156,7 @@ def begin_step(self, positions: torch.Tensor) -> str | None: elif n not in self.widths: self.why = "width" else: - self.why = self.spec.skip_step(self.ctx, self.state, n) + self.why = self.spec.skip_step(self.vllm_config, self.state, n) if self.why is None: self._refuse_caches() self.width = n diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 0956da7443da..3a559df5ebc5 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -7,8 +7,6 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar -import torch - if TYPE_CHECKING: from vllm.config import VllmConfig @@ -41,20 +39,6 @@ class MonoUnsupported(Exception): """The configuration or the loaded model is outside what a spec serves.""" -@dataclass(frozen=True) -class MonoContext: - """One rank's view of the deployment, built once at attach. - - Attributes: - vllm_config: The engine's config. - device: The device the kernels launch on. - - """ - - vllm_config: "VllmConfig" - device: torch.device - - class LayerKernel(Protocol): """One decoder layer's launch, at one width, on this rank.""" @@ -116,10 +100,11 @@ class MonoSpec(Generic[LaunchStateT]): skip_step: Why one step takes vLLM's path, None when it takes the kernel's. Takes the step's width; anything else it needs it reads from the forward context, which every rank sees alike. - launch_state: The weights and launch state, built once at attach. The - only place a weight copy is ever made. Opaque to the library, which - only hands it back to the spec's own callbacks. - build: One ``LayerKernel`` for (context, state, layer, width). + launch_state: The weights and launch state, built once at attach, off + the loaded model (whose parameters carry the device the kernels + launch on). The only place a weight copy is ever made. Opaque to + the library, which only hands it back to the spec's own callbacks. + build: One ``LayerKernel`` for (config, state, layer, width). """ @@ -128,8 +113,8 @@ class MonoSpec(Generic[LaunchStateT]): widths: Callable[["VllmConfig"], tuple[int, ...]] refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] - refuse_caches: Callable[[MonoContext, Mapping[str, Any]], list[str]] + refuse_caches: Callable[["VllmConfig", Mapping[str, Any]], list[str]] skip_layer: Callable[[Any], str | None] - skip_step: Callable[[MonoContext, LaunchStateT, int], str | None] - launch_state: Callable[[Any, MonoContext], LaunchStateT] - build: Callable[[MonoContext, LaunchStateT, Any, int], LayerKernel] + skip_step: Callable[["VllmConfig", LaunchStateT, int], str | None] + launch_state: Callable[[Any, "VllmConfig"], LaunchStateT] + build: Callable[["VllmConfig", LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 826fd4db7088..8d59ac1b9644 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -11,11 +11,7 @@ from vllm.config import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.common.mono.spec import ( - Feature, - MonoContext, - MonoSpec, -) +from vllm.models.common.mono.spec import Feature, MonoSpec from vllm.platforms import current_platform if TYPE_CHECKING: @@ -331,7 +327,7 @@ def _ffn_only(layer) -> str | None: return None -def _skip_step(ctx: MonoContext, launch_state: LaunchState, width: int) -> str | None: +def _skip_step(vllm_config, launch_state: LaunchState, width: int) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP group sees alike.""" @@ -367,17 +363,18 @@ def _skip_step(ctx: MonoContext, launch_state: LaunchState, width: int) -> str | return None -def _launch_state(model, ctx: MonoContext) -> LaunchState: +def _launch_state(model, vllm_config) -> LaunchState: """This rank's runner, collective over the TP group.""" from vllm.models.common.mono.guards import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS rank, world, cpu_group = tp_group() + device = next(model.parameters()).device return LaunchState( - runner=DSV41MonoLayer(world, rank, cpu_group, ctx.device), + runner=DSV41MonoLayer(world, rank, cpu_group, device), first=getattr(model, "start_layer", 0), - words=torch.zeros(world, dtype=torch.int32, device=ctx.device), + words=torch.zeros(world, dtype=torch.int32, device=device), ) @@ -435,7 +432,7 @@ def _attn_weights(a, world: int): return attn -def _refuse_caches(ctx: MonoContext, caches) -> list[str]: +def _refuse_caches(vllm_config, caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as [blocks, block, RECORD] bytes.""" reasons = [] @@ -453,9 +450,9 @@ def _refuse_caches(ctx: MonoContext, caches) -> list[str]: def _build( - ctx: MonoContext, launch_state: LaunchState, layer, width: int + vllm_config, launch_state: LaunchState, layer, width: int ) -> DSV41LayerKernel: - world = ctx.vllm_config.parallel_config.tensor_parallel_size + world = vllm_config.parallel_config.tensor_parallel_size why = _ffn_only(layer) if layer.attn.layer_id not in launch_state.weights: launch_state.weights[layer.attn.layer_id] = _weights( From 43eb5e4b0a92622a2cc6b133006ecaf89f7bf2d5 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 20:00:06 -0500 Subject: [PATCH 25/79] feat: parametrize the runtime on layerkernel Signed-off-by: simondanielsson --- vllm/models/common/mono/runtime.py | 28 ++++++++++++++++++--------- vllm/models/deepseek_v41/amd/model.py | 9 +++------ 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index ca0262ade25d..a4eacb8e2a59 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -5,7 +5,7 @@ """ from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast import torch @@ -29,6 +29,11 @@ logger = init_logger(__name__) +KernelT = TypeVar("KernelT", bound=LayerKernel) +"""The kernel type one spec's ``build`` returns. ``LayerKernel`` is only the +lower bound, so a model annotates its runtime with its own kernel class +(``MonoRuntime[MyLayerKernel]``) and reads what that class adds.""" + class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): """A spec's layer-kernel builds as vLLM JIT warm-up keys.""" @@ -51,7 +56,7 @@ def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: self._compiled_cache[compile_key] = True -class MonoRuntime: +class MonoRuntime(Generic[KernelT]): """A model's MonoKernels on this rank: built once after load.""" def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): @@ -65,7 +70,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): raise MonoUnsupported(f"{refused}: another TP rank refused") self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) - self.kernels: dict[tuple[int, int], LayerKernel] = {} + self.kernels: dict[tuple[int, int], KernelT] = {} self.state: Any = None self.width: int | None = None self.why: str | None = None @@ -79,8 +84,8 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.state = spec.launch_state(model, vllm_config) for i in indices: for w in self.widths: - self.kernels[i, w] = spec.build( - vllm_config, self.state, model.layers[i], w + self.kernels[i, w] = cast( + KernelT, spec.build(vllm_config, self.state, model.layers[i], w) ) self.cache_layers = sorted( {n for k in self.kernels.values() for n in k.cache_layers} @@ -163,7 +168,9 @@ def begin_step(self, positions: torch.Tensor) -> str | None: return self.why -def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | None: +def attach( + model, vllm_config: "VllmConfig", specs_path: str +) -> MonoRuntime[Any] | None: """Picks the variant, builds the launch state and this rank's layer kernels. Collective over the TP group. @@ -174,7 +181,8 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | N Returns: This rank's ``MonoRuntime``, or None with the switch off or when no layer - takes a kernel. + takes a kernel. Annotate it with the kernel class the spec's ``build`` + returns, which is what ``kernel_for_layer`` then hands back. Raises: MonoUnsupported: Every variant refused the config, or this rank (or @@ -186,14 +194,16 @@ def attach(model, vllm_config: "VllmConfig", specs_path: str) -> MonoRuntime | N spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) if spec is None: raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) - mono = MonoRuntime(spec, model, vllm_config) + mono: MonoRuntime[Any] = MonoRuntime(spec, model, vllm_config) if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) return None return mono -def kernel_for_layer(mono: MonoRuntime | None, i: int, positions) -> LayerKernel | None: +def kernel_for_layer( + mono: MonoRuntime[KernelT] | None, i: int, positions +) -> KernelT | None: """Layer ``i``'s kernel when this step runs on the kernels, else None. The first call under a forward context decides for every layer. The model diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index a364e6cf9eec..8e1c5f990616 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -277,12 +277,9 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - mono: MonoRuntime | None = None, + mono: MonoRuntime[DSV41LayerKernel] | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - kernel = typing.cast( - DSV41LayerKernel | None, - kernel_for_layer(mono, self.attn.layer_id, positions), - ) + kernel = kernel_for_layer(mono, self.attn.layer_id, positions) if kernel is not None and not kernel.ffn_only: out = kernel.launch( x, @@ -483,7 +480,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps self.vllm_config = vllm_config - self.mono: MonoRuntime | None = None + self.mono: MonoRuntime[DSV41LayerKernel] | None = None configure(vllm_config, SPECS_PATH) # Three aux streams: one per non-default input GEMM in From af9947e67c957efaf51e59dcedf1e7ac2af675f1 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 20:13:50 -0500 Subject: [PATCH 26/79] fix: cleanup the runtime Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 11 ++++--- tests/models/test_dsv41_mono_decode.py | 2 +- vllm/models/common/mono/runtime.py | 40 ++++++++++++++------------ 3 files changed, 27 insertions(+), 26 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index d2c374bc302e..53e8ae99157c 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -170,23 +170,22 @@ def test_every_rank_decides_alike(self, config, model): with forward( config, CUDAGraphMode.NONE, self.decode_metadata(config, model) ): - decisions.append(other.begin_step(self._positions(other))) + decisions.append(other.skip_reason(self._positions(other))) assert decisions[0] == decisions[1] is None def test_skipped_step_runs_vllm_layers(self, config, model, mono): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): - mono.begin_step(self._positions(mono)) - assert mono.width is None + assert mono.skip_reason(self._positions(mono)) is not None def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): off = max(mono.widths) + 1 with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - assert mono.begin_step(torch.zeros(off, dtype=torch.long)) == "width" + assert mono.skip_reason(torch.zeros(off, dtype=torch.long)) == "width" def test_piecewise_runs_vllm_layers(self, config, model, mono): meta = self.decode_metadata(config, model) with forward(config, CUDAGraphMode.PIECEWISE, meta): - assert mono.begin_step(self._positions(mono)) == "piecewise" + assert mono.skip_reason(self._positions(mono)) == "piecewise" def test_caches_rechecked_when_they_move(self, config, model, mono): checks: list = [] @@ -239,7 +238,7 @@ def _positions(self, mono) -> torch.Tensor: def _decode(self, config, model, mono) -> None: with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - mono.begin_step(self._positions(mono)) + mono.skip_reason(self._positions(mono)) def _run_layers(self, mono, model) -> None: hidden = torch.zeros(min(mono.widths), model.config.hidden_size) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 7a814ae8fcd7..c5bc9b05b89a 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -370,7 +370,7 @@ def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): def _step(mono, model, metadata, width=1): with forward(mono.vllm_config, CUDAGraphMode.FULL, metadata): - return mono.begin_step(torch.zeros(width, dtype=torch.long)) + return mono.skip_reason(torch.zeros(width, dtype=torch.long)) @pytest.mark.parametrize( diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index a4eacb8e2a59..c3486ad6b1b0 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -72,8 +72,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.widths = spec.widths(vllm_config) self.kernels: dict[tuple[int, int], KernelT] = {} self.state: Any = None - self.width: int | None = None - self.why: str | None = None + self._why: str | None = None self._fc: Any = None self._checked: tuple | None = None @@ -147,25 +146,31 @@ def _refuse_caches(self) -> None: ) self._checked = key - def begin_step(self, positions: torch.Tensor) -> str | None: - """This forward's decision, taken on the first kernel layer to ask and - reused by the rest: None runs the step on the kernels (and sets - ``width``), else the reason it runs vLLM's layers.""" + def skip_reason(self, positions: torch.Tensor) -> str | None: + """Why this forward runs vLLM's layers, None when it runs on the + kernels. Decided on the first kernel layer to ask and reused by the + rest, so every layer of one forward is answered alike. + + Raises: + MonoUnsupported: The caches the kernels would read moved to + something this spec cannot serve. + + """ fc = get_forward_context() if fc is self._fc: - return self.why + return self._why n = positions.size(-1) - self._fc, self.width = fc, None + why: str | None if fc.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: - self.why = "piecewise" + why = "piecewise" elif n not in self.widths: - self.why = "width" + why = "width" else: - self.why = self.spec.skip_step(self.vllm_config, self.state, n) - if self.why is None: + why = self.spec.skip_step(self.vllm_config, self.state, n) + if why is None: self._refuse_caches() - self.width = n - return self.why + self._fc, self._why = fc, why + return why def attach( @@ -210,9 +215,6 @@ def kernel_for_layer( calls ``launch`` where it suits the layer, and runs the vLLM layer unchanged when this returns None. """ - if mono is None: - return None - mono.begin_step(positions) - if mono.width is None: + if mono is None or mono.skip_reason(positions) is not None: return None - return mono.kernels.get((i, mono.width)) + return mono.kernels.get((i, positions.size(-1))) From 255886a995d5a4e5ca2f20ecd9b94aff3bf3bc9c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 20:46:06 -0500 Subject: [PATCH 27/79] fix: add layer indices to build Signed-off-by: simondanielsson --- tests/models/common/mono/reference.py | 6 ++--- vllm/models/common/mono/runtime.py | 2 +- vllm/models/common/mono/spec.py | 7 ++--- vllm/models/deepseek_v41/amd/mono_spec.py | 33 ++++++++++++----------- 4 files changed, 25 insertions(+), 23 deletions(-) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 8656a389253e..7cada6f63ab7 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -101,10 +101,8 @@ def skip_step(vllm_config, state: LaunchState, width: int) -> str | None: return "prefill" if md.get("prefill") else None -def launch_state(model, vllm_config) -> LaunchState: - return LaunchState( - weights={i: layer.weight.data for i, layer in enumerate(model.layers)} - ) +def launch_state(model, vllm_config, layers: tuple[int, ...]) -> LaunchState: + return LaunchState(weights={i: model.layers[i].weight.data for i in layers}) def refuse_caches(vllm_config, caches) -> list[str]: diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index c3486ad6b1b0..67b51a8a3b2f 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -80,7 +80,7 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): if not indices: self.cache_layers: list[str] = [] return - self.state = spec.launch_state(model, vllm_config) + self.state = spec.launch_state(model, vllm_config, tuple(indices)) for i in indices: for w in self.widths: self.kernels[i, w] = cast( diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 3a559df5ebc5..0c3ae3c889ad 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -102,8 +102,9 @@ class MonoSpec(Generic[LaunchStateT]): from the forward context, which every rank sees alike. launch_state: The weights and launch state, built once at attach, off the loaded model (whose parameters carry the device the kernels - launch on). The only place a weight copy is ever made. Opaque to - the library, which only hands it back to the spec's own callbacks. + launch on) and the indices ``skip_layer`` left, in order. The only + place a weight copy is ever made. Opaque to the library, which only + hands it back to the spec's own callbacks. build: One ``LayerKernel`` for (config, state, layer, width). """ @@ -116,5 +117,5 @@ class MonoSpec(Generic[LaunchStateT]): refuse_caches: Callable[["VllmConfig", Mapping[str, Any]], list[str]] skip_layer: Callable[[Any], str | None] skip_step: Callable[["VllmConfig", LaunchStateT, int], str | None] - launch_state: Callable[[Any, "VllmConfig"], LaunchStateT] + launch_state: Callable[[Any, "VllmConfig", tuple[int, ...]], LaunchStateT] build: Callable[["VllmConfig", LaunchStateT, Any, int], LayerKernel] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 8d59ac1b9644..31968852fa60 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -81,7 +81,7 @@ class LaunchState: """ runner: Any - first: int + embed_layer: int words: torch.Tensor weights: dict[int, Any] = field(default_factory=dict) swa_keys: list[str] = field(default_factory=list) @@ -131,7 +131,7 @@ def __init__( ) # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == launch_state.first or layer.engram is not None + settled = attn.layer_id == launch_state.embed_layer or layer.engram is not None self.folds = (True, False) if settled else (True,) def warmup_keys(self) -> list[Any]: @@ -363,19 +363,30 @@ def _skip_step(vllm_config, launch_state: LaunchState, width: int) -> str | None return None -def _launch_state(model, vllm_config) -> LaunchState: - """This rank's runner, collective over the TP group.""" +def _launch_state(model, vllm_config, layers: tuple[int, ...]) -> LaunchState: + """This rank's runner, collective over the TP group, and the metadata keys + of the layers that take a whole-layer launch.""" from vllm.models.common.mono.guards import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS rank, world, cpu_group = tp_group() device = next(model.parameters()).device - return LaunchState( + state = LaunchState( runner=DSV41MonoLayer(world, rank, cpu_group, device), - first=getattr(model, "start_layer", 0), + embed_layer=getattr(model, "start_layer", 0), words=torch.zeros(world, dtype=torch.int32, device=device), ) + whole = [model.layers[i] for i in layers if _ffn_only(model.layers[i]) is None] + state.need_swa = bool(whole) + for layer in whole: + for key, keys in ( + (layer.attn.swa_cache_layer.prefix, state.swa_keys), + (layer.attn.compressed_cache_prefix, state.comp_keys), + ): + if key is not None and key not in keys: + keys.append(key) + return state def _weights(layer, ffn_only: bool, world: int): @@ -464,15 +475,7 @@ def _build( # reduces it on the steps that keep vLLM's path. layer.attn.wo_b.reduce_results = False return FfnSeam(launch_state, layer, width) - kernel = WholeLayer(launch_state, layer, width) - launch_state.need_swa = True - for key, keys in ( - (kernel.swa_key, launch_state.swa_keys), - (kernel.comp_key, launch_state.comp_keys), - ): - if key is not None and key not in keys: - keys.append(key) - return kernel + return WholeLayer(launch_state, layer, width) MONO_SPEC = MonoSpec[LaunchState]( From ee150461d290c8018561d8a1f3c8498b0d3ad29b Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 20:54:51 -0500 Subject: [PATCH 28/79] return the selected monospec Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 16 ++++++++----- tests/models/common/mono/test_reference.py | 14 ++++++------ tests/models/test_dsv41_mono_decode.py | 12 +++++----- vllm/models/common/mono/__init__.py | 2 ++ vllm/models/common/mono/guards.py | 14 +++++++++--- vllm/models/common/mono/runtime.py | 26 ++++++---------------- vllm/models/deepseek_v41/amd/model.py | 16 +++++++++---- 7 files changed, 57 insertions(+), 43 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 53e8ae99157c..f0175d4e86cb 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -71,7 +71,7 @@ class MonoConformance: module and fill in the hooks below.""" specs_path: str - """``"package.module:MONO_SPECS"``, as the model passes it to ``attach``.""" + """``"package.module:MONO_SPECS"``, as the model passes it to ``configure``.""" def make_config(self, **overrides) -> Any: """A config this spec serves, with ``overrides`` applied on top.""" @@ -104,7 +104,13 @@ def model(self, config): @pytest.fixture def mono(self, config, model): - return attach(model, config, self.specs_path) + return self.make_mono(config, model) + + def make_mono(self, config, model=None): + """This rank's runtime, built the way a model builds it: the variant + picked once by ``configure``, then handed to ``attach``.""" + model = self.make_model(config) if model is None else model + return attach(model, config, configure(config, self.specs_path)) def spec(self, config): spec, refusals = _select_spec(_load_specs(self.specs_path), config) @@ -166,7 +172,7 @@ def test_decided_once_a_forward(self, config, model, mono): def test_every_rank_decides_alike(self, config, model): decisions = [] for _ in range(2): - other = attach(self.make_model(config), config, self.specs_path) + other = self.make_mono(config) with forward( config, CUDAGraphMode.NONE, self.decode_metadata(config, model) ): @@ -217,14 +223,14 @@ def test_mismatched_caches_refuse_the_step(self, config, model, mono): def test_attach_does_not_move_weights(self, config): model = self.make_model(config) before = {n: p.data_ptr() for n, p in model.named_parameters()} - attach(model, config, self.specs_path) + self.make_mono(config, model) after = {n: p.data_ptr() for n, p in model.named_parameters()} assert before == after def test_warmup_registered_for_every_width(self, config): registry = JitWarmupRegistry(config) with registry.activate(): - mono = attach(self.make_model(config), config, self.specs_path) + mono = self.make_mono(config) assert len(registry) >= 1 by_width: dict[int, set] = defaultdict(set) for (_, width), kernel in mono.kernels.items(): diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index 8b62d1f59672..fbae85c8ae37 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -8,7 +8,7 @@ from tests.models.common.mono import reference from tests.models.common.mono.kit import MonoConformance, forward, forward_layer from vllm.config.compilation import CUDAGraphMode -from vllm.models.common.mono.guards import _enabled, _select_spec +from vllm.models.common.mono.guards import _select_spec, configure from vllm.models.common.mono.runtime import attach from vllm.models.common.mono.spec import MonoUnsupported @@ -39,8 +39,8 @@ def bad_caches(self, config, model): def test_switch_off_means_no_attach(): config = reference.make_config() config.kernel_config.monokernel = False - assert not _enabled(config) - assert attach(reference.make_model(config), config, SPECS_PATH) is None + assert configure(config, SPECS_PATH) is None + assert attach(reference.make_model(config), config, None) is None def test_the_first_servable_variant_wins(): @@ -59,7 +59,7 @@ def test_the_first_servable_variant_wins(): def test_skipped_layers_keep_vllm_path(): config = reference.make_config() model = reference.make_model(config) - mono = attach(model, config, SPECS_PATH) + mono = attach(model, config, configure(config, SPECS_PATH)) assert {i for i, _ in mono.kernels} == {1, 2, 3} @@ -67,14 +67,14 @@ def test_kernel_layers_follow_the_pipeline_slice(): config = reference.make_config() model = reference.make_model(config) model.start_layer, model.end_layer = 2, 4 - mono = attach(model, config, SPECS_PATH) + mono = attach(model, config, configure(config, SPECS_PATH)) assert {i for i, _ in mono.kernels} == {2, 3} def test_kernels_run_the_layers_they_replace(): config = reference.make_config() model = reference.make_model(config) - mono = attach(model, config, SPECS_PATH) + mono = attach(model, config, configure(config, SPECS_PATH)) width = min(mono.widths) positions = torch.zeros(width, dtype=torch.long) hidden = torch.zeros(width, model.config.hidden_size) @@ -88,4 +88,4 @@ def test_kernels_run_the_layers_they_replace(): def test_model_refusal_is_raised(): config = reference.make_config() with pytest.raises(MonoUnsupported, match="not a ReferenceModel"): - attach(torch.nn.Linear(2, 2), config, SPECS_PATH) + attach(torch.nn.Linear(2, 2), config, configure(config, SPECS_PATH)) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index c5bc9b05b89a..5cf7f40f20c6 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -13,6 +13,7 @@ from tests.models.common.mono.kit import MonoConformance, forward from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.models.common.mono.guards import configure from vllm.models.common.mono.runtime import attach, kernel_for_layer from vllm.models.common.mono.spec import MonoUnsupported from vllm.models.deepseek_v41.amd import mono_spec as ms @@ -291,12 +292,13 @@ def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS) -> None: def _attach(monkeypatch, config=None, model=None, **layer_overrides): _gpu(monkeypatch) + kit = TestDsv41Spec() config = config or make_config() - model = model if model is not None else TestDsv41Spec().make_model(config) + model = model if model is not None else kit.make_model(config) for layer in model.layers: for key, value in layer_overrides.items(): setattr(layer.attn, key, value) - return attach(model, config, SPECS_PATH), model + return kit.make_mono(config, model), model def _indexer(): @@ -434,7 +436,7 @@ def test_other_isa_refuses(monkeypatch, cdna): model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cdna=cdna) with pytest.raises(MonoUnsupported, match="CDNA4"): - attach(model, config, SPECS_PATH) + attach(model, config, configure(config, SPECS_PATH)) def test_too_few_compute_units_refuses(monkeypatch): @@ -442,7 +444,7 @@ def test_too_few_compute_units_refuses(monkeypatch): model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cus=BLOCKS // 2) with pytest.raises(MonoUnsupported, match="compute units"): - attach(model, config, SPECS_PATH) + attach(model, config, configure(config, SPECS_PATH)) def test_other_moe_backend_refuses(monkeypatch): @@ -453,7 +455,7 @@ def test_other_moe_backend_refuses(monkeypatch): experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON _gpu(monkeypatch) with pytest.raises(MonoUnsupported, match="AITER_MXFP4_BF16"): - attach(model, config, SPECS_PATH) + attach(model, config, configure(config, SPECS_PATH)) def test_off_without_the_flag(monkeypatch): diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index 510326126c91..b98bd5e4e26a 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -4,8 +4,10 @@ from vllm.models.common.mono.guards import configure from vllm.models.common.mono.runtime import attach, kernel_for_layer +from vllm.models.common.mono.spec import MonoSpec __all__ = [ + "MonoSpec", "attach", "configure", "kernel_for_layer", diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index 28a431398cda..6f55ca74e574 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -168,20 +168,27 @@ def _set_capture_sizes(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> None: cc.cudagraph_capture_sizes = new -def configure(vllm_config: "VllmConfig", specs_path: str) -> None: +def configure(vllm_config: "VllmConfig", specs_path: str) -> MonoSpec[Any] | None: """Picks the variant, sets the capture sizes, or raises with every reason. - A no-op with ``--monokernel`` off. + + The variant is picked here, once: the model holds what this returns and + hands it to ``attach`` after the weights are loaded. A model that prepares + anything of its own at load time reads the returned spec to know which + variant it is preparing for. Args: vllm_config: The engine's config, before the weights are loaded. specs_path: ``"package.module:MONO_SPECS"``. + Returns: + The variant this config runs, or None with ``--monokernel`` off. + Raises: MonoUnsupported: Every variant refused this config. """ if not _enabled(vllm_config): - return + return None spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) if spec is None: raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) @@ -192,3 +199,4 @@ def configure(vllm_config: "VllmConfig", specs_path: str) -> None: "CUDA graphs", spec.name, ) + return spec diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 67b51a8a3b2f..283edaf2066f 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -13,15 +13,7 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel -from vllm.models.common.mono.guards import ( - _enabled, - _load_specs, - _refusal_message, - _refuse_model, - _select_spec, - tp_agree, - tp_group, -) +from vllm.models.common.mono.guards import _refuse_model, tp_agree, tp_group from vllm.models.common.mono.spec import LayerKernel, MonoSpec, MonoUnsupported if TYPE_CHECKING: @@ -174,15 +166,15 @@ def skip_reason(self, positions: torch.Tensor) -> str | None: def attach( - model, vllm_config: "VllmConfig", specs_path: str + model, vllm_config: "VllmConfig", spec: MonoSpec[Any] | None ) -> MonoRuntime[Any] | None: - """Picks the variant, builds the launch state and this rank's layer - kernels. Collective over the TP group. + """Builds the launch state and this rank's layer kernels. Collective over + the TP group. Args: model: The loaded model, in ``process_weights_after_loading``. vllm_config: The engine's config. - specs_path: ``"package.module:MONO_SPECS"``. + spec: The variant ``configure`` picked, None with the switch off. Returns: This rank's ``MonoRuntime``, or None with the switch off or when no layer @@ -190,15 +182,11 @@ def attach( returns, which is what ``kernel_for_layer`` then hands back. Raises: - MonoUnsupported: Every variant refused the config, or this rank (or - another) refused the loaded model. + MonoUnsupported: This rank, or another, refused the loaded model. """ - if not _enabled(vllm_config): - return None - spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) if spec is None: - raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) + return None mono: MonoRuntime[Any] = MonoRuntime(spec, model, vllm_config) if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 8e1c5f990616..272ddcd8077d 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -50,7 +50,7 @@ make_layers, maybe_prefix, ) -from vllm.models.common.mono import attach, configure, kernel_for_layer +from vllm.models.common.mono import MonoSpec, attach, configure, kernel_for_layer from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, @@ -61,7 +61,11 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v41.amd.mono_spec import SPECS_PATH, DSV41LayerKernel +from vllm.models.deepseek_v41.amd.mono_spec import ( + SPECS_PATH, + DSV41LayerKernel, + LaunchState, +) from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers @@ -481,7 +485,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.rms_norm_eps = config.rms_norm_eps self.vllm_config = vllm_config self.mono: MonoRuntime[DSV41LayerKernel] | None = None - configure(vllm_config, SPECS_PATH) + self.mono_spec: MonoSpec[LaunchState] | None = configure( + vllm_config, SPECS_PATH + ) # Three aux streams: one per non-default input GEMM in # DeepseekV4Attention._run_parallel_input_projections @@ -1413,7 +1419,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() - self.model.mono = attach(self.model, self.model.vllm_config, SPECS_PATH) + self.model.mono = attach( + self.model, self.model.vllm_config, self.model.mono_spec + ) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() From d0aecb7b98642959931ae27acebf9012e0ffb443 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 21:02:45 -0500 Subject: [PATCH 29/79] fix: call build() once per layer, andnot per layer*width Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 11 ++++++++++- tests/models/common/mono/reference.py | 8 +++++--- vllm/models/common/mono/runtime.py | 13 +++++++++---- vllm/models/common/mono/spec.py | 20 +++++++++++++------- vllm/models/deepseek_v41/amd/mono_spec.py | 16 +++++++--------- 5 files changed, 44 insertions(+), 24 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index f0175d4e86cb..79068149d940 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -208,7 +208,7 @@ def test_caches_rechecked_when_they_move(self, config, model, mono): def test_refuse_caches_reports_every_reason(self, config, model, mono): reasons = mono.spec.refuse_caches( - mono.vllm_config, self.bad_caches(config, model) + mono.vllm_config, mono.state, self.bad_caches(config, model) ) assert len(reasons) >= 2, reasons @@ -227,6 +227,15 @@ def test_attach_does_not_move_weights(self, config): after = {n: p.data_ptr() for n, p in model.named_parameters()} assert before == after + def test_build_runs_once_a_layer(self, config, model): + calls: list = [] + spec = _spy(configure(config, self.specs_path), "build", calls) + mono = attach(model, config, spec) + layers = sorted({i for i, _ in mono.kernels}) + assert [args[2] for args in calls] == [model.layers[i] for i in layers] + assert all(set(args[3]) == set(mono.widths) for args in calls) + assert set(mono.kernels) == {(i, w) for i in layers for w in mono.widths} + def test_warmup_registered_for_every_width(self, config): registry = JitWarmupRegistry(config) with registry.activate(): diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 7cada6f63ab7..11c1cb2a2399 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -105,7 +105,7 @@ def launch_state(model, vllm_config, layers: tuple[int, ...]) -> LaunchState: return LaunchState(weights={i: model.layers[i].weight.data for i in layers}) -def refuse_caches(vllm_config, caches) -> list[str]: +def refuse_caches(vllm_config, state: LaunchState, caches) -> list[str]: reasons = [] for name, cache in caches.items(): if cache.dtype is not CACHE_DTYPE: @@ -115,8 +115,10 @@ def refuse_caches(vllm_config, caches) -> list[str]: return reasons -def build(vllm_config, state: LaunchState, layer, width: int) -> ReferenceKernel: - return ReferenceKernel(vllm_config, state, layer, width) +def build( + vllm_config, state: LaunchState, layer, widths: tuple[int, ...] +) -> dict[int, ReferenceKernel]: + return {w: ReferenceKernel(vllm_config, state, layer, w) for w in widths} REFERENCE = MonoSpec[LaunchState]( diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 283edaf2066f..9a22a73d858e 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -74,10 +74,15 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): return self.state = spec.launch_state(model, vllm_config, tuple(indices)) for i in indices: - for w in self.widths: - self.kernels[i, w] = cast( - KernelT, spec.build(vllm_config, self.state, model.layers[i], w) + built = spec.build(vllm_config, self.state, model.layers[i], self.widths) + missing = set(self.widths) - set(built) + if missing: + raise MonoUnsupported( + f"MonoKernel {spec.name}: layer {i} has no kernel for widths " + f"{sorted(missing)}" ) + for w in self.widths: + self.kernels[i, w] = cast(KernelT, built[w]) self.cache_layers = sorted( {n for k in self.kernels.values() for n in k.cache_layers} ) @@ -131,7 +136,7 @@ def _refuse_caches(self) -> None: raise MonoUnsupported( f"MonoKernel {self.spec.name}: caches moved during CUDA graph capture" ) - reasons = self.spec.refuse_caches(self.vllm_config, caches) + reasons = self.spec.refuse_caches(self.vllm_config, self.state, caches) if reasons: raise MonoUnsupported( f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 0c3ae3c889ad..734fa0612986 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -72,9 +72,9 @@ class MonoSpec(Generic[LaunchStateT]): as a ``MONO_SPECS`` tuple the library selects from. The parameter is the spec's own launch state: ``launch_state`` builds it, - and ``skip_step`` and ``build`` are handed back the same object. Annotate - the spec with it (``MonoSpec[MyLaunchState](...)``) and those three - callbacks are checked against each other. + and ``refuse_caches``, ``skip_step`` and ``build`` are handed back the same + object. Annotate the spec with it (``MonoSpec[MyLaunchState](...)``) and + those four callbacks are checked against each other. Gating runs in tiers, and the prefix is the consequence. ``refuse_*`` says the deployment asked for something it cannot have, so it may take the @@ -93,7 +93,8 @@ class MonoSpec(Generic[LaunchStateT]): the weights are loaded. refuse_model: Deployment refusals only the loaded model answers. refuse_caches: Deployment refusals the KV caches the kernels will read - answer, re-run whenever those caches move. + answer, re-run whenever those caches move. Takes the launch state, + so the caches are checked against what the kernels were built for. skip_layer: Why one layer takes vLLM's path, None when it takes the kernel's. Takes only the layer, so what it reads is what every rank of the same model sees alike. @@ -105,7 +106,10 @@ class MonoSpec(Generic[LaunchStateT]): launch on) and the indices ``skip_layer`` left, in order. The only place a weight copy is ever made. Opaque to the library, which only hands it back to the spec's own callbacks. - build: One ``LayerKernel`` for (config, state, layer, width). + build: One layer's kernels, by width, called once a layer with every + width the spec serves. What a layer prepares once (its weight + views, anything the layer itself must be told) is prepared here, + and the widths share it. Must return a kernel for every width. """ @@ -114,8 +118,10 @@ class MonoSpec(Generic[LaunchStateT]): widths: Callable[["VllmConfig"], tuple[int, ...]] refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] - refuse_caches: Callable[["VllmConfig", Mapping[str, Any]], list[str]] + refuse_caches: Callable[["VllmConfig", LaunchStateT, Mapping[str, Any]], list[str]] skip_layer: Callable[[Any], str | None] skip_step: Callable[["VllmConfig", LaunchStateT, int], str | None] launch_state: Callable[[Any, "VllmConfig", tuple[int, ...]], LaunchStateT] - build: Callable[["VllmConfig", LaunchStateT, Any, int], LayerKernel] + build: Callable[ + ["VllmConfig", LaunchStateT, Any, tuple[int, ...]], Mapping[int, LayerKernel] + ] diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 31968852fa60..0545b23ed005 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -443,7 +443,7 @@ def _attn_weights(a, world: int): return attn -def _refuse_caches(vllm_config, caches) -> list[str]: +def _refuse_caches(vllm_config, launch_state: LaunchState, caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as [blocks, block, RECORD] bytes.""" reasons = [] @@ -461,21 +461,19 @@ def _refuse_caches(vllm_config, caches) -> list[str]: def _build( - vllm_config, launch_state: LaunchState, layer, width: int -) -> DSV41LayerKernel: + vllm_config, launch_state: LaunchState, layer, widths: tuple[int, ...] +) -> dict[int, DSV41LayerKernel]: world = vllm_config.parallel_config.tensor_parallel_size why = _ffn_only(layer) - if layer.attn.layer_id not in launch_state.weights: - launch_state.weights[layer.attn.layer_id] = _weights( - layer, why is not None, world - ) + launch_state.weights[layer.attn.layer_id] = _weights(layer, why is not None, world) + kernel: type[DSV41LayerKernel] = WholeLayer if why is not None: logger.info_once("DSv4.1 MonoKernel: the FFN launch for layers with %s", why) # A decode step hands wo_b's TP partial to the FFN launch; the layer # reduces it on the steps that keep vLLM's path. layer.attn.wo_b.reduce_results = False - return FfnSeam(launch_state, layer, width) - return WholeLayer(launch_state, layer, width) + kernel = FfnSeam + return {w: kernel(launch_state, layer, w) for w in widths} MONO_SPEC = MonoSpec[LaunchState]( From eb05ec84b053b3231f7f68fd3227ae4611e05859 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 21:07:13 -0500 Subject: [PATCH 30/79] parameterize the runtime also by launch state type Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 2 +- tests/models/test_dsv41_mono_layers.py | 2 +- vllm/models/common/mono/runtime.py | 48 ++++++++++++++++++-------- vllm/models/deepseek_v41/amd/model.py | 4 +-- 4 files changed, 37 insertions(+), 19 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 79068149d940..82947d779a54 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -208,7 +208,7 @@ def test_caches_rechecked_when_they_move(self, config, model, mono): def test_refuse_caches_reports_every_reason(self, config, model, mono): reasons = mono.spec.refuse_caches( - mono.vllm_config, mono.state, self.bad_caches(config, model) + mono.vllm_config, mono.launch_state, self.bad_caches(config, model) ) assert len(reasons) >= 2, reasons diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index 7008346ffe3a..c3832d85f85e 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -206,7 +206,7 @@ def mono_layers(self) -> dict: ) for i in self._decoder_layers() } - runner = None if mono is None else mono.state.runner + runner = None if mono is None else mono.launch_state.runner return { "layers": layers, "epoch": -1 if runner is None else int(runner.epoch[0]), diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 9a22a73d858e..ebca792f3950 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -14,7 +14,12 @@ from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel from vllm.models.common.mono.guards import _refuse_model, tp_agree, tp_group -from vllm.models.common.mono.spec import LayerKernel, MonoSpec, MonoUnsupported +from vllm.models.common.mono.spec import ( + LaunchStateT, + LayerKernel, + MonoSpec, + MonoUnsupported, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -24,7 +29,8 @@ KernelT = TypeVar("KernelT", bound=LayerKernel) """The kernel type one spec's ``build`` returns. ``LayerKernel`` is only the lower bound, so a model annotates its runtime with its own kernel class -(``MonoRuntime[MyLayerKernel]``) and reads what that class adds.""" +(``MonoRuntime[MyLayerKernel, MyLaunchState]``) and reads what that class +adds.""" class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): @@ -48,10 +54,20 @@ def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: self._compiled_cache[compile_key] = True -class MonoRuntime(Generic[KernelT]): - """A model's MonoKernels on this rank: built once after load.""" +class MonoRuntime(Generic[KernelT, LaunchStateT]): + """A model's MonoKernels on this rank: built once after load. - def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): + ``launch_state`` is what the spec built at attach, and the one handle the + model reads outside a launch: anything a model needs between steps (the + runner the kernels share, its peer buffers, the error words) lives there + rather than in a module global. + """ + + launch_state: LaunchStateT + """What ``MonoSpec.launch_state`` built, set whenever a layer takes a + kernel, which is the only case ``attach`` hands the runtime back.""" + + def __init__(self, spec: MonoSpec[LaunchStateT], model, vllm_config: "VllmConfig"): self.spec = spec reasons = _refuse_model(spec, model) agreed = tp_agree(not reasons, tp_group()[2]) @@ -63,7 +79,6 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) self.kernels: dict[tuple[int, int], KernelT] = {} - self.state: Any = None self._why: str | None = None self._fc: Any = None self._checked: tuple | None = None @@ -72,9 +87,11 @@ def __init__(self, spec: MonoSpec[Any], model, vllm_config: "VllmConfig"): if not indices: self.cache_layers: list[str] = [] return - self.state = spec.launch_state(model, vllm_config, tuple(indices)) + self.launch_state = spec.launch_state(model, vllm_config, tuple(indices)) for i in indices: - built = spec.build(vllm_config, self.state, model.layers[i], self.widths) + built = spec.build( + vllm_config, self.launch_state, model.layers[i], self.widths + ) missing = set(self.widths) - set(built) if missing: raise MonoUnsupported( @@ -136,7 +153,7 @@ def _refuse_caches(self) -> None: raise MonoUnsupported( f"MonoKernel {self.spec.name}: caches moved during CUDA graph capture" ) - reasons = self.spec.refuse_caches(self.vllm_config, self.state, caches) + reasons = self.spec.refuse_caches(self.vllm_config, self.launch_state, caches) if reasons: raise MonoUnsupported( f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) @@ -163,7 +180,7 @@ def skip_reason(self, positions: torch.Tensor) -> str | None: elif n not in self.widths: why = "width" else: - why = self.spec.skip_step(self.vllm_config, self.state, n) + why = self.spec.skip_step(self.vllm_config, self.launch_state, n) if why is None: self._refuse_caches() self._fc, self._why = fc, why @@ -171,8 +188,8 @@ def skip_reason(self, positions: torch.Tensor) -> str | None: def attach( - model, vllm_config: "VllmConfig", spec: MonoSpec[Any] | None -) -> MonoRuntime[Any] | None: + model, vllm_config: "VllmConfig", spec: MonoSpec[LaunchStateT] | None +) -> MonoRuntime[Any, LaunchStateT] | None: """Builds the launch state and this rank's layer kernels. Collective over the TP group. @@ -184,7 +201,8 @@ def attach( Returns: This rank's ``MonoRuntime``, or None with the switch off or when no layer takes a kernel. Annotate it with the kernel class the spec's ``build`` - returns, which is what ``kernel_for_layer`` then hands back. + returns (which is what ``kernel_for_layer`` then hands back) and the + spec's launch state. Raises: MonoUnsupported: This rank, or another, refused the loaded model. @@ -192,7 +210,7 @@ def attach( """ if spec is None: return None - mono: MonoRuntime[Any] = MonoRuntime(spec, model, vllm_config) + mono: MonoRuntime[Any, LaunchStateT] = MonoRuntime(spec, model, vllm_config) if not mono.kernels: logger.info("MonoKernel %s: no layer takes a kernel", spec.name) return None @@ -200,7 +218,7 @@ def attach( def kernel_for_layer( - mono: MonoRuntime[KernelT] | None, i: int, positions + mono: MonoRuntime[KernelT, Any] | None, i: int, positions ) -> KernelT | None: """Layer ``i``'s kernel when this step runs on the kernels, else None. diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 272ddcd8077d..9b2e89b91848 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -281,7 +281,7 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - mono: MonoRuntime[DSV41LayerKernel] | None = None, + mono: MonoRuntime[DSV41LayerKernel, LaunchState] | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: kernel = kernel_for_layer(mono, self.attn.layer_id, positions) if kernel is not None and not kernel.ffn_only: @@ -484,7 +484,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps self.vllm_config = vllm_config - self.mono: MonoRuntime[DSV41LayerKernel] | None = None + self.mono: MonoRuntime[DSV41LayerKernel, LaunchState] | None = None self.mono_spec: MonoSpec[LaunchState] | None = configure( vllm_config, SPECS_PATH ) From 1bd3fd82943daeb04ab6fb1ca9a31b094a9d3c90 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 13 Oct 2026 21:21:48 -0500 Subject: [PATCH 31/79] chore: ccleanup docstrings Signed-off-by: simondanielsson --- vllm/models/common/mono/runtime.py | 13 ++++++++----- vllm/models/common/mono/spec.py | 30 ++++++------------------------ 2 files changed, 14 insertions(+), 29 deletions(-) diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index ebca792f3950..43c1630d07ef 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -55,12 +55,15 @@ def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: class MonoRuntime(Generic[KernelT, LaunchStateT]): - """A model's MonoKernels on this rank: built once after load. + """A model's MonoKernels on this rank. - ``launch_state`` is what the spec built at attach, and the one handle the - model reads outside a launch: anything a model needs between steps (the - runner the kernels share, its peer buffers, the error words) lives there - rather than in a module global. + The runtime reads the spec and builds the layer kernels. It also + registers them for warmup and handles caching, and drivers the + per-step decision whether to run the kernels or the model's layers. + + To build the kernels, it first creates the necessary launch state + LaunchStateT. The kernels built by the runtime are of type KernelT + and are stored in the `kernels` attribute. """ launch_state: LaunchStateT diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index 734fa0612986..dd331f607c6b 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -47,9 +47,10 @@ class LayerKernel(Protocol): def launch(self, *args, **kwargs) -> Any: """Takes and returns what the vLLM layer it replaces does, or None to - decline this call, which sends the layer down vLLM's path. Decline only - on what ``skip_step`` cannot answer, such as the layout of the tensors - the layer was handed.""" + decline this call. + + The decline should only trigger on what ``skip_step`` cannot answer, + such as the layout of the tensors the layer was handed.""" ... def warmup_keys(self) -> list[Any]: @@ -64,24 +65,7 @@ def compile(self, key: Any) -> None: @dataclass(frozen=True) class MonoSpec(Generic[LaunchStateT]): - """One (model, vendor, ISA)'s declaration: values and callbacks, no - behaviour of its own. - - A spec is a value, so a second ISA is ``dataclasses.replace`` of the first - at module scope rather than a subclass, and a model exports its candidates - as a ``MONO_SPECS`` tuple the library selects from. - - The parameter is the spec's own launch state: ``launch_state`` builds it, - and ``refuse_caches``, ``skip_step`` and ``build`` are handed back the same - object. Annotate the spec with it (``MonoSpec[MyLaunchState](...)``) and - those four callbacks are checked against each other. - - Gating runs in tiers, and the prefix is the consequence. ``refuse_*`` says - the deployment asked for something it cannot have, so it may take the - engine down; ``skip_*`` says only that this layer, or this step, takes - vLLM's path, and never raises. If the same model with the same flags could - legitimately have another layer, or a later step, that passes the check, it - is a ``skip_*``. Every tier reports all of its reasons, never the first. + """Definition of when a monokernel can run, and how to build it. Attributes: name: Log and metric name. @@ -103,9 +87,7 @@ class MonoSpec(Generic[LaunchStateT]): from the forward context, which every rank sees alike. launch_state: The weights and launch state, built once at attach, off the loaded model (whose parameters carry the device the kernels - launch on) and the indices ``skip_layer`` left, in order. The only - place a weight copy is ever made. Opaque to the library, which only - hands it back to the spec's own callbacks. + launch on). The only place a weight copy is ever made. build: One layer's kernels, by width, called once a layer with every width the spec serves. What a layer prepares once (its weight views, anything the layer itself must be told) is prepared here, From 06914f29dcb701018ec5d8c6106a7accadff6f83 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 14:50:23 -0500 Subject: [PATCH 32/79] refactor: convert the monoruntime to one per model instead of one per layer Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 123 +++++----- tests/models/common/mono/reference.py | 82 +++---- tests/models/common/mono/test_reference.py | 42 ++-- tests/models/test_dsv41_mono_decode.py | 67 +++--- vllm/models/common/mono/__init__.py | 14 +- vllm/models/common/mono/guards.py | 6 +- vllm/models/common/mono/runtime.py | 247 +++++++------------- vllm/models/common/mono/spec.py | 72 +++--- vllm/models/deepseek_v41/amd/model.py | 43 ++-- vllm/models/deepseek_v41/amd/model_state.py | 20 ++ vllm/models/deepseek_v41/amd/mono_spec.py | 216 ++++++++++------- 11 files changed, 462 insertions(+), 470 deletions(-) create mode 100644 vllm/models/deepseek_v41/amd/model_state.py diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index 82947d779a54..a79b342ce59b 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -4,7 +4,6 @@ import subprocess import sys -from collections import defaultdict from collections.abc import Iterator from contextlib import contextmanager from dataclasses import replace @@ -21,9 +20,9 @@ _refuse_config, _select_spec, _set_capture_sizes, - configure, + select_mono_spec, ) -from vllm.models.common.mono.runtime import attach, kernel_for_layer +from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.mono.spec import Feature, MonoUnsupported KERNEL_BACKENDS = ("aiter", "flydsl") @@ -44,17 +43,6 @@ def forward(config, mode: CUDAGraphMode, attn_metadata: Any) -> Iterator[None]: yield -def forward_layer(mono, i: int, layer, positions, *args, **kwargs): - """Layer ``i`` on its kernel when this step runs on the kernels, else the - vLLM layer unchanged: the layer loop a model writes for itself.""" - k = kernel_for_layer(mono, i, positions) - if k is not None: - out = k.launch(positions, *args, **kwargs) - if out is not None: - return out - return layer(positions, *args, **kwargs) - - def _spy(spec, field: str, calls: list): """``spec`` with one callback recording its calls.""" inner = getattr(spec, field) @@ -71,7 +59,8 @@ class MonoConformance: module and fill in the hooks below.""" specs_path: str - """``"package.module:MONO_SPECS"``, as the model passes it to ``configure``.""" + """``"package.module:MONO_SPECS"``, as the model passes it to + ``select_mono_spec``.""" def make_config(self, **overrides) -> Any: """A config this spec serves, with ``overrides`` applied on top.""" @@ -106,11 +95,17 @@ def model(self, config): def mono(self, config, model): return self.make_mono(config, model) - def make_mono(self, config, model=None): + def make_mono(self, config, model=None, bind: bool = True): """This rank's runtime, built the way a model builds it: the variant - picked once by ``configure``, then handed to ``attach``.""" + picked once by ``select_mono_spec``, then the runtime over the loaded model, + then the caches the runner binds.""" model = self.make_model(config) if model is None else model - return attach(model, config, configure(config, self.specs_path)) + spec = select_mono_spec(config, self.specs_path) + assert spec is not None + mono = MonoRuntime(model, config, spec) + if bind: + mono.bind_caches() + return mono def spec(self, config): spec, refusals = _select_spec(_load_specs(self.specs_path), config) @@ -131,7 +126,7 @@ def test_refusals_report_every_reason(self, config): def test_refusal_names_every_variant(self, config): config.use_v2_model_runner = False with pytest.raises(MonoUnsupported) as excinfo: - configure(config, self.specs_path) + select_mono_spec(config, self.specs_path) for spec in _load_specs(self.specs_path): assert spec.name in str(excinfo.value) @@ -161,102 +156,98 @@ def test_capture_sizes_are_the_widths(self, config): sizes = config.compilation_config.cudagraph_capture_sizes assert [s for s in sizes if s <= max(widths)] == sorted(widths) - def test_decided_once_a_forward(self, config, model, mono): + def test_decided_once_a_step(self, config, model, mono): asked: list = [] mono.spec = _spy(mono.spec, "skip_step", asked) - with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - self._run_layers(mono, model) + step = self._plan(config, model, mono) + assert step is not None and step.width == min(mono.widths) assert len(asked) == 1 - assert len(mono.kernels) > 1 def test_every_rank_decides_alike(self, config, model): decisions = [] for _ in range(2): other = self.make_mono(config) - with forward( - config, CUDAGraphMode.NONE, self.decode_metadata(config, model) - ): - decisions.append(other.skip_reason(self._positions(other))) - assert decisions[0] == decisions[1] is None + decisions.append(self._plan(config, model, other)) + assert decisions[0] is not None and decisions[1] is not None def test_skipped_step_runs_vllm_layers(self, config, model, mono): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): - assert mono.skip_reason(self._positions(mono)) is not None + assert mono.plan(self._positions(mono)) is None + assert mono.last_reason is not None + + def test_pipeline_stage_input_runs_vllm_layers(self, config, model, mono): + with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): + assert mono.plan(self._positions(mono), object()) is None + assert mono.last_reason == "pipeline stage input" def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): - off = max(mono.widths) + 1 + off = torch.zeros(max(mono.widths) + 1, dtype=torch.long) with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - assert mono.skip_reason(torch.zeros(off, dtype=torch.long)) == "width" + assert mono.plan(off) is None + assert mono.last_reason == "width" def test_piecewise_runs_vllm_layers(self, config, model, mono): meta = self.decode_metadata(config, model) with forward(config, CUDAGraphMode.PIECEWISE, meta): - assert mono.skip_reason(self._positions(mono)) == "piecewise" + assert mono.plan(self._positions(mono)) is None + assert mono.last_reason == "piecewise" - def test_caches_rechecked_when_they_move(self, config, model, mono): + def test_unbound_caches_run_vllm_layers(self, config, model): + mono = self.make_mono(config, model, bind=False) + assert self._plan(config, model, mono) is None + assert mono.last_reason == "caches not bound" + + def test_caches_checked_on_every_bind(self, config, model, mono): checks: list = [] mono.spec = _spy(mono.spec, "refuse_caches", checks) - self._decode(config, model, mono) - assert len(checks) == 1 - self._decode(config, model, mono) + mono.bind_caches() assert len(checks) == 1 sfc = config.compilation_config.static_forward_context - for name in mono.cache_layers: + for name in mono.kernel.cache_layers: sfc[name].kv_cache = sfc[name].kv_cache.clone() - self._decode(config, model, mono) + mono.bind_caches() assert len(checks) == 2 def test_refuse_caches_reports_every_reason(self, config, model, mono): reasons = mono.spec.refuse_caches( - mono.vllm_config, mono.launch_state, self.bad_caches(config, model) + mono.vllm_config, mono.kernel, self.bad_caches(config, model) ) assert len(reasons) >= 2, reasons - def test_mismatched_caches_refuse_the_step(self, config, model, mono): + def test_mismatched_caches_refuse_the_bind(self, config, model, mono): sfc = config.compilation_config.static_forward_context bad = self.bad_caches(config, model) - for name, cache in zip(mono.cache_layers, bad.values(), strict=False): + for name, cache in zip(mono.kernel.cache_layers, bad.values(), strict=False): sfc[name].kv_cache = cache with pytest.raises(MonoUnsupported): - self._decode(config, model, mono) + mono.bind_caches() - def test_attach_does_not_move_weights(self, config): + 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_mono(config, model) after = {n: p.data_ptr() for n, p in model.named_parameters()} assert before == after - def test_build_runs_once_a_layer(self, config, model): + def test_build_runs_once_over_every_width(self, config, model): calls: list = [] - spec = _spy(configure(config, self.specs_path), "build", calls) - mono = attach(model, config, spec) - layers = sorted({i for i, _ in mono.kernels}) - assert [args[2] for args in calls] == [model.layers[i] for i in layers] - assert all(set(args[3]) == set(mono.widths) for args in calls) - assert set(mono.kernels) == {(i, w) for i in layers for w in mono.widths} - - def test_warmup_registered_for_every_width(self, config): + spec = _spy(select_mono_spec(config, self.specs_path), "build", calls) + mono = MonoRuntime(model, config, spec) + assert [args[1] for args in calls] == [model] + assert set(calls[0][2]) == set(mono.widths) + + def test_warmup_registered_for_every_key(self, config): registry = JitWarmupRegistry(config) with registry.activate(): mono = self.make_mono(config) - assert len(registry) >= 1 - by_width: dict[int, set] = defaultdict(set) - for (_, width), kernel in mono.kernels.items(): - by_width[width] |= set(kernel.warmup_keys()) - assert set(by_width) == set(mono.widths) - assert all(by_width.values()) + keys = mono.kernel.warmup_keys() + assert len(keys) == len(set(keys)) >= 1 + assert len(registry) == len(keys) registry.warmup() def _positions(self, mono) -> torch.Tensor: return torch.zeros(min(mono.widths), dtype=torch.long) - def _decode(self, config, model, mono) -> None: + def _plan(self, config, model, mono): with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - mono.skip_reason(self._positions(mono)) - - def _run_layers(self, mono, model) -> None: - hidden = torch.zeros(min(mono.widths), model.config.hidden_size) - positions = self._positions(mono) - for i, layer in enumerate(model.layers): - hidden = forward_layer(mono, i, layer, positions, hidden) + return mono.plan(self._positions(mono)) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/mono/reference.py index 11c1cb2a2399..3b8d1316804f 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/mono/reference.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The reference MonoKernel: a ``LayerKernel`` over the layer it replaces.""" +"""The reference MonoKernel: a ``ModelKernel`` over the layers it replaces.""" -from dataclasses import dataclass, replace +from dataclasses import replace from types import SimpleNamespace from typing import Any @@ -10,7 +10,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context -from vllm.models.common.mono.spec import MonoSpec +from vllm.models.common.mono.spec import MonoSpec, Step CACHE_DTYPE = torch.bfloat16 CACHE_NDIM = 4 @@ -43,32 +43,48 @@ def __init__(self, layers: int = 4, hidden: int = 8): self.start_layer = 0 self.end_layer = layers - -@dataclass -class LaunchState: - """What ``launch_state`` builds once at attach: views of the model's - weights, never a copy.""" - - weights: dict[int, torch.Tensor] + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden = inputs_embeds + for layer in self.layers: + hidden = layer(positions, hidden) + return hidden class ReferenceKernel: - """One layer's launch at one width: the work of the layer it replaces, off - the weight view its launch state holds.""" - - def __init__(self, vllm_config, state: LaunchState, layer, width: int): - self.cache_layers = (layer.attn_name,) - self.index = layer.index - self.weight = state.weights[layer.index] - self.width = width + """The model's decode in one launch: the work of every layer it replaces, + off views of the model's weights, with the dense layers left to vLLM.""" + + def __init__(self, vllm_config, model: ReferenceModel, widths: tuple[int, ...]): + self.model = model + self.widths = widths + served = [ + model.layers[i] + for i in range(model.start_layer, model.end_layer) + if not model.layers[i].dense + ] + self.weights = {layer.index: layer.weight.data for layer in served} + self.cache_layers = tuple(layer.attn_name for layer in served) self.launched: list[int] = [] - def launch(self, positions: torch.Tensor, hidden: torch.Tensor): - self.launched.append(self.width) - return hidden + self.weight + def run(self, step: Step, input_ids, inputs_embeds, **kwargs): + self.launched.append(step.width) + hidden = inputs_embeds + for layer in self.model.layers: + weight = self.weights.get(layer.index) + if weight is None: + hidden = layer(step.positions, hidden) + else: + hidden = hidden + weight + return hidden def warmup_keys(self) -> list[Any]: - return [(self.index, self.width)] + return [(i, w) for i in self.weights for w in self.widths] def compile(self, key: Any) -> None: pass @@ -92,20 +108,12 @@ def refuse_model(model) -> list[str]: return [] if isinstance(model, ReferenceModel) else ["not a ReferenceModel"] -def skip_layer(layer) -> str | None: - return "dense layer" if layer.dense else None - - -def skip_step(vllm_config, state: LaunchState, width: int) -> str | None: +def skip_step(vllm_config, kernel: ReferenceKernel, width: int) -> str | None: md = get_forward_context().attn_metadata return "prefill" if md.get("prefill") else None -def launch_state(model, vllm_config, layers: tuple[int, ...]) -> LaunchState: - return LaunchState(weights={i: model.layers[i].weight.data for i in layers}) - - -def refuse_caches(vllm_config, state: LaunchState, caches) -> list[str]: +def refuse_caches(vllm_config, kernel: ReferenceKernel, caches) -> list[str]: reasons = [] for name, cache in caches.items(): if cache.dtype is not CACHE_DTYPE: @@ -115,22 +123,18 @@ def refuse_caches(vllm_config, state: LaunchState, caches) -> list[str]: return reasons -def build( - vllm_config, state: LaunchState, layer, widths: tuple[int, ...] -) -> dict[int, ReferenceKernel]: - return {w: ReferenceKernel(vllm_config, state, layer, w) for w in widths} +def build(vllm_config, model, widths: tuple[int, ...]) -> ReferenceKernel: + return ReferenceKernel(vllm_config, model, widths) -REFERENCE = MonoSpec[LaunchState]( +REFERENCE = MonoSpec[ReferenceKernel]( name="reference", supports=frozenset(), widths=widths, refuse_config=refuse_config, refuse_model=refuse_model, refuse_caches=refuse_caches, - skip_layer=skip_layer, skip_step=skip_step, - launch_state=launch_state, build=build, ) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index fbae85c8ae37..769151826530 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -6,10 +6,10 @@ import torch from tests.models.common.mono import reference -from tests.models.common.mono.kit import MonoConformance, forward, forward_layer +from tests.models.common.mono.kit import MonoConformance, forward from vllm.config.compilation import CUDAGraphMode -from vllm.models.common.mono.guards import _select_spec, configure -from vllm.models.common.mono.runtime import attach +from vllm.models.common.mono.guards import _select_spec, select_mono_spec +from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.mono.spec import MonoUnsupported SPECS_PATH = "tests.models.common.mono.reference:MONO_SPECS" @@ -36,11 +36,10 @@ def bad_caches(self, config, model): } -def test_switch_off_means_no_attach(): +def test_switch_off_means_no_spec(): config = reference.make_config() config.kernel_config.monokernel = False - assert configure(config, SPECS_PATH) is None - assert attach(reference.make_model(config), config, None) is None + assert select_mono_spec(config, SPECS_PATH) is None def test_the_first_servable_variant_wins(): @@ -56,36 +55,43 @@ def test_the_first_servable_variant_wins(): assert refusals == {} +def _make_mono(config, model): + mono = MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) + mono.bind_caches() + return mono + + def test_skipped_layers_keep_vllm_path(): config = reference.make_config() model = reference.make_model(config) - mono = attach(model, config, configure(config, SPECS_PATH)) - assert {i for i, _ in mono.kernels} == {1, 2, 3} + mono = _make_mono(config, model) + assert set(mono.kernel.weights) == {1, 2, 3} def test_kernel_layers_follow_the_pipeline_slice(): config = reference.make_config() model = reference.make_model(config) model.start_layer, model.end_layer = 2, 4 - mono = attach(model, config, configure(config, SPECS_PATH)) - assert {i for i, _ in mono.kernels} == {2, 3} + mono = _make_mono(config, model) + assert set(mono.kernel.weights) == {2, 3} -def test_kernels_run_the_layers_they_replace(): +def test_the_kernel_runs_the_layers_it_replaces(): config = reference.make_config() model = reference.make_model(config) - mono = attach(model, config, configure(config, SPECS_PATH)) + mono = _make_mono(config, model) width = min(mono.widths) positions = torch.zeros(width, dtype=torch.long) - hidden = torch.zeros(width, model.config.hidden_size) + embeds = torch.zeros(width, model.config.hidden_size) with forward(config, CUDAGraphMode.NONE, {"prefill": False}): - for i, layer in enumerate(model.layers): - hidden = forward_layer(mono, i, layer, positions, hidden) - assert torch.equal(hidden, torch.full_like(hidden, 0 + 1 + 2 + 3)) - assert [k.launched for k in mono.kernels.values() if k.launched] == [[width]] * 3 + step = mono.plan(positions) + assert step is not None + hidden = mono.run(step, None, embeds) + assert torch.equal(hidden, model(None, positions, None, embeds)) + assert mono.kernel.launched == [width] def test_model_refusal_is_raised(): config = reference.make_config() with pytest.raises(MonoUnsupported, match="not a ReferenceModel"): - attach(torch.nn.Linear(2, 2), config, configure(config, SPECS_PATH)) + MonoRuntime(torch.nn.Linear(2, 2), config, select_mono_spec(config, SPECS_PATH)) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 5cf7f40f20c6..78207f84b14b 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -13,8 +13,8 @@ from tests.models.common.mono.kit import MonoConformance, forward from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -from vllm.models.common.mono.guards import configure -from vllm.models.common.mono.runtime import attach, kernel_for_layer +from vllm.models.common.mono.guards import select_mono_spec +from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.mono.spec import MonoUnsupported from vllm.models.deepseek_v41.amd import mono_spec as ms @@ -162,7 +162,6 @@ def __init__(self, config): self.layers = torch.nn.ModuleList(_Layer(i) for i in range(LAYERS)) self.start_layer = 0 self.end_layer = LAYERS - self.mono = None def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: @@ -264,11 +263,6 @@ def skipped_metadata(self, config, model): def bad_caches(self, config, model): return {"swa": torch.zeros(4, dtype=torch.bfloat16)} - def _run_layers(self, mono, model) -> None: - positions = self._positions(mono) - for i, layer in enumerate(model.layers): - kernel_for_layer(mono, i, positions) - def _rocm_platform() -> ModuleType: """The real ROCm platform module, whose arch probe needs a GPU to answer.""" @@ -290,7 +284,7 @@ def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS) -> None: monkeypatch.setattr(_rocm_platform(), "get_cdna_version", lambda: cdna) -def _attach(monkeypatch, config=None, model=None, **layer_overrides): +def _build(monkeypatch, config=None, model=None, **layer_overrides): _gpu(monkeypatch) kit = TestDsv41Spec() config = config or make_config() @@ -301,6 +295,10 @@ def _attach(monkeypatch, config=None, model=None, **layer_overrides): return kit.make_mono(config, model), model +def _launch(mono, layer_id: int, width: int = 1): + return ms.DSV41Dispatch(mono.kernel.kernels, width)(layer_id) + + def _indexer(): from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer @@ -337,9 +335,8 @@ def test_layers_take_whole_or_ffn_launch(monkeypatch, overrides, ffn_only): index source's compressor and indexer run between them); a layer outside them keeps vLLM's attention and runs its FFN half as one launch, from wo_b's unreduced output.""" - mono, model = _attach(monkeypatch, **overrides()) - kernels = {k.ffn_only for k in mono.kernels.values()} - assert kernels == {ffn_only} + mono, model = _build(monkeypatch, **overrides()) + assert {k.ffn_only for k in mono.kernel.kernels.values()} == {ffn_only} assert all(layer.attn.wo_b.reduce_results is not ffn_only for layer in model.layers) @@ -357,8 +354,8 @@ def test_skipped_layers_keep_vllm_path(monkeypatch, skip): for layer in model.layers: for key, value in skip.items(): setattr(layer, key, value) - mono, _ = _attach(monkeypatch, config=config, model=model) - assert mono is None + mono, _ = _build(monkeypatch, config=config, model=model) + assert not mono.kernel.kernels def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): @@ -366,13 +363,14 @@ def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): model = TestDsv41Spec().make_model(config) for layer in model.layers: layer.ffn.n_routed_experts = 256 - mono, _ = _attach(monkeypatch, config=config, model=model) - assert mono is None + mono, _ = _build(monkeypatch, config=config, model=model) + assert not mono.kernel.kernels def _step(mono, model, metadata, width=1): with forward(mono.vllm_config, CUDAGraphMode.FULL, metadata): - return mono.skip_reason(torch.zeros(width, dtype=torch.long)) + mono.plan(torch.zeros(width, dtype=torch.long)) + return mono.last_reason @pytest.mark.parametrize( @@ -387,22 +385,22 @@ def _step(mono, model, metadata, width=1): ids=["prefill", "no-decode", "tokens-past-rows", "noncausal-window", "no-swa"], ) def test_other_steps_take_vllm_path(monkeypatch, case, why): - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) assert why in _step(mono, model, _metadata(model, **case)) def test_decode_step_takes_the_kernels(monkeypatch): - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) assert _step(mono, model, _metadata(model)) is None def test_profile_run_takes_vllm_path(monkeypatch): - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) assert _step(mono, model, None) == "no attention metadata" def test_rows_past_the_kernels_take_vllm_path(monkeypatch): - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) assert _step(mono, model, _metadata(model), width=ms.MAX_ROWS + 1) == "width" @@ -436,7 +434,7 @@ def test_other_isa_refuses(monkeypatch, cdna): model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cdna=cdna) with pytest.raises(MonoUnsupported, match="CDNA4"): - attach(model, config, configure(config, SPECS_PATH)) + MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) def test_too_few_compute_units_refuses(monkeypatch): @@ -444,7 +442,7 @@ def test_too_few_compute_units_refuses(monkeypatch): model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cus=BLOCKS // 2) with pytest.raises(MonoUnsupported, match="compute units"): - attach(model, config, configure(config, SPECS_PATH)) + MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) def test_other_moe_backend_refuses(monkeypatch): @@ -455,21 +453,21 @@ def test_other_moe_backend_refuses(monkeypatch): experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON _gpu(monkeypatch) with pytest.raises(MonoUnsupported, match="AITER_MXFP4_BF16"): - attach(model, config, configure(config, SPECS_PATH)) + MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) def test_off_without_the_flag(monkeypatch): + _gpu(monkeypatch) config = make_config() config.kernel_config.monokernel = False - mono, _ = _attach(monkeypatch, config=config) - assert mono is None + assert select_mono_spec(config, SPECS_PATH) is None def test_whole_launch_runs_the_layer(monkeypatch): - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for_layer(mono, 1, positions) + kernel = _launch(mono, 1) out = kernel.launch( torch.zeros(1, HIDDEN), positions, @@ -488,20 +486,19 @@ def test_whole_launch_runs_the_layer(monkeypatch): def test_first_layer_takes_the_embedding(monkeypatch): """Layer 0's residual is the embedding broadcast to the four streams, which K1 takes as it is (no owed post to fold).""" - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for_layer(mono, 0, positions) + kernel = _launch(mono, 0) out = kernel.launch(torch.zeros(1, HIDDEN), positions, None, None, None, None) assert out[0] == "whole" and not out[2]["fold"] assert out[1][2].shape == (1, 4, HIDDEN) def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): - mono, model = _attach(monkeypatch, compressor=object()) - positions = torch.zeros(1, dtype=torch.long) + mono, model = _build(monkeypatch, compressor=object()) with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for_layer(mono, 1, positions) + kernel = _launch(mono, 1) assert kernel.ffn_only out = kernel.launch( torch.zeros(1, HIDDEN), @@ -516,10 +513,10 @@ def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): def test_launch_declines_a_layout_it_cannot_read(monkeypatch): """A later pipeline stage's first layer gets a stream without the mixes: ``skip_step`` cannot see that, so the launch declines it.""" - mono, model = _attach(monkeypatch) + mono, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = kernel_for_layer(mono, 1, positions) + kernel = _launch(mono, 1) out = kernel.launch( torch.zeros(1, 4, HIDDEN), positions, None, None, None, None ) diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py index b98bd5e4e26a..a0d65eec7eea 100644 --- a/vllm/models/common/mono/__init__.py +++ b/vllm/models/common/mono/__init__.py @@ -1,14 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MonoKernels (RFC #60904): a decoder layer's decode step in one launch.""" +"""MonoKernels (RFC #60904): a model's decode step in one launch.""" -from vllm.models.common.mono.guards import configure -from vllm.models.common.mono.runtime import attach, kernel_for_layer -from vllm.models.common.mono.spec import MonoSpec +from vllm.models.common.mono.guards import select_mono_spec +from vllm.models.common.mono.runtime import MonoRuntime +from vllm.models.common.mono.spec import MonoSpec, Step __all__ = [ + "MonoRuntime", "MonoSpec", - "attach", - "configure", - "kernel_for_layer", + "Step", + "select_mono_spec", ] diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index 6f55ca74e574..325585d654a0 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -168,11 +168,13 @@ def _set_capture_sizes(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> None: cc.cudagraph_capture_sizes = new -def configure(vllm_config: "VllmConfig", specs_path: str) -> MonoSpec[Any] | None: +def select_mono_spec( + vllm_config: "VllmConfig", specs_path: str +) -> MonoSpec[Any] | None: """Picks the variant, sets the capture sizes, or raises with every reason. The variant is picked here, once: the model holds what this returns and - hands it to ``attach`` after the weights are loaded. A model that prepares + hands it to ``MonoRuntime`` after the weights are loaded. A model that prepares anything of its own at load time reads the returned spec to know which variant it is preparing for. diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/mono/runtime.py index 43c1630d07ef..085096d73717 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/mono/runtime.py @@ -1,24 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""One ``MonoRuntime`` a rank: the per-step decision and the cache re-checks, -behind the hooks a model calls (``attach``, then ``kernel_for_layer`` per layer). +"""One ``MonoRuntime`` a rank: the per-step decision and the cache checks, +behind the two hooks a model's forward calls (``plan``, then ``run``). """ from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, Generic import torch from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import get_forward_context +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel from vllm.models.common.mono.guards import _refuse_model, tp_agree, tp_group from vllm.models.common.mono.spec import ( - LaunchStateT, - LayerKernel, + KernelT, + ModelKernel, MonoSpec, MonoUnsupported, + Step, ) if TYPE_CHECKING: @@ -26,51 +27,49 @@ logger = init_logger(__name__) -KernelT = TypeVar("KernelT", bound=LayerKernel) -"""The kernel type one spec's ``build`` returns. ``LayerKernel`` is only the -lower bound, so a model annotates its runtime with its own kernel class -(``MonoRuntime[MyLayerKernel, MyLaunchState]``) and reads what that class -adds.""" - class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): - """A spec's layer-kernel builds as vLLM JIT warm-up keys.""" + """A spec's kernel builds as vLLM JIT warm-up keys.""" @dataclass(frozen=True) class CompileKey: spec: str build: Any - def __init__(self, spec: str, owners: dict[Any, LayerKernel]): + def __init__(self, spec: str, kernel: ModelKernel): super().__init__() self.spec = spec - self.owners = owners + self.kernel = kernel def get_warmup_keys(self) -> list["MonoWarmup.CompileKey"]: - return [self.CompileKey(self.spec, b) for b in self.owners] + return [self.CompileKey(self.spec, b) for b in self.kernel.warmup_keys()] def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: - self.owners[compile_key.build].compile(compile_key.build) + self.kernel.compile(compile_key.build) self._compiled_cache[compile_key] = True -class MonoRuntime(Generic[KernelT, LaunchStateT]): - """A model's MonoKernels on this rank. - - The runtime reads the spec and builds the layer kernels. It also - registers them for warmup and handles caching, and drivers the - per-step decision whether to run the kernels or the model's layers. +class MonoRuntime(Generic[KernelT]): + """A model's MonoKernel on this rank. - To build the kernels, it first creates the necessary launch state - LaunchStateT. The kernels built by the runtime are of type KernelT - and are stored in the `kernels` attribute. + The runtime reads the spec and builds the kernel, registers it for warm-up, + checks the caches whenever the runner binds them, and decides each step + between the kernel and vLLM's own forward. The kernel it builds is of type + KernelT and is held in the ``kernel`` attribute. """ - launch_state: LaunchStateT - """What ``MonoSpec.launch_state`` built, set whenever a layer takes a - kernel, which is the only case ``attach`` hands the runtime back.""" + def __init__(self, model, vllm_config: "VllmConfig", spec: MonoSpec[KernelT]): + """Builds this rank's kernel. Collective over the TP group. - def __init__(self, spec: MonoSpec[LaunchStateT], model, vllm_config: "VllmConfig"): + Args: + model: The loaded model, in ``process_weights_after_loading``. + vllm_config: The engine's config. + spec: The variant ``select_mono_spec`` picked. + + Raises: + MonoUnsupported: This rank, or another, refused the loaded model. + + """ self.spec = spec reasons = _refuse_model(spec, model) agreed = tp_agree(not reasons, tp_group()[2]) @@ -81,154 +80,66 @@ def __init__(self, spec: MonoSpec[LaunchStateT], model, vllm_config: "VllmConfig raise MonoUnsupported(f"{refused}: another TP rank refused") self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) - self.kernels: dict[tuple[int, int], KernelT] = {} - self._why: str | None = None - self._fc: Any = None - self._checked: tuple | None = None - - indices = self._kernel_layers(model) - if not indices: - self.cache_layers: list[str] = [] - return - self.launch_state = spec.launch_state(model, vllm_config, tuple(indices)) - for i in indices: - built = spec.build( - vllm_config, self.launch_state, model.layers[i], self.widths - ) - missing = set(self.widths) - set(built) - if missing: - raise MonoUnsupported( - f"MonoKernel {spec.name}: layer {i} has no kernel for widths " - f"{sorted(missing)}" - ) - for w in self.widths: - self.kernels[i, w] = cast(KernelT, built[w]) - self.cache_layers = sorted( - {n for k in self.kernels.values() for n in k.cache_layers} - ) - self._register_warmup() - logger.info( - "MonoKernel %s: %d layers, widths %s", - spec.name, - len(indices), - list(self.widths), - ) - - def _kernel_layers(self, model) -> list[int]: - start = getattr(model, "start_layer", 0) - end = getattr(model, "end_layer", len(model.layers)) - out = [] - for i in range(start, end): - why = self.spec.skip_layer(model.layers[i]) - if why is None: - out.append(i) - else: - logger.info_once( - "MonoKernel %s: layer %d on vLLM's path (%s)", - self.spec.name, - i, - why, - ) - return out - - def _register_warmup(self) -> None: - owners: dict[Any, LayerKernel] = {} - for k in self.kernels.values(): - for key in k.warmup_keys(): - owners.setdefault(key, k) - if owners: - MonoWarmup(self.spec.name, owners).register_warmup() - - def _refuse_caches(self) -> None: - """Re-checks the caches whenever they moved (keyed on every tensor's - data_ptr and numel): memory profiling runs real steps on throwaway - caches before the real ones exist.""" + self.last_reason: str | None = None + self._bound = False + self.kernel: KernelT = spec.build(vllm_config, model, self.widths) + if self.kernel.warmup_keys(): + MonoWarmup(spec.name, self.kernel).register_warmup() + logger.info("MonoKernel %s: widths %s", spec.name, list(self.widths)) + + def bind_caches(self) -> None: + """Checks the caches the kernel will read, as the runner binds them. + + Called from ``ModelState.initialize_kv_cache``, which runs once for the + throwaway cache memory profiling allocates and again for the real one. + + Raises: + MonoUnsupported: The caches are not what this spec serves. + + """ sfc = self.vllm_config.compilation_config.static_forward_context - caches = {n: sfc[n].kv_cache for n in self.cache_layers} - key = tuple( - (t.data_ptr(), t.numel()) - for c in caches.values() - for t in (c if isinstance(c, (list, tuple)) else (c,)) - ) - if key == self._checked: - return - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - raise MonoUnsupported( - f"MonoKernel {self.spec.name}: caches moved during CUDA graph capture" - ) - reasons = self.spec.refuse_caches(self.vllm_config, self.launch_state, caches) + caches = {n: sfc[n].kv_cache for n in self.kernel.cache_layers} + reasons = self.spec.refuse_caches(self.vllm_config, self.kernel, caches) if reasons: raise MonoUnsupported( f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) ) - self._checked = key + self._bound = True - def skip_reason(self, positions: torch.Tensor) -> str | None: - """Why this forward runs vLLM's layers, None when it runs on the - kernels. Decided on the first kernel layer to ask and reused by the - rest, so every layer of one forward is answered alike. + def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | None: + """This step's plan, or None to run the model's own forward instead. - Raises: - MonoUnsupported: The caches the kernels would read moved to - something this spec cannot serve. + Steps that run the model's forward: a pipeline stage's inputs, a width + off the spec's widths, a PIECEWISE step (keyed by token count alone, so + it is replayed for mixed batches too), and whatever ``skip_step`` + refuses. ``last_reason`` holds why. + Every predicate is the same on every TP rank: a MonoKernel waits + in-kernel on its peers, so the ranks must take the same path. """ - fc = get_forward_context() - if fc is self._fc: - return self._why - n = positions.size(-1) + width = positions.size(-1) why: str | None - if fc.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: - why = "piecewise" - elif n not in self.widths: + if intermediate_tensors is not None: + why = "pipeline stage input" + elif width not in self.widths: why = "width" + elif not is_forward_context_available(): + why = "no forward context" + elif get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + why = "piecewise" + elif not self._bound: + why = "caches not bound" else: - why = self.spec.skip_step(self.vllm_config, self.launch_state, n) - if why is None: - self._refuse_caches() - self._fc, self._why = fc, why - return why - - -def attach( - model, vllm_config: "VllmConfig", spec: MonoSpec[LaunchStateT] | None -) -> MonoRuntime[Any, LaunchStateT] | None: - """Builds the launch state and this rank's layer kernels. Collective over - the TP group. - - Args: - model: The loaded model, in ``process_weights_after_loading``. - vllm_config: The engine's config. - spec: The variant ``configure`` picked, None with the switch off. - - Returns: - This rank's ``MonoRuntime``, or None with the switch off or when no layer - takes a kernel. Annotate it with the kernel class the spec's ``build`` - returns (which is what ``kernel_for_layer`` then hands back) and the - spec's launch state. - - Raises: - MonoUnsupported: This rank, or another, refused the loaded model. + why = self.spec.skip_step(self.vllm_config, self.kernel, width) + self.last_reason = why + if why is not None: + logger.info_once( + "MonoKernel %s: step on vLLM's path (%s)", self.spec.name, why + ) + return None + return Step(width=width, positions=positions) - """ - if spec is None: - return None - mono: MonoRuntime[Any, LaunchStateT] = MonoRuntime(spec, model, vllm_config) - if not mono.kernels: - logger.info("MonoKernel %s: no layer takes a kernel", spec.name) - return None - return mono - - -def kernel_for_layer( - mono: MonoRuntime[KernelT, Any] | None, i: int, positions -) -> KernelT | None: - """Layer ``i``'s kernel when this step runs on the kernels, else None. - - The first call under a forward context decides for every layer. The model - calls ``launch`` where it suits the layer, and runs the vLLM layer - unchanged when this returns None. - """ - if mono is None or mono.skip_reason(positions) is not None: - return None - return mono.kernels.get((i, positions.size(-1))) + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + """A planned step on the kernel, returning what the model's own forward + would have returned.""" + return self.kernel.run(step, input_ids, inputs_embeds, **kwargs) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index dd331f607c6b..a3dde031e888 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -8,6 +8,8 @@ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar if TYPE_CHECKING: + import torch + from vllm.config import VllmConfig @@ -30,31 +32,37 @@ class Feature(str, Enum): WEIGHT_UPDATES = "weight_updates" -LaunchStateT = TypeVar("LaunchStateT") -"""Whatever one spec's ``launch_state`` builds. The library never looks inside -it, so a spec chooses the type and the callbacks that receive it agree.""" - - class MonoUnsupported(Exception): """The configuration or the loaded model is outside what a spec serves.""" -class LayerKernel(Protocol): - """One decoder layer's launch, at one width, on this rank.""" +@dataclass(frozen=True) +class Step: + """One decode step the kernel takes, as ``MonoRuntime.plan`` decided it.""" + + width: int + """The step's rows, one of the widths the spec serves.""" + + positions: "torch.Tensor" + """The step's positions, as the model's forward was handed them.""" + + +class ModelKernel(Protocol): + """One model's decode on this rank, for every width the spec serves. + + The kernel owns its own layer loop, so the library has no opinion about + layers: a kernel may replace one layer, two, or the whole backbone. + """ cache_layers: tuple[str, ...] """Attention layers whose caches ``MonoSpec.refuse_caches`` checks.""" - def launch(self, *args, **kwargs) -> Any: - """Takes and returns what the vLLM layer it replaces does, or None to - decline this call. - - The decline should only trigger on what ``skip_step`` cannot answer, - such as the layout of the tensors the layer was handed.""" + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + """Takes and returns what the model's own forward does.""" ... def warmup_keys(self) -> list[Any]: - """Hashable build keys of this width's launches, compiled before CUDA + """Hashable build keys of this kernel's launches, compiled before CUDA graph capture by vLLM's JIT warm-up.""" ... @@ -63,8 +71,14 @@ def compile(self, key: Any) -> None: ... +KernelT = TypeVar("KernelT", bound=ModelKernel) +"""What one spec's ``build`` returns. ``ModelKernel`` is only the lower bound, +so a model annotates its runtime with its own kernel class +(``MonoRuntime[MyKernel]``) and reads what that class adds.""" + + @dataclass(frozen=True) -class MonoSpec(Generic[LaunchStateT]): +class MonoSpec(Generic[KernelT]): """Definition of when a monokernel can run, and how to build it. Attributes: @@ -76,22 +90,14 @@ class MonoSpec(Generic[LaunchStateT]): refuse_config: Deployment refusals the engine's config answers, before the weights are loaded. refuse_model: Deployment refusals only the loaded model answers. - refuse_caches: Deployment refusals the KV caches the kernels will read - answer, re-run whenever those caches move. Takes the launch state, - so the caches are checked against what the kernels were built for. - skip_layer: Why one layer takes vLLM's path, None when it takes the - kernel's. Takes only the layer, so what it reads is what every rank - of the same model sees alike. + refuse_caches: Deployment refusals the KV caches the kernel will read + answer, re-run whenever the runner binds caches. skip_step: Why one step takes vLLM's path, None when it takes the kernel's. Takes the step's width; anything else it needs it reads from the forward context, which every rank sees alike. - launch_state: The weights and launch state, built once at attach, off - the loaded model (whose parameters carry the device the kernels - launch on). The only place a weight copy is ever made. - build: One layer's kernels, by width, called once a layer with every - width the spec serves. What a layer prepares once (its weight - views, anything the layer itself must be told) is prepared here, - and the widths share it. Must return a kernel for every width. + build: This rank's kernel, built once off the loaded model (whose + parameters carry the device it launches on) and every width the + spec serves. The only place a weight copy is ever made. """ @@ -100,10 +106,6 @@ class MonoSpec(Generic[LaunchStateT]): widths: Callable[["VllmConfig"], tuple[int, ...]] refuse_config: Callable[["VllmConfig"], list[str]] refuse_model: Callable[[Any], list[str]] - refuse_caches: Callable[["VllmConfig", LaunchStateT, Mapping[str, Any]], list[str]] - skip_layer: Callable[[Any], str | None] - skip_step: Callable[["VllmConfig", LaunchStateT, int], str | None] - launch_state: Callable[[Any, "VllmConfig", tuple[int, ...]], LaunchStateT] - build: Callable[ - ["VllmConfig", LaunchStateT, Any, tuple[int, ...]], Mapping[int, LayerKernel] - ] + refuse_caches: Callable[["VllmConfig", KernelT, Mapping[str, Any]], list[str]] + skip_step: Callable[["VllmConfig", KernelT, int], str | None] + build: Callable[["VllmConfig", Any, tuple[int, ...]], KernelT] diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 9b2e89b91848..185d0e12a9f8 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -50,8 +50,7 @@ make_layers, maybe_prefix, ) -from vllm.models.common.mono import MonoSpec, attach, configure, kernel_for_layer -from vllm.models.common.mono.runtime import MonoRuntime +from vllm.models.common.mono import MonoRuntime, MonoSpec, select_mono_spec from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -63,8 +62,8 @@ ) from vllm.models.deepseek_v41.amd.mono_spec import ( SPECS_PATH, - DSV41LayerKernel, - LaunchState, + DSV41Dispatch, + DSV41Kernel, ) from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention @@ -281,9 +280,9 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - mono: MonoRuntime[DSV41LayerKernel, LaunchState] | None = None, + mono: DSV41Dispatch | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - kernel = kernel_for_layer(mono, self.attn.layer_id, positions) + kernel = mono(self.attn.layer_id) if mono is not None else None if kernel is not None and not kernel.ffn_only: out = kernel.launch( x, @@ -484,10 +483,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps self.vllm_config = vllm_config - self.mono: MonoRuntime[DSV41LayerKernel, LaunchState] | None = None - self.mono_spec: MonoSpec[LaunchState] | None = configure( - vllm_config, SPECS_PATH - ) # Three aux streams: one per non-default input GEMM in # DeepseekV4Attention._run_parallel_input_projections @@ -659,6 +654,7 @@ def forward( intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, lookback_token_ids: torch.Tensor | None = None, + mono: DSV41Dispatch | None = None, ) -> torch.Tensor | IntermediateTensors: if get_pp_group().is_first_rank: if inputs_embeds is not None: @@ -761,6 +757,7 @@ def forward( engram_hashes, engram_mask, full_num_tokens, + mono, ) # Replay collapses inside its own batch. The early loop still sets # ran_layers, so that flag alone must not post again on the full batch. @@ -825,6 +822,7 @@ def _run_layers( engram_hashes: torch.Tensor | None, engram_mask: torch.Tensor | None, full_num_tokens: int, + mono: DSV41Dispatch | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -850,7 +848,7 @@ def _run_layers( residual, engram_hashes, engram_mask, - self.mono, + mono, ) # The aux id of the first replay layer is the post of the last # full-batch layer. Capture it on the trimmed rows instead. @@ -1312,6 +1310,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.config = config + self.vllm_config = vllm_config + self.mono_spec: MonoSpec[DSV41Kernel] | None = select_mono_spec( + vllm_config, SPECS_PATH + ) + self.mono: MonoRuntime[DSV41Kernel] | None = None expert_dtype = getattr(config, "expert_dtype", "fp4") self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype) @@ -1372,9 +1375,9 @@ def compute_logits_local( @staticmethod def get_model_state_cls(): - from ..nvidia.model_state import DeepseekV41ModelState + from .model_state import DeepseekV41ROCMModelState - return DeepseekV41ModelState + return DeepseekV41ROCMModelState @property def token_lookback_depth(self) -> int: @@ -1395,6 +1398,15 @@ def forward( inputs_embeds: torch.Tensor | None = None, lookback_token_ids: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: + if self.mono is not None: + step = self.mono.plan(positions, intermediate_tensors) + if step is not None: + return self.mono.run( + step, + input_ids, + inputs_embeds, + lookback_token_ids=lookback_token_ids, + ) hidden_states = self.model( input_ids, positions, @@ -1419,9 +1431,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() - self.model.mono = attach( - self.model, self.model.vllm_config, self.model.mono_spec - ) + if self.mono_spec is not None: + self.mono = MonoRuntime(self.model, self.vllm_config, self.mono_spec) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v41/amd/model_state.py b/vllm/models/deepseek_v41/amd/model_state.py new file mode 100644 index 000000000000..84d545158c73 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/model_state.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1's model runner state on ROCm: the shared state plus the +MonoKernel's cache checks.""" + +from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.block_table import BlockTables + + +class DeepseekV41ROCMModelState(DeepseekV41ModelState): + """Checks the MonoKernel's caches for every cache the runner binds: the + throwaway one memory profiling allocates, then the real one.""" + + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, block_tables: BlockTables + ) -> None: + super().initialize_kv_cache(kv_cache_config, block_tables) + if self.model.mono is not None: + self.model.mono.bind_caches() diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 0545b23ed005..83abdea66d68 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek-V4.1's MonoKernels on ROCm CDNA4 (``--monokernel``, RFC #60904).""" -from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast import torch @@ -11,7 +10,7 @@ from vllm.config import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.common.mono.spec import Feature, MonoSpec +from vllm.models.common.mono.spec import Feature, MonoSpec, Step from vllm.platforms import current_platform if TYPE_CHECKING: @@ -69,39 +68,19 @@ def _select( ) -@dataclass -class LaunchState: - """What this rank's launches hold across every step: the runner (one a - rank, its peer memory collective over the TP group), the error words, each - kernel layer's weight views, and the metadata keys ``skip_step`` validates - for every kernel layer at once. - - Built once at attach and never moved: the kernels and the captured graphs - hold its addresses. - """ - - runner: Any - embed_layer: int - words: torch.Tensor - weights: dict[int, Any] = field(default_factory=dict) - swa_keys: list[str] = field(default_factory=list) - comp_keys: list[str] = field(default_factory=list) - need_swa: bool = False - - class DSV41LayerKernel: - """The ``LayerKernel`` both launches implement: the layer, its width and - its weight views, with the launch shape left to the subclass.""" + """What both launches share: the layer, its width and its weight views, + with the launch shape left to the subclass.""" ffn_only = False def __init__( - self, launch_state: LaunchState, layer: "DeepseekV4DecoderLayer", width: int + self, mono: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int ): - self.launch_state = launch_state + self.mono = mono self.layer = layer self.width = width - self.weights = launch_state.weights[layer.attn.layer_id] + self.weights = mono.weights[layer.attn.layer_id] self.cache_layers: tuple[str, ...] = () def warmup_keys(self) -> list[Any]: @@ -118,9 +97,9 @@ class WholeLayer(DSV41LayerKernel): """The whole backbone layer as K1 and K2, from the layer's own inputs.""" def __init__( - self, launch_state: LaunchState, layer: "DeepseekV4DecoderLayer", width: int + self, mono: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int ): - super().__init__(launch_state, layer, width) + super().__init__(mono, layer, width) attn = layer.attn self.ratio = attn.compress_ratio self.index = attn.indexer is not None @@ -131,7 +110,7 @@ def __init__( ) # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == launch_state.embed_layer or layer.engram is not None + settled = attn.layer_id == mono.embed_layer or layer.engram is not None self.folds = (True, False) if settled else (True,) def warmup_keys(self) -> list[Any]: @@ -139,8 +118,8 @@ def warmup_keys(self) -> list[Any]: def compile(self, key: Any) -> None: width, ratio, fold, index = key - self.launch_state.runner.scratch(width) - self.launch_state.runner.kernels(width, ratio, fold=fold, index=index) + self.mono.runner.scratch(width) + self.mono.runner.kernels(width, ratio, fold=fold, index=index) def launch( self, @@ -154,7 +133,7 @@ def launch( engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: layer, attn = self.layer, self.layer.attn - runner = self.launch_state.runner + runner = self.mono.runner first = residual is None if first and x.dim() != 2: return None @@ -214,8 +193,8 @@ def warmup_keys(self) -> list[Any]: return [(self.width, "ffn")] def compile(self, key: Any) -> None: - self.launch_state.runner.scratch(key[0]) - self.launch_state.runner.ffn_kernel(key[0]) + self.mono.runner.scratch(key[0]) + self.mono.runner.ffn_kernel(key[0]) def launch( self, @@ -228,7 +207,114 @@ def launch( tensors = (part, residual, post_mix, res_mix, pre_mix) if not all(t.is_contiguous() for t in tensors): return None - return self.launch_state.runner.ffn(self.weights, *tensors) + return self.mono.runner.ffn(self.weights, *tensors) + + +class DSV41Dispatch: + """One step's launches, as the backbone's layers reach them. Built per step + from the step's width, so a layer only ever sees the launch its rows fit.""" + + def __init__(self, kernels: dict[tuple[int, int], DSV41LayerKernel], width: int): + self.kernels = kernels + self.width = width + + def __call__(self, layer_id: int) -> DSV41LayerKernel | None: + return self.kernels.get((layer_id, self.width)) + + +class DSV41Kernel: + """DeepSeek-V4.1's decode on this rank: the runner (one a rank, its peer + memory collective over the TP group), each kernel layer's weight views, and + a launch per (layer, width). + + ``run`` takes the backbone's own forward, which keeps the engram hashes, + the sequence-parallel shard and the final collapse on one code path, and + hands it the step's dispatch for its layers to launch from. + + Built once at load and never moved: the launches and the captured graphs + hold its addresses. + """ + + def __init__(self, vllm_config, model, widths: tuple[int, ...]): + from vllm.models.common.mono.guards import tp_group + from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer + + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + rank, world, cpu_group = tp_group() + device = next(model.parameters()).device + self.model = model + self.runner = DSV41MonoLayer(world, rank, cpu_group, device) + self.embed_layer = getattr(model, "start_layer", 0) + self.words = torch.zeros(world, dtype=torch.int32, device=device) + self.weights: dict[int, Any] = {} + self.swa_keys: list[str] = [] + self.comp_keys: list[str] = [] + self.need_swa = False + self.kernels: dict[tuple[int, int], DSV41LayerKernel] = {} + self._owners: dict[Any, DSV41LayerKernel] = {} + + tp_size = vllm_config.parallel_config.tensor_parallel_size + for layer in self._kernel_layers(model): + self._build_layer(layer, tp_size, widths) + self.cache_layers = tuple( + sorted({n for k in self.kernels.values() for n in k.cache_layers}) + ) + + def _kernel_layers(self, model) -> list["DeepseekV4DecoderLayer"]: + end = getattr(model, "end_layer", len(model.layers)) + layers = [] + for i in range(self.embed_layer, end): + why = _skip_layer(model.layers[i]) + if why is None: + layers.append(model.layers[i]) + else: + logger.info_once( + "DSv4.1 MonoKernel: layer %d on vLLM's path (%s)", i, why + ) + return layers + + def _build_layer(self, layer, tp_size: int, widths: tuple[int, ...]) -> None: + why = _ffn_only(layer) + layer_id = layer.attn.layer_id + self.weights[layer_id] = _weights(layer, why is not None, tp_size) + launch: type[DSV41LayerKernel] = WholeLayer + if why is not None: + logger.info_once( + "DSv4.1 MonoKernel: the FFN launch for layers with %s", why + ) + # A decode step hands wo_b's TP partial to the FFN launch; the layer + # reduces it on the steps that keep vLLM's path. + layer.attn.wo_b.reduce_results = False + launch = FfnSeam + else: + self.need_swa = True + for key, keys in ( + (layer.attn.swa_cache_layer.prefix, self.swa_keys), + (layer.attn.compressed_cache_prefix, self.comp_keys), + ): + if key is not None and key not in keys: + keys.append(key) + for width in widths: + kernel = launch(self, layer, width) + self.kernels[layer_id, width] = kernel + for key in kernel.warmup_keys(): + self._owners.setdefault(key, kernel) + + def warmup_keys(self) -> list[Any]: + return list(self._owners) + + def compile(self, key: Any) -> None: + self._owners[key].compile(key) + + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + return self.model( + input_ids, + step.positions, + None, + inputs_embeds, + mono=DSV41Dispatch(self.kernels, step.width), + **kwargs, + ) def _compute_units() -> int: @@ -327,12 +413,14 @@ def _ffn_only(layer) -> str | None: return None -def _skip_step(vllm_config, launch_state: LaunchState, width: int) -> str | None: +def _skip_step(vllm_config, mono: DSV41Kernel, width: int) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP group sees alike.""" fc = get_forward_context() md = fc.attn_metadata + if not mono.kernels: + return "no layer takes a launch" if not 1 <= width <= MAX_ROWS: return f"{width} rows" if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): @@ -340,9 +428,9 @@ def _skip_step(vllm_config, launch_state: LaunchState, width: int) -> str | None return f"{fc.cudagraph_runtime_mode} step" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches return "no attention metadata" - if any(md.get(k) is None for k in launch_state.comp_keys): + if any(md.get(k) is None for k in mono.comp_keys): return "layer metadata missing" - for key in launch_state.swa_keys: + for key in mono.swa_keys: swa = cast("DeepseekV4ROCMAiterSparseSWAMetadata | None", md.get(key)) if swa is None: return "layer metadata missing" @@ -350,7 +438,7 @@ def _skip_step(vllm_config, launch_state: LaunchState, width: int) -> str | None return "not a decode-only step" if swa.num_decode_tokens > width: return "more decode tokens than rows" - if not launch_state.need_swa: + if not mono.need_swa: continue if swa.decode_swa_width != SWA_WIDTH: return f"SWA width {swa.decode_swa_width}" @@ -363,32 +451,6 @@ def _skip_step(vllm_config, launch_state: LaunchState, width: int) -> str | None return None -def _launch_state(model, vllm_config, layers: tuple[int, ...]) -> LaunchState: - """This rank's runner, collective over the TP group, and the metadata keys - of the layers that take a whole-layer launch.""" - from vllm.models.common.mono.guards import tp_group - from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer - - assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS - rank, world, cpu_group = tp_group() - device = next(model.parameters()).device - state = LaunchState( - runner=DSV41MonoLayer(world, rank, cpu_group, device), - embed_layer=getattr(model, "start_layer", 0), - words=torch.zeros(world, dtype=torch.int32, device=device), - ) - whole = [model.layers[i] for i in layers if _ffn_only(model.layers[i]) is None] - state.need_swa = bool(whole) - for layer in whole: - for key, keys in ( - (layer.attn.swa_cache_layer.prefix, state.swa_keys), - (layer.attn.compressed_cache_prefix, state.comp_keys), - ): - if key is not None and key not in keys: - keys.append(key) - return state - - def _weights(layer, ffn_only: bool, world: int): from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights @@ -443,7 +505,7 @@ def _attn_weights(a, world: int): return attn -def _refuse_caches(vllm_config, launch_state: LaunchState, caches) -> list[str]: +def _refuse_caches(vllm_config, mono: DSV41Kernel, caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as [blocks, block, RECORD] bytes.""" reasons = [] @@ -460,32 +522,18 @@ def _refuse_caches(vllm_config, launch_state: LaunchState, caches) -> list[str]: return reasons -def _build( - vllm_config, launch_state: LaunchState, layer, widths: tuple[int, ...] -) -> dict[int, DSV41LayerKernel]: - world = vllm_config.parallel_config.tensor_parallel_size - why = _ffn_only(layer) - launch_state.weights[layer.attn.layer_id] = _weights(layer, why is not None, world) - kernel: type[DSV41LayerKernel] = WholeLayer - if why is not None: - logger.info_once("DSv4.1 MonoKernel: the FFN launch for layers with %s", why) - # A decode step hands wo_b's TP partial to the FFN launch; the layer - # reduces it on the steps that keep vLLM's path. - layer.attn.wo_b.reduce_results = False - kernel = FfnSeam - return {w: kernel(launch_state, layer, w) for w in widths} +def _build(vllm_config, model, widths: tuple[int, ...]) -> DSV41Kernel: + return DSV41Kernel(vllm_config, model, widths) -MONO_SPEC = MonoSpec[LaunchState]( +MONO_SPEC = MonoSpec[DSV41Kernel]( name="dsv41-cdna4", supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), widths=_widths, refuse_config=_refuse_config, refuse_model=_refuse_model, refuse_caches=_refuse_caches, - skip_layer=_skip_layer, skip_step=_skip_step, - launch_state=_launch_state, build=_build, ) From e7c2d3b7186d2f828821257365274d21f9f28ba8 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 16:30:52 -0500 Subject: [PATCH 33/79] fix: remove lazy import of the spec Signed-off-by: simondanielsson --- tests/models/common/mono/kit.py | 23 +++++++++--------- tests/models/common/mono/test_reference.py | 11 +++++---- tests/models/test_dsv41_mono_decode.py | 13 ++++++----- vllm/models/common/mono/guards.py | 27 ++++++++++------------ vllm/models/deepseek_v41/amd/model.py | 8 +++---- vllm/models/deepseek_v41/amd/mono_spec.py | 2 -- 6 files changed, 40 insertions(+), 44 deletions(-) diff --git a/tests/models/common/mono/kit.py b/tests/models/common/mono/kit.py index a79b342ce59b..cf954db77503 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/mono/kit.py @@ -16,7 +16,6 @@ from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.models.common.mono.guards import ( - _load_specs, _refuse_config, _select_spec, _set_capture_sizes, @@ -58,9 +57,11 @@ class MonoConformance: """What every ``MonoSpec`` must satisfy. Subclass it in the model's test module and fill in the hooks below.""" - specs_path: str - """``"package.module:MONO_SPECS"``, as the model passes it to - ``select_mono_spec``.""" + specs: tuple[Any, ...] + """The model's variants, as it passes them to ``select_mono_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.""" @@ -100,7 +101,7 @@ def make_mono(self, config, model=None, bind: bool = True): picked once by ``select_mono_spec``, then the runtime over the loaded model, then the caches the runner binds.""" model = self.make_model(config) if model is None else model - spec = select_mono_spec(config, self.specs_path) + spec = select_mono_spec(config, self.specs) assert spec is not None mono = MonoRuntime(model, config, spec) if bind: @@ -108,7 +109,7 @@ def make_mono(self, config, model=None, bind: bool = True): return mono def spec(self, config): - spec, refusals = _select_spec(_load_specs(self.specs_path), config) + spec, refusals = _select_spec(self.specs, config) assert spec is not None, refusals return spec @@ -116,7 +117,7 @@ def test_refusals_report_every_reason(self, config): config.lora_config = object() config.speculative_config = object() config.use_v2_model_runner = False - for spec in _load_specs(self.specs_path): + for spec in self.specs: reasons = _refuse_config(spec, config) assert sum("Model Runner V2" in r for r in reasons) == 1, reasons for feature in (Feature.LORA, Feature.SPECULATIVE_DECODING): @@ -126,12 +127,12 @@ def test_refusals_report_every_reason(self, config): def test_refusal_names_every_variant(self, config): config.use_v2_model_runner = False with pytest.raises(MonoUnsupported) as excinfo: - select_mono_spec(config, self.specs_path) - for spec in _load_specs(self.specs_path): + select_mono_spec(config, self.specs) + for spec in self.specs: assert spec.name in str(excinfo.value) def test_spec_module_imports_no_kernel_backend(self): - module = self.specs_path.split(":")[0] + module = self.specs_module source = ( "import sys; sys.path.insert(0, '.'); import importlib;" f"importlib.import_module({module!r});" @@ -231,7 +232,7 @@ def test_build_does_not_move_weights(self, config): def test_build_runs_once_over_every_width(self, config, model): calls: list = [] - spec = _spy(select_mono_spec(config, self.specs_path), "build", calls) + spec = _spy(select_mono_spec(config, self.specs), "build", calls) mono = MonoRuntime(model, config, spec) assert [args[1] for args in calls] == [model] assert set(calls[0][2]) == set(mono.widths) diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/mono/test_reference.py index 769151826530..4244fb6aa3dc 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/mono/test_reference.py @@ -12,11 +12,12 @@ from vllm.models.common.mono.runtime import MonoRuntime from vllm.models.common.mono.spec import MonoUnsupported -SPECS_PATH = "tests.models.common.mono.reference:MONO_SPECS" +SPECS = reference.MONO_SPECS class TestReferenceSpec(MonoConformance): - specs_path = SPECS_PATH + specs = SPECS + specs_module = "tests.models.common.mono.reference" def make_config(self, **overrides): return reference.make_config(**overrides) @@ -39,7 +40,7 @@ def bad_caches(self, config, model): def test_switch_off_means_no_spec(): config = reference.make_config() config.kernel_config.monokernel = False - assert select_mono_spec(config, SPECS_PATH) is None + assert select_mono_spec(config, SPECS) is None def test_the_first_servable_variant_wins(): @@ -56,7 +57,7 @@ def test_the_first_servable_variant_wins(): def _make_mono(config, model): - mono = MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) + mono = MonoRuntime(model, config, select_mono_spec(config, SPECS)) mono.bind_caches() return mono @@ -94,4 +95,4 @@ def test_the_kernel_runs_the_layers_it_replaces(): def test_model_refusal_is_raised(): config = reference.make_config() with pytest.raises(MonoUnsupported, match="not a ReferenceModel"): - MonoRuntime(torch.nn.Linear(2, 2), config, select_mono_spec(config, SPECS_PATH)) + MonoRuntime(torch.nn.Linear(2, 2), config, select_mono_spec(config, SPECS)) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 78207f84b14b..6f0879feff6d 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -18,7 +18,7 @@ from vllm.models.common.mono.spec import MonoUnsupported from vllm.models.deepseek_v41.amd import mono_spec as ms -SPECS_PATH = ms.SPECS_PATH +SPECS = ms.MONO_SPECS LAYERS = 6 HIDDEN = 16 BLOCKS = 256 @@ -234,7 +234,8 @@ def _metadata(model, **overrides) -> dict: class TestDsv41Spec(MonoConformance): - specs_path = SPECS_PATH + specs = SPECS + specs_module = "vllm.models.deepseek_v41.amd.mono_spec" @pytest.fixture(autouse=True) def _stubs(self, monkeypatch): @@ -434,7 +435,7 @@ def test_other_isa_refuses(monkeypatch, cdna): model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cdna=cdna) with pytest.raises(MonoUnsupported, match="CDNA4"): - MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) + MonoRuntime(model, config, select_mono_spec(config, SPECS)) def test_too_few_compute_units_refuses(monkeypatch): @@ -442,7 +443,7 @@ def test_too_few_compute_units_refuses(monkeypatch): model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cus=BLOCKS // 2) with pytest.raises(MonoUnsupported, match="compute units"): - MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) + MonoRuntime(model, config, select_mono_spec(config, SPECS)) def test_other_moe_backend_refuses(monkeypatch): @@ -453,14 +454,14 @@ def test_other_moe_backend_refuses(monkeypatch): experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON _gpu(monkeypatch) with pytest.raises(MonoUnsupported, match="AITER_MXFP4_BF16"): - MonoRuntime(model, config, select_mono_spec(config, SPECS_PATH)) + MonoRuntime(model, config, select_mono_spec(config, SPECS)) def test_off_without_the_flag(monkeypatch): _gpu(monkeypatch) config = make_config() config.kernel_config.monokernel = False - assert select_mono_spec(config, SPECS_PATH) is None + assert select_mono_spec(config, SPECS) is None def test_whole_launch_runs_the_layer(monkeypatch): diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/mono/guards.py index 325585d654a0..5f516522a5d8 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/mono/guards.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Refusals: the variant a config can serve, and TP agreement on every stage.""" -import importlib from collections.abc import Sequence from typing import TYPE_CHECKING, Any @@ -11,7 +10,12 @@ from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger -from vllm.models.common.mono.spec import Feature, MonoSpec, MonoUnsupported +from vllm.models.common.mono.spec import ( + Feature, + KernelT, + MonoSpec, + MonoUnsupported, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -64,15 +68,6 @@ def _enabled(vllm_config: "VllmConfig") -> bool: return bool(vllm_config.kernel_config.monokernel) -def _load_specs(path: str) -> tuple[MonoSpec[Any], ...]: - """``"package.module:NAME"`` -> the spec variants, importing only that - module. ``NAME`` is a ``MonoSpec`` or a tuple of them, in the order they - are preferred.""" - module, name = path.split(":") - specs = getattr(importlib.import_module(module), name) - return (specs,) if isinstance(specs, MonoSpec) else tuple(specs) - - def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: """The deployment's features a spec has to list in ``supports``.""" pc = vllm_config.parallel_config @@ -169,8 +164,9 @@ def _set_capture_sizes(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> None: def select_mono_spec( - vllm_config: "VllmConfig", specs_path: str -) -> MonoSpec[Any] | None: + vllm_config: "VllmConfig", + specs: MonoSpec[KernelT] | Sequence[MonoSpec[KernelT]], +) -> MonoSpec[KernelT] | None: """Picks the variant, sets the capture sizes, or raises with every reason. The variant is picked here, once: the model holds what this returns and @@ -180,7 +176,7 @@ def select_mono_spec( Args: vllm_config: The engine's config, before the weights are loaded. - specs_path: ``"package.module:MONO_SPECS"``. + specs: The model's variants, in the order they are preferred. Returns: The variant this config runs, or None with ``--monokernel`` off. @@ -191,7 +187,8 @@ def select_mono_spec( """ if not _enabled(vllm_config): return None - spec, refusals = _select_spec(_load_specs(specs_path), vllm_config) + ordered = (specs,) if isinstance(specs, MonoSpec) else tuple(specs) + spec, refusals = _select_spec(ordered, vllm_config) if spec is None: raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) _set_capture_sizes(spec, vllm_config) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 6d960167680a..c4972e196bf4 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -49,7 +49,7 @@ make_layers, maybe_prefix, ) -from vllm.models.common.mono import MonoRuntime, MonoSpec, select_mono_spec +from vllm.models.common.mono import MonoRuntime, select_mono_spec from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -60,7 +60,7 @@ DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v41.amd.mono_spec import ( - SPECS_PATH, + MONO_SPECS, DSV41Dispatch, DSV41Kernel, ) @@ -1302,9 +1302,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self.mono_spec: MonoSpec[DSV41Kernel] | None = select_mono_spec( - vllm_config, SPECS_PATH - ) + self.mono_spec = select_mono_spec(vllm_config, MONO_SPECS) self.mono: MonoRuntime[DSV41Kernel] | None = None expert_dtype = getattr(config, "expert_dtype", "fp4") self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 83abdea66d68..34c63b08ebc0 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -19,8 +19,6 @@ logger = init_logger(__name__) -SPECS_PATH = "vllm.models.deepseek_v41.amd.mono_spec:MONO_SPECS" - RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes SWA_WIDTH = 128 # a causal decode token's window slots MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) From 2b683ce29dc40e60262067ae4c494dee25b17bdd Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 16:39:43 -0500 Subject: [PATCH 34/79] fix: move the layer loop to the monokernel itself Signed-off-by: simondanielsson --- tests/models/test_dsv41_mono_decode.py | 61 ++++++- vllm/models/common/mono/spec.py | 8 +- vllm/models/deepseek_v41/amd/model.py | 190 ++++++++++++++++------ vllm/models/deepseek_v41/amd/mono_spec.py | 97 ++++++++--- 4 files changed, 278 insertions(+), 78 deletions(-) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 6f0879feff6d..276bb70683c4 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -15,7 +15,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend from vllm.models.common.mono.guards import select_mono_spec from vllm.models.common.mono.runtime import MonoRuntime -from vllm.models.common.mono.spec import MonoUnsupported +from vllm.models.common.mono.spec import MonoUnsupported, Step from vllm.models.deepseek_v41.amd import mono_spec as ms SPECS = ms.MONO_SPECS @@ -152,6 +152,22 @@ def __init__(self, index: int, **overrides): for key, value in overrides.items(): setattr(self, key, value) self.weight = torch.nn.Parameter(torch.zeros(HIDDEN)) + self.taken: list = [] + + def forward(self, x, positions, input_ids, *rest, mono_ffn=None): + self.taken.append(("layer", self.attn.layer_id, mono_ffn is not None)) + return _stream() + + +def _stream() -> tuple[torch.Tensor, ...]: + """What a decoder layer returns: the stream and its three mixes.""" + return ( + torch.zeros(1, HIDDEN), + torch.zeros(1, 4, HIDDEN), + torch.zeros(1, 4, 1), + torch.zeros(1, 4, 4), + torch.zeros(1, 4), + ) class _Model(torch.nn.Module): @@ -163,6 +179,15 @@ def __init__(self, config): self.start_layer = 0 self.end_layer = LAYERS + def prologue(self, input_ids, positions, intermediate_tensors, embeds, lookback): + return embeds, input_ids, None, None, None, positions.shape[0] + + def after_layer(self, idx, hidden, residual, post_mix, res_mix, tokens, aux): + return None + + def epilogue(self, hidden, *rest): + return hidden + def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: return _bytes(blocks, block * ms.RECORD) @@ -297,7 +322,7 @@ def _build(monkeypatch, config=None, model=None, **layer_overrides): def _launch(mono, layer_id: int, width: int = 1): - return ms.DSV41Dispatch(mono.kernel.kernels, width)(layer_id) + return mono.kernel.kernels[layer_id, width] def _indexer(): @@ -511,6 +536,38 @@ def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): assert out[0] == "ffn" +def test_run_takes_every_layer_the_way_its_launch_allows(monkeypatch): + """One loop over the backbone, from the model's prologue to its epilogue: + a whole-layer launch replaces the layer, a declined one and a layer without + a launch fall through to it, and an FFN-only launch rides into it.""" + config = make_config() + model = TestDsv41Spec().make_model(config) + model.layers[2].attn.compressor = object() + model.layers[4].use_sequence_parallel = True + mono, _ = _build(monkeypatch, config=config, model=model) + for (layer_id, width), kernel in mono.kernel.kernels.items(): + if width != 1 or kernel.ffn_only: + continue + taken = model.layers[layer_id].taken + declines = layer_id == 1 + kernel.launch = lambda *a, _t=taken, _id=layer_id, _d=declines, **kw: ( + _t.append(("launch", _id)) or (None if _d else _stream()) + ) + + step = Step(width=1, positions=torch.zeros(1, dtype=torch.long)) + out = mono.kernel.run(step, None, torch.zeros(1, HIDDEN)) + + assert out.shape == (1, HIDDEN) + assert [layer.taken for layer in model.layers] == [ + [("launch", 0)], + [("launch", 1), ("layer", 1, False)], + [("layer", 2, True)], + [("launch", 3)], + [("layer", 4, False)], + [("launch", 5)], + ] + + def test_launch_declines_a_layout_it_cannot_read(monkeypatch): """A later pipeline stage's first layer gets a stream without the mixes: ``skip_step`` cannot see that, so the launch declines it.""" diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/mono/spec.py index a3dde031e888..e7a8665ba08b 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/mono/spec.py @@ -58,7 +58,13 @@ class ModelKernel(Protocol): """Attention layers whose caches ``MonoSpec.refuse_caches`` checks.""" def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - """Takes and returns what the model's own forward does.""" + """Takes and returns what the model's own forward does. + + A launch replaces a *whole* decoder layer: it returns what that layer's + forward returns, or declines, and the kernel runs the layer on vLLM's + path instead. Covering part of a layer is the model's own problem, and + is discouraged: a new kernel should cover a whole layer or fall back. + """ ... def warmup_keys(self) -> list[Any]: diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index c4972e196bf4..992c25ab960c 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -59,11 +59,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v41.amd.mono_spec import ( - MONO_SPECS, - DSV41Dispatch, - DSV41Kernel, -) +from vllm.models.deepseek_v41.amd.mono_spec import MONO_SPECS, DSV41Kernel from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers @@ -287,22 +283,8 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - mono: DSV41Dispatch | None = None, + mono_ffn: typing.Any | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - kernel = mono(self.attn.layer_id) if mono is not None else None - if kernel is not None and not kernel.ffn_only: - out = kernel.launch( - x, - positions, - residual, - post_mix, - res_mix, - pre_mix, - engram_hashes, - engram_mask, - ) - if out is not None: - return out # Layer 0's attention seam projects the 2-D embedding with the folded # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn. # The fused kernel only takes the latter, so that seam keeps the @@ -402,8 +384,12 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) elif not self.attn.wo_b.reduce_results: - if kernel is not None: - out = kernel.launch(x, residual, post_mix, res_mix, attn_pre) + # A partial-layer MonoKernel: the FFN launch takes wo_b's TP partial + # and the rest of the layer. Whole-layer launches never come through + # here, and a new kernel should cover a whole layer instead of + # reaching into one. + if mono_ffn is not None: + out = mono_ffn.launch(x, residual, post_mix, res_mix, attn_pre) if out is not None: return out x = tensor_model_parallel_all_reduce(x) @@ -661,8 +647,80 @@ def forward( intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, lookback_token_ids: torch.Tensor | None = None, - mono: DSV41Dispatch | None = None, ) -> torch.Tensor | IntermediateTensors: + ( + hidden_states, + input_ids, + pre_mix, + engram_hashes, + engram_mask, + full_num_tokens, + ) = self.prologue( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + lookback_token_ids, + ) + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.start_layer, self.decoder_replay_start), + hidden_states, + positions, + input_ids, + pre_mix, + None, + None, + None, + engram_hashes, + engram_mask, + full_num_tokens, + ) + return self.epilogue( + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_states, + final_aux_recon, + ran_layers, + full_num_tokens, + ) + + def prologue( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None, + lookback_token_ids: torch.Tensor | None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + int, + ]: + """Everything the decoder layers need, before the first of them runs. + + Returns: + The stream the first layer takes, the (possibly sharded) input ids, + the carried pre-mix, the Engram hashes and mask, and the batch's + token count before any sequence-parallel shard. + + """ if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -738,34 +796,35 @@ def forward( hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) - residual, post_mix, res_mix = None, None, None pre_mix: torch.Tensor | None = None if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - ( - hidden_states, - residual, - post_mix, - res_mix, - pre_mix, - aux_hidden_states, - final_aux_recon, - ran_layers, - ) = self._run_layers( - range(self.start_layer, self.decoder_replay_start), + return ( hidden_states, - positions, input_ids, pre_mix, - post_mix, - res_mix, - residual, engram_hashes, engram_mask, full_num_tokens, - mono, ) + + def epilogue( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + residual: torch.Tensor | None, + aux_hidden_states: list[torch.Tensor], + final_aux_recon: torch.Tensor | None, + ran_layers: bool, + full_num_tokens: int, + ) -> torch.Tensor | IntermediateTensors: + """Everything after the last decoder layer: the replay batch, the final + post and collapse, and this rank's return.""" # Replay collapses inside its own batch. The early loop still sets # ran_layers, so that flag alone must not post again on the full batch. collapsed = False @@ -829,7 +888,6 @@ def _run_layers( engram_hashes: torch.Tensor | None, engram_mask: torch.Tensor | None, full_num_tokens: int, - mono: DSV41Dispatch | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -855,19 +913,17 @@ def _run_layers( residual, engram_hashes, engram_mask, - mono, ) - # The aux id of the first replay layer is the post of the last - # full-batch layer. Capture it on the trimmed rows instead. - if idx + 1 in self.aux_hidden_state_layers and not ( - self.decoder_replay_layers is not None - and idx + 1 == self.decoder_replay_start - ): - aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) - aux_hidden_state = aux_recon.mean(dim=1) - if self.use_sequence_parallel: - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] - aux_hidden_states.append(aux_hidden_state) + aux_recon = self.after_layer( + idx, + hidden_states, + residual, + post_mix, + res_mix, + full_num_tokens, + aux_hidden_states, + ) + if aux_recon is not None: final_aux_recon = aux_recon return ( hidden_states, @@ -880,6 +936,34 @@ def _run_layers( ran_layers, ) + def after_layer( + self, + idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + full_num_tokens: int, + aux_hidden_states: list[torch.Tensor], + ) -> torch.Tensor | None: + """This layer's aux hidden state, appended to ``aux_hidden_states``, or + None when layer ``idx + 1`` does not want one. + + The aux id of the first replay layer is the post of the last full-batch + layer; that one is captured on the trimmed rows instead. + """ + if idx + 1 not in self.aux_hidden_state_layers or ( + self.decoder_replay_layers is not None + and idx + 1 == self.decoder_replay_start + ): + return None + aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) + aux_hidden_state = aux_recon.mean(dim=1) + if self.use_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + return aux_recon + def _run_replay_layers( self, hidden_states: torch.Tensor, diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/mono_spec.py index 34c63b08ebc0..80372ebd4b7e 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/mono_spec.py @@ -208,26 +208,14 @@ def launch( return self.mono.runner.ffn(self.weights, *tensors) -class DSV41Dispatch: - """One step's launches, as the backbone's layers reach them. Built per step - from the step's width, so a layer only ever sees the launch its rows fit.""" - - def __init__(self, kernels: dict[tuple[int, int], DSV41LayerKernel], width: int): - self.kernels = kernels - self.width = width - - def __call__(self, layer_id: int) -> DSV41LayerKernel | None: - return self.kernels.get((layer_id, self.width)) - - class DSV41Kernel: """DeepSeek-V4.1's decode on this rank: the runner (one a rank, its peer memory collective over the TP group), each kernel layer's weight views, and a launch per (layer, width). - ``run`` takes the backbone's own forward, which keeps the engram hashes, - the sequence-parallel shard and the final collapse on one code path, and - hands it the step's dispatch for its layers to launch from. + ``run`` owns the backbone's layer loop, between the model's own prologue + and epilogue: those keep the engram hashes, the sequence-parallel shard and + the final collapse on one code path. Built once at load and never moved: the launches and the captured graphs hold its addresses. @@ -243,6 +231,10 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): self.model = model self.runner = DSV41MonoLayer(world, rank, cpu_group, device) self.embed_layer = getattr(model, "start_layer", 0) + # The replay layers run on their own trimmed batch, outside the loop + # ``run`` owns, so no launch is built for them. + end = getattr(model, "end_layer", len(model.layers)) + self.layer_end = getattr(model, "decoder_replay_start", end) self.words = torch.zeros(world, dtype=torch.int32, device=device) self.weights: dict[int, Any] = {} self.swa_keys: list[str] = [] @@ -259,7 +251,7 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): ) def _kernel_layers(self, model) -> list["DeepseekV4DecoderLayer"]: - end = getattr(model, "end_layer", len(model.layers)) + end = self.layer_end layers = [] for i in range(self.embed_layer, end): why = _skip_layer(model.layers[i]) @@ -305,13 +297,74 @@ def compile(self, key: Any) -> None: self._owners[key].compile(key) def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - return self.model( + model, positions = self.model, step.positions + ( + hidden, + input_ids, + pre_mix, + engram_hashes, + engram_mask, + full_num_tokens, + ) = model.prologue( + input_ids, positions, None, inputs_embeds, kwargs.get("lookback_token_ids") + ) + residual = post_mix = res_mix = None + aux_hidden_states: list[torch.Tensor] = [] + final_aux_recon = None + for idx in range(self.embed_layer, self.layer_end): + layer = model.layers[idx] + launch = self.kernels.get((layer.attn.layer_id, step.width)) + ffn = launch if launch is not None and launch.ffn_only else None + whole = launch if launch is not None and ffn is None else None + out = None + if whole is not None: + out = whole.launch( + hidden, + positions, + residual, + post_mix, + res_mix, + pre_mix, + engram_hashes, + engram_mask, + ) + if out is None: + out = layer( + hidden, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + mono_ffn=ffn, + ) + hidden, residual, post_mix, res_mix, pre_mix = out + aux_recon = model.after_layer( + idx, + hidden, + residual, + post_mix, + res_mix, + full_num_tokens, + aux_hidden_states, + ) + if aux_recon is not None: + final_aux_recon = aux_recon + return model.epilogue( + hidden, + positions, input_ids, - step.positions, - None, - inputs_embeds, - mono=DSV41Dispatch(self.kernels, step.width), - **kwargs, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_states, + final_aux_recon, + self.layer_end > self.embed_layer, + full_num_tokens, ) From 6b20e1e68539dc5e14b2f7684733ae84552dd021 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 16:54:30 -0500 Subject: [PATCH 35/79] fix: rename from 'mono' to 'SpecializedDecode' Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 12 +- .../{mono => specialized_decode}/__init__.py | 0 .../{mono => specialized_decode}/kit.py | 123 +++++++++-------- .../{mono => specialized_decode}/reference.py | 10 +- .../test_reference.py | 57 ++++---- tests/models/test_dsv41_mono_layers.py | 6 +- ...de.py => test_dsv41_specialized_decode.py} | 130 +++++++++--------- vllm/models/common/mono/__init__.py | 14 -- .../common/specialized_decode/__init__.py | 15 ++ .../{mono => specialized_decode}/guards.py | 57 ++++---- .../{mono => specialized_decode}/runtime.py | 50 +++---- .../{mono => specialized_decode}/spec.py | 15 +- vllm/models/deepseek_v41/amd/model.py | 34 ++--- vllm/models/deepseek_v41/amd/model_state.py | 8 +- .../amd/mono/common/peer_memory.py | 8 +- .../{mono_spec.py => specialized_decode.py} | 57 ++++---- 16 files changed, 308 insertions(+), 288 deletions(-) rename tests/models/common/{mono => specialized_decode}/__init__.py (100%) rename tests/models/common/{mono => specialized_decode}/kit.py (68%) rename tests/models/common/{mono => specialized_decode}/reference.py (95%) rename tests/models/common/{mono => specialized_decode}/test_reference.py (58%) rename tests/models/{test_dsv41_mono_decode.py => test_dsv41_specialized_decode.py} (81%) delete mode 100644 vllm/models/common/mono/__init__.py create mode 100644 vllm/models/common/specialized_decode/__init__.py rename vllm/models/common/{mono => specialized_decode}/guards.py (77%) rename vllm/models/common/{mono => specialized_decode}/runtime.py (73%) rename vllm/models/common/{mono => specialized_decode}/spec.py (89%) rename vllm/models/deepseek_v41/amd/{mono_spec.py => specialized_decode.py} (92%) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index b678c43d743d..cd417098bfad 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -71,9 +71,9 @@ steps: - tests/models/test_deepseek_v41_replay_batch.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py - - tests/models/common/mono/ + - 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/common/mono + - 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/common/specialized_decode mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -95,11 +95,11 @@ 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 - - vllm/models/common/mono/ - - tests/models/common/mono/ + - tests/models/test_dsv41_specialized_decode.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/common/mono + - 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/common/specialized_decode - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/common/mono/__init__.py b/tests/models/common/specialized_decode/__init__.py similarity index 100% rename from tests/models/common/mono/__init__.py rename to tests/models/common/specialized_decode/__init__.py diff --git a/tests/models/common/mono/kit.py b/tests/models/common/specialized_decode/kit.py similarity index 68% rename from tests/models/common/mono/kit.py rename to tests/models/common/specialized_decode/kit.py index cf954db77503..c39bc0760433 100644 --- a/tests/models/common/mono/kit.py +++ b/tests/models/common/specialized_decode/kit.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The MonoKernel conformance kit (RFC #60904): what a ``MonoSpec`` must satisfy.""" +"""The specialized decode conformance kit (RFC #60904): what a ``DecodeSpec`` +must satisfy.""" import subprocess import sys @@ -15,14 +16,14 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry -from vllm.models.common.mono.guards import ( +from vllm.models.common.specialized_decode.guards import ( _refuse_config, _select_spec, _set_capture_sizes, - select_mono_spec, + select_decode_spec, ) -from vllm.models.common.mono.runtime import MonoRuntime -from vllm.models.common.mono.spec import Feature, MonoUnsupported +from vllm.models.common.specialized_decode.runtime import SpecializedDecode +from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Feature KERNEL_BACKENDS = ("aiter", "flydsl") """Backends a spec module must not pull in: it is imported at config time, on @@ -53,12 +54,12 @@ def recorded(*args): return replace(spec, **{field: recorded}) -class MonoConformance: - """What every ``MonoSpec`` must satisfy. Subclass it in the model's test +class DecodeConformance: + """What every ``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_mono_spec``.""" + """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.""" @@ -93,20 +94,20 @@ def model(self, config): return self.make_model(config) @pytest.fixture - def mono(self, config, model): - return self.make_mono(config, model) + def decode(self, config, model): + return self.make_decode(config, model) - def make_mono(self, config, model=None, bind: bool = True): + def make_decode(self, config, model=None, bind: bool = True): """This rank's runtime, built the way a model builds it: the variant - picked once by ``select_mono_spec``, then the runtime over the loaded model, + picked once by ``select_decode_spec``, then the runtime over the loaded model, then the caches the runner binds.""" model = self.make_model(config) if model is None else model - spec = select_mono_spec(config, self.specs) + spec = select_decode_spec(config, self.specs) assert spec is not None - mono = MonoRuntime(model, config, spec) + decode = SpecializedDecode(model, config, spec) if bind: - mono.bind_caches() - return mono + decode.bind_caches() + return decode def spec(self, config): spec, refusals = _select_spec(self.specs, config) @@ -126,8 +127,8 @@ def test_refusals_report_every_reason(self, config): def test_refusal_names_every_variant(self, config): config.use_v2_model_runner = False - with pytest.raises(MonoUnsupported) as excinfo: - select_mono_spec(config, self.specs) + with pytest.raises(DecodeUnsupported) as excinfo: + select_decode_spec(config, self.specs) for spec in self.specs: assert spec.name in str(excinfo.value) @@ -157,98 +158,98 @@ def test_capture_sizes_are_the_widths(self, config): sizes = config.compilation_config.cudagraph_capture_sizes assert [s for s in sizes if s <= max(widths)] == sorted(widths) - def test_decided_once_a_step(self, config, model, mono): + def test_decided_once_a_step(self, config, model, decode): asked: list = [] - mono.spec = _spy(mono.spec, "skip_step", asked) - step = self._plan(config, model, mono) - assert step is not None and step.width == min(mono.widths) + decode.spec = _spy(decode.spec, "skip_step", asked) + step = self._plan(config, model, decode) + assert step is not None and step.width == min(decode.widths) assert len(asked) == 1 def test_every_rank_decides_alike(self, config, model): decisions = [] for _ in range(2): - other = self.make_mono(config) + other = self.make_decode(config) decisions.append(self._plan(config, model, other)) assert decisions[0] is not None and decisions[1] is not None - def test_skipped_step_runs_vllm_layers(self, config, model, mono): + def test_skipped_step_runs_vllm_layers(self, config, model, decode): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): - assert mono.plan(self._positions(mono)) is None - assert mono.last_reason is not None + assert decode.plan(self._positions(decode)) is None + assert decode.last_reason is not None - def test_pipeline_stage_input_runs_vllm_layers(self, config, model, mono): + def test_pipeline_stage_input_runs_vllm_layers(self, config, model, decode): with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - assert mono.plan(self._positions(mono), object()) is None - assert mono.last_reason == "pipeline stage input" + assert decode.plan(self._positions(decode), object()) is None + assert decode.last_reason == "pipeline stage input" - def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, mono): - off = torch.zeros(max(mono.widths) + 1, dtype=torch.long) + def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, decode): + off = torch.zeros(max(decode.widths) + 1, dtype=torch.long) with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - assert mono.plan(off) is None - assert mono.last_reason == "width" + assert decode.plan(off) is None + assert decode.last_reason == "width" - def test_piecewise_runs_vllm_layers(self, config, model, mono): + def test_piecewise_runs_vllm_layers(self, config, model, decode): meta = self.decode_metadata(config, model) with forward(config, CUDAGraphMode.PIECEWISE, meta): - assert mono.plan(self._positions(mono)) is None - assert mono.last_reason == "piecewise" + assert decode.plan(self._positions(decode)) is None + assert decode.last_reason == "piecewise" def test_unbound_caches_run_vllm_layers(self, config, model): - mono = self.make_mono(config, model, bind=False) - assert self._plan(config, model, mono) is None - assert mono.last_reason == "caches not bound" + decode = self.make_decode(config, model, bind=False) + assert self._plan(config, model, decode) is None + assert decode.last_reason == "caches not bound" - def test_caches_checked_on_every_bind(self, config, model, mono): + def test_caches_checked_on_every_bind(self, config, model, decode): checks: list = [] - mono.spec = _spy(mono.spec, "refuse_caches", checks) - mono.bind_caches() + decode.spec = _spy(decode.spec, "refuse_caches", checks) + decode.bind_caches() assert len(checks) == 1 sfc = config.compilation_config.static_forward_context - for name in mono.kernel.cache_layers: + for name in decode.kernel.cache_layers: sfc[name].kv_cache = sfc[name].kv_cache.clone() - mono.bind_caches() + decode.bind_caches() assert len(checks) == 2 - def test_refuse_caches_reports_every_reason(self, config, model, mono): - reasons = mono.spec.refuse_caches( - mono.vllm_config, mono.kernel, self.bad_caches(config, model) + 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_mismatched_caches_refuse_the_bind(self, config, model, mono): + def test_mismatched_caches_refuse_the_bind(self, config, model, decode): sfc = config.compilation_config.static_forward_context bad = self.bad_caches(config, model) - for name, cache in zip(mono.kernel.cache_layers, bad.values(), strict=False): + for name, cache in zip(decode.kernel.cache_layers, bad.values(), strict=False): sfc[name].kv_cache = cache - with pytest.raises(MonoUnsupported): - mono.bind_caches() + with pytest.raises(DecodeUnsupported): + decode.bind_caches() 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_mono(config, model) + self.make_decode(config, model) after = {n: p.data_ptr() for n, p in model.named_parameters()} assert before == after def test_build_runs_once_over_every_width(self, config, model): calls: list = [] - spec = _spy(select_mono_spec(config, self.specs), "build", calls) - mono = MonoRuntime(model, config, spec) + spec = _spy(select_decode_spec(config, self.specs), "build", calls) + decode = SpecializedDecode(model, config, spec) assert [args[1] for args in calls] == [model] - assert set(calls[0][2]) == set(mono.widths) + assert set(calls[0][2]) == set(decode.widths) def test_warmup_registered_for_every_key(self, config): registry = JitWarmupRegistry(config) with registry.activate(): - mono = self.make_mono(config) - keys = mono.kernel.warmup_keys() + decode = self.make_decode(config) + keys = decode.kernel.warmup_keys() assert len(keys) == len(set(keys)) >= 1 assert len(registry) == len(keys) registry.warmup() - def _positions(self, mono) -> torch.Tensor: - return torch.zeros(min(mono.widths), dtype=torch.long) + def _positions(self, decode) -> torch.Tensor: + return torch.zeros(min(decode.widths), dtype=torch.long) - def _plan(self, config, model, mono): + def _plan(self, config, model, decode): with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - return mono.plan(self._positions(mono)) + return decode.plan(self._positions(decode)) diff --git a/tests/models/common/mono/reference.py b/tests/models/common/specialized_decode/reference.py similarity index 95% rename from tests/models/common/mono/reference.py rename to tests/models/common/specialized_decode/reference.py index 3b8d1316804f..23d5eb87b76d 100644 --- a/tests/models/common/mono/reference.py +++ b/tests/models/common/specialized_decode/reference.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The reference MonoKernel: a ``ModelKernel`` over the layers it replaces.""" +"""The reference kernel: a ``ModelKernel`` over the layers it replaces.""" from dataclasses import replace from types import SimpleNamespace @@ -10,7 +10,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import get_forward_context -from vllm.models.common.mono.spec import MonoSpec, Step +from vllm.models.common.specialized_decode.spec import DecodeSpec, Step CACHE_DTYPE = torch.bfloat16 CACHE_NDIM = 4 @@ -32,7 +32,7 @@ def forward(self, positions: torch.Tensor, hidden: torch.Tensor): class ReferenceModel(torch.nn.Module): - """A model shaped like the layer loop a MonoKernel hooks into.""" + """A model shaped like the layer loop a kernel hooks into.""" def __init__(self, layers: int = 4, hidden: int = 8): super().__init__() @@ -127,7 +127,7 @@ def build(vllm_config, model, widths: tuple[int, ...]) -> ReferenceKernel: return ReferenceKernel(vllm_config, model, widths) -REFERENCE = MonoSpec[ReferenceKernel]( +REFERENCE = DecodeSpec[ReferenceKernel]( name="reference", supports=frozenset(), widths=widths, @@ -147,7 +147,7 @@ def build(vllm_config, model, widths: tuple[int, ...]) -> ReferenceKernel: ), ) -MONO_SPECS = (REFERENCE_WIDE, REFERENCE) +DECODE_SPECS = (REFERENCE_WIDE, REFERENCE) def make_config(**overrides) -> Any: diff --git a/tests/models/common/mono/test_reference.py b/tests/models/common/specialized_decode/test_reference.py similarity index 58% rename from tests/models/common/mono/test_reference.py rename to tests/models/common/specialized_decode/test_reference.py index 4244fb6aa3dc..e05a1d7a27b9 100644 --- a/tests/models/common/mono/test_reference.py +++ b/tests/models/common/specialized_decode/test_reference.py @@ -5,19 +5,22 @@ import pytest import torch -from tests.models.common.mono import reference -from tests.models.common.mono.kit import MonoConformance, forward +from tests.models.common.specialized_decode import reference +from tests.models.common.specialized_decode.kit import DecodeConformance, forward from vllm.config.compilation import CUDAGraphMode -from vllm.models.common.mono.guards import _select_spec, select_mono_spec -from vllm.models.common.mono.runtime import MonoRuntime -from vllm.models.common.mono.spec import MonoUnsupported +from vllm.models.common.specialized_decode.guards import ( + _select_spec, + select_decode_spec, +) +from vllm.models.common.specialized_decode.runtime import SpecializedDecode +from vllm.models.common.specialized_decode.spec import DecodeUnsupported -SPECS = reference.MONO_SPECS +SPECS = reference.DECODE_SPECS -class TestReferenceSpec(MonoConformance): +class TestReferenceSpec(DecodeConformance): specs = SPECS - specs_module = "tests.models.common.mono.reference" + specs_module = "tests.models.common.specialized_decode.reference" def make_config(self, **overrides): return reference.make_config(**overrides) @@ -40,59 +43,61 @@ def bad_caches(self, config, model): def test_switch_off_means_no_spec(): config = reference.make_config() config.kernel_config.monokernel = False - assert select_mono_spec(config, SPECS) is None + assert select_decode_spec(config, SPECS) is None def test_the_first_servable_variant_wins(): config = reference.make_config() config.scheduler_config.max_num_seqs = 4 - spec, refusals = _select_spec(reference.MONO_SPECS, config) + spec, refusals = _select_spec(reference.DECODE_SPECS, config) assert spec is reference.REFERENCE assert refusals["reference-wide"] == ["needs max_num_seqs >= 8"] config.scheduler_config.max_num_seqs = 8 - spec, refusals = _select_spec(reference.MONO_SPECS, config) + spec, refusals = _select_spec(reference.DECODE_SPECS, config) assert spec is reference.REFERENCE_WIDE assert refusals == {} -def _make_mono(config, model): - mono = MonoRuntime(model, config, select_mono_spec(config, SPECS)) - mono.bind_caches() - return mono +def _make_decode(config, model): + decode = SpecializedDecode(model, config, select_decode_spec(config, SPECS)) + decode.bind_caches() + return decode def test_skipped_layers_keep_vllm_path(): config = reference.make_config() model = reference.make_model(config) - mono = _make_mono(config, model) - assert set(mono.kernel.weights) == {1, 2, 3} + decode = _make_decode(config, model) + assert set(decode.kernel.weights) == {1, 2, 3} def test_kernel_layers_follow_the_pipeline_slice(): config = reference.make_config() model = reference.make_model(config) model.start_layer, model.end_layer = 2, 4 - mono = _make_mono(config, model) - assert set(mono.kernel.weights) == {2, 3} + decode = _make_decode(config, model) + assert set(decode.kernel.weights) == {2, 3} def test_the_kernel_runs_the_layers_it_replaces(): config = reference.make_config() model = reference.make_model(config) - mono = _make_mono(config, model) - width = min(mono.widths) + decode = _make_decode(config, model) + width = min(decode.widths) positions = torch.zeros(width, dtype=torch.long) embeds = torch.zeros(width, model.config.hidden_size) with forward(config, CUDAGraphMode.NONE, {"prefill": False}): - step = mono.plan(positions) + step = decode.plan(positions) assert step is not None - hidden = mono.run(step, None, embeds) + hidden = decode.run(step, None, embeds) assert torch.equal(hidden, model(None, positions, None, embeds)) - assert mono.kernel.launched == [width] + assert decode.kernel.launched == [width] def test_model_refusal_is_raised(): config = reference.make_config() - with pytest.raises(MonoUnsupported, match="not a ReferenceModel"): - MonoRuntime(torch.nn.Linear(2, 2), config, select_mono_spec(config, SPECS)) + with pytest.raises(DecodeUnsupported, match="not a ReferenceModel"): + SpecializedDecode( + torch.nn.Linear(2, 2), config, select_decode_spec(config, SPECS) + ) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index c3832d85f85e..bb1e573ab28c 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -193,8 +193,8 @@ def save_capture(self, path: str) -> None: def mono_layers(self) -> dict: model = self.model_runner.model # type: ignore[attr-defined] - mono = getattr(model.model, "mono", None) - kernels = {} if mono is None else mono.kernels + decode = getattr(model, "specialized_decode", None) + kernels = {} if decode is None else decode.kernel.kernels layers = { i: next( ( @@ -206,7 +206,7 @@ def mono_layers(self) -> dict: ) for i in self._decoder_layers() } - runner = None if mono is None else mono.launch_state.runner + runner = None if decode is None else decode.kernel.runner return { "layers": layers, "epoch": -1 if runner is None else int(runner.epoch[0]), diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_specialized_decode.py similarity index 81% rename from tests/models/test_dsv41_mono_decode.py rename to tests/models/test_dsv41_specialized_decode.py index 276bb70683c4..eab098b12e71 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's MonoSpec against the conformance kit, plus its own choices.""" +"""DeepSeek-V4.1's DecodeSpec against the conformance kit, plus its own choices.""" import importlib import sys @@ -10,15 +10,15 @@ import pytest import torch -from tests.models.common.mono.kit import MonoConformance, forward +from tests.models.common.specialized_decode.kit import DecodeConformance, forward from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -from vllm.models.common.mono.guards import select_mono_spec -from vllm.models.common.mono.runtime import MonoRuntime -from vllm.models.common.mono.spec import MonoUnsupported, Step -from vllm.models.deepseek_v41.amd import mono_spec as ms +from vllm.models.common.specialized_decode.guards import select_decode_spec +from vllm.models.common.specialized_decode.runtime import SpecializedDecode +from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Step +from vllm.models.deepseek_v41.amd import specialized_decode as sd -SPECS = ms.MONO_SPECS +SPECS = sd.DECODE_SPECS LAYERS = 6 HIDDEN = 16 BLOCKS = 256 @@ -154,8 +154,8 @@ def __init__(self, index: int, **overrides): self.weight = torch.nn.Parameter(torch.zeros(HIDDEN)) self.taken: list = [] - def forward(self, x, positions, input_ids, *rest, mono_ffn=None): - self.taken.append(("layer", self.attn.layer_id, mono_ffn is not None)) + def forward(self, x, positions, input_ids, *rest, decode_ffn=None): + self.taken.append(("layer", self.attn.layer_id, decode_ffn is not None)) return _stream() @@ -190,7 +190,7 @@ def epilogue(self, hidden, *rest): def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: - return _bytes(blocks, block * ms.RECORD) + return _bytes(blocks, block * sd.RECORD) def make_config(**overrides) -> Any: @@ -219,7 +219,7 @@ def make_config(**overrides) -> Any: aux_output_config=SimpleNamespace(enable_return_routed_experts=False), compilation_config=SimpleNamespace( cudagraph_mode=CUDAGraphMode.FULL, - cudagraph_capture_sizes=list(range(1, ms.MAX_ROWS + 1)), + cudagraph_capture_sizes=list(range(1, sd.MAX_ROWS + 1)), static_forward_context={}, ), lora_config=None, @@ -238,8 +238,8 @@ def _swa(**overrides) -> Any: num_prefills=0, num_decodes=1, num_decode_tokens=1, - decode_swa_width=ms.SWA_WIDTH, - decode_swa_indices=torch.zeros(8, 1, ms.SWA_WIDTH, dtype=torch.int32), + decode_swa_width=sd.SWA_WIDTH, + decode_swa_indices=torch.zeros(8, 1, sd.SWA_WIDTH, dtype=torch.int32), decode_swa_lens=torch.zeros(8, dtype=torch.int32), token_to_req_indices=torch.zeros(8, dtype=torch.int32), slot_mapping=torch.zeros(8, dtype=torch.int64), @@ -258,9 +258,9 @@ def _metadata(model, **overrides) -> dict: return meta -class TestDsv41Spec(MonoConformance): +class TestDsv41Spec(DecodeConformance): specs = SPECS - specs_module = "vllm.models.deepseek_v41.amd.mono_spec" + specs_module = "vllm.models.deepseek_v41.amd.specialized_decode" @pytest.fixture(autouse=True) def _stubs(self, monkeypatch): @@ -306,7 +306,7 @@ def _rocm_platform() -> ModuleType: def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS) -> None: _stub_modules(monkeypatch) - monkeypatch.setattr(ms, "_compute_units", lambda: cus) + monkeypatch.setattr(sd, "_compute_units", lambda: cus) monkeypatch.setattr(_rocm_platform(), "get_cdna_version", lambda: cdna) @@ -318,11 +318,11 @@ def _build(monkeypatch, config=None, model=None, **layer_overrides): for layer in model.layers: for key, value in layer_overrides.items(): setattr(layer.attn, key, value) - return kit.make_mono(config, model), model + return kit.make_decode(config, model), model -def _launch(mono, layer_id: int, width: int = 1): - return mono.kernel.kernels[layer_id, width] +def _launch(decode, layer_id: int, width: int = 1): + return decode.kernel.kernels[layer_id, width] def _indexer(): @@ -361,8 +361,8 @@ def test_layers_take_whole_or_ffn_launch(monkeypatch, overrides, ffn_only): index source's compressor and indexer run between them); a layer outside them keeps vLLM's attention and runs its FFN half as one launch, from wo_b's unreduced output.""" - mono, model = _build(monkeypatch, **overrides()) - assert {k.ffn_only for k in mono.kernel.kernels.values()} == {ffn_only} + decode, model = _build(monkeypatch, **overrides()) + assert {k.ffn_only for k in decode.kernel.kernels.values()} == {ffn_only} assert all(layer.attn.wo_b.reduce_results is not ffn_only for layer in model.layers) @@ -380,8 +380,8 @@ def test_skipped_layers_keep_vllm_path(monkeypatch, skip): for layer in model.layers: for key, value in skip.items(): setattr(layer, key, value) - mono, _ = _build(monkeypatch, config=config, model=model) - assert not mono.kernel.kernels + decode, _ = _build(monkeypatch, config=config, model=model) + assert not decode.kernel.kernels def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): @@ -389,14 +389,14 @@ def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): model = TestDsv41Spec().make_model(config) for layer in model.layers: layer.ffn.n_routed_experts = 256 - mono, _ = _build(monkeypatch, config=config, model=model) - assert not mono.kernel.kernels + decode, _ = _build(monkeypatch, config=config, model=model) + assert not decode.kernel.kernels -def _step(mono, model, metadata, width=1): - with forward(mono.vllm_config, CUDAGraphMode.FULL, metadata): - mono.plan(torch.zeros(width, dtype=torch.long)) - return mono.last_reason +def _step(decode, model, metadata, width=1): + with forward(decode.vllm_config, CUDAGraphMode.FULL, metadata): + decode.plan(torch.zeros(width, dtype=torch.long)) + return decode.last_reason @pytest.mark.parametrize( @@ -411,47 +411,47 @@ def _step(mono, model, metadata, width=1): ids=["prefill", "no-decode", "tokens-past-rows", "noncausal-window", "no-swa"], ) def test_other_steps_take_vllm_path(monkeypatch, case, why): - mono, model = _build(monkeypatch) - assert why in _step(mono, model, _metadata(model, **case)) + decode, model = _build(monkeypatch) + assert why in _step(decode, model, _metadata(model, **case)) def test_decode_step_takes_the_kernels(monkeypatch): - mono, model = _build(monkeypatch) - assert _step(mono, model, _metadata(model)) is None + decode, model = _build(monkeypatch) + assert _step(decode, model, _metadata(model)) is None def test_profile_run_takes_vllm_path(monkeypatch): - mono, model = _build(monkeypatch) - assert _step(mono, model, None) == "no attention metadata" + decode, model = _build(monkeypatch) + assert _step(decode, model, None) == "no attention metadata" def test_rows_past_the_kernels_take_vllm_path(monkeypatch): - mono, model = _build(monkeypatch) - assert _step(mono, model, _metadata(model), width=ms.MAX_ROWS + 1) == "width" + decode, model = _build(monkeypatch) + assert _step(decode, model, _metadata(model), width=sd.MAX_ROWS + 1) == "width" def test_widths_follow_the_capture_sizes(monkeypatch): config = make_config() config.compilation_config.cudagraph_capture_sizes = [1, 2, 8, 64, 512] - assert ms._widths(config) == (1, 2, 8) + assert sd._widths(config) == (1, 2, 8) def test_widths_default_to_every_row(monkeypatch): config = make_config() config.compilation_config.cudagraph_capture_sizes = [512] - assert ms._widths(config) == tuple(range(1, ms.MAX_ROWS + 1)) + assert sd._widths(config) == tuple(range(1, sd.MAX_ROWS + 1)) @pytest.mark.parametrize("tp", [1, 3, 8]) def test_tensor_parallel_outside_the_kernels_refuses(tp): config = make_config() config.parallel_config.tensor_parallel_size = tp - assert any("tensor parallel" in r for r in ms._refuse_config(config)) + assert any("tensor parallel" in r for r in sd._refuse_config(config)) def test_a4w4_moe_refuses(monkeypatch): monkeypatch.setattr("vllm.envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", True) - assert any("A8W4" in r for r in ms._refuse_config(make_config())) + assert any("A8W4" in r for r in sd._refuse_config(make_config())) @pytest.mark.parametrize("cdna", [3, 5]) @@ -459,16 +459,16 @@ def test_other_isa_refuses(monkeypatch, cdna): config = make_config() model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cdna=cdna) - with pytest.raises(MonoUnsupported, match="CDNA4"): - MonoRuntime(model, config, select_mono_spec(config, SPECS)) + with pytest.raises(DecodeUnsupported, match="CDNA4"): + SpecializedDecode(model, config, select_decode_spec(config, SPECS)) def test_too_few_compute_units_refuses(monkeypatch): config = make_config() model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cus=BLOCKS // 2) - with pytest.raises(MonoUnsupported, match="compute units"): - MonoRuntime(model, config, select_mono_spec(config, SPECS)) + with pytest.raises(DecodeUnsupported, match="compute units"): + SpecializedDecode(model, config, select_decode_spec(config, SPECS)) def test_other_moe_backend_refuses(monkeypatch): @@ -478,22 +478,22 @@ def test_other_moe_backend_refuses(monkeypatch): experts = layer.ffn.experts.routed_experts experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON _gpu(monkeypatch) - with pytest.raises(MonoUnsupported, match="AITER_MXFP4_BF16"): - MonoRuntime(model, config, select_mono_spec(config, SPECS)) + with pytest.raises(DecodeUnsupported, match="AITER_MXFP4_BF16"): + SpecializedDecode(model, config, select_decode_spec(config, SPECS)) def test_off_without_the_flag(monkeypatch): _gpu(monkeypatch) config = make_config() config.kernel_config.monokernel = False - assert select_mono_spec(config, SPECS) is None + assert select_decode_spec(config, SPECS) is None def test_whole_launch_runs_the_layer(monkeypatch): - mono, model = _build(monkeypatch) + decode, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(mono, 1) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = _launch(decode, 1) out = kernel.launch( torch.zeros(1, HIDDEN), positions, @@ -503,28 +503,28 @@ def test_whole_launch_runs_the_layer(monkeypatch): torch.zeros(1, 4), ) assert out[0] == "whole" - records = [a for a in out[1] if torch.is_tensor(a) and a.shape[-1:] == (ms.RECORD,)] + records = [a for a in out[1] if torch.is_tensor(a) and a.shape[-1:] == (sd.RECORD,)] assert len(records) == 1 - assert out[2]["comp_cache"].shape[-1] == ms.RECORD + assert out[2]["comp_cache"].shape[-1] == sd.RECORD assert out[2]["fold"] def test_first_layer_takes_the_embedding(monkeypatch): """Layer 0's residual is the embedding broadcast to the four streams, which K1 takes as it is (no owed post to fold).""" - mono, model = _build(monkeypatch) + decode, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(mono, 0) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = _launch(decode, 0) out = kernel.launch(torch.zeros(1, HIDDEN), positions, None, None, None, None) assert out[0] == "whole" and not out[2]["fold"] assert out[1][2].shape == (1, 4, HIDDEN) def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): - mono, model = _build(monkeypatch, compressor=object()) - with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(mono, 1) + decode, model = _build(monkeypatch, compressor=object()) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = _launch(decode, 1) assert kernel.ffn_only out = kernel.launch( torch.zeros(1, HIDDEN), @@ -544,8 +544,8 @@ def test_run_takes_every_layer_the_way_its_launch_allows(monkeypatch): model = TestDsv41Spec().make_model(config) model.layers[2].attn.compressor = object() model.layers[4].use_sequence_parallel = True - mono, _ = _build(monkeypatch, config=config, model=model) - for (layer_id, width), kernel in mono.kernel.kernels.items(): + decode, _ = _build(monkeypatch, config=config, model=model) + for (layer_id, width), kernel in decode.kernel.kernels.items(): if width != 1 or kernel.ffn_only: continue taken = model.layers[layer_id].taken @@ -555,7 +555,7 @@ def test_run_takes_every_layer_the_way_its_launch_allows(monkeypatch): ) step = Step(width=1, positions=torch.zeros(1, dtype=torch.long)) - out = mono.kernel.run(step, None, torch.zeros(1, HIDDEN)) + out = decode.kernel.run(step, None, torch.zeros(1, HIDDEN)) assert out.shape == (1, HIDDEN) assert [layer.taken for layer in model.layers] == [ @@ -571,10 +571,10 @@ def test_run_takes_every_layer_the_way_its_launch_allows(monkeypatch): def test_launch_declines_a_layout_it_cannot_read(monkeypatch): """A later pipeline stage's first layer gets a stream without the mixes: ``skip_step`` cannot see that, so the launch declines it.""" - mono, model = _build(monkeypatch) + decode, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) - with forward(mono.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(mono, 1) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + kernel = _launch(decode, 1) out = kernel.launch( torch.zeros(1, 4, HIDDEN), positions, None, None, None, None ) diff --git a/vllm/models/common/mono/__init__.py b/vllm/models/common/mono/__init__.py deleted file mode 100644 index a0d65eec7eea..000000000000 --- a/vllm/models/common/mono/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MonoKernels (RFC #60904): a model's decode step in one launch.""" - -from vllm.models.common.mono.guards import select_mono_spec -from vllm.models.common.mono.runtime import MonoRuntime -from vllm.models.common.mono.spec import MonoSpec, Step - -__all__ = [ - "MonoRuntime", - "MonoSpec", - "Step", - "select_mono_spec", -] diff --git a/vllm/models/common/specialized_decode/__init__.py b/vllm/models/common/specialized_decode/__init__.py new file mode 100644 index 000000000000..a8d333824080 --- /dev/null +++ b/vllm/models/common/specialized_decode/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Specialized decode (RFC #60904, ``--monokernel``): a model's decode step +in one launch.""" + +from vllm.models.common.specialized_decode.guards import select_decode_spec +from vllm.models.common.specialized_decode.runtime import SpecializedDecode +from vllm.models.common.specialized_decode.spec import DecodeSpec, Step + +__all__ = [ + "DecodeSpec", + "SpecializedDecode", + "Step", + "select_decode_spec", +] diff --git a/vllm/models/common/mono/guards.py b/vllm/models/common/specialized_decode/guards.py similarity index 77% rename from vllm/models/common/mono/guards.py rename to vllm/models/common/specialized_decode/guards.py index 5f516522a5d8..6b20d4e54acc 100644 --- a/vllm/models/common/mono/guards.py +++ b/vllm/models/common/specialized_decode/guards.py @@ -10,11 +10,11 @@ from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger -from vllm.models.common.mono.spec import ( +from vllm.models.common.specialized_decode.spec import ( + DecodeSpec, + DecodeUnsupported, Feature, KernelT, - MonoSpec, - MonoUnsupported, ) if TYPE_CHECKING: @@ -44,9 +44,10 @@ def tp_group() -> tuple[int, int, object | None]: def tp_agree(ok: bool, group) -> bool: """Agrees a verdict every TP rank must take alike. - A MonoKernel waits in-kernel on every peer rank, so a rank that turned mono - off alone, or failed to build what the others launch, leaves the others - spinning on the GPU for ever. Every rank must call this at the same point. + The kernel waits in-kernel on every peer rank, so a rank that turned + specialized decode off alone, or failed to build what the others launch, + leaves the others spinning on the GPU for ever. Every rank must call this + at the same point. Args: ok: This rank's local verdict. @@ -64,7 +65,7 @@ def tp_agree(ok: bool, group) -> bool: def _enabled(vllm_config: "VllmConfig") -> bool: - """Whether ``--monokernel`` asked for the model's MonoKernels.""" + """Whether ``--monokernel`` asked for the model's specialized decode.""" return bool(vllm_config.kernel_config.monokernel) @@ -94,7 +95,7 @@ def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: return {f for f, used in on.items() if used} -def _refuse_config(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> list[str]: +def _refuse_config(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> list[str]: """Every reason this config cannot run ``spec``: default-deny on features, then the runner and the graph mode, then the spec's own.""" reasons = [ @@ -109,8 +110,8 @@ def _refuse_config(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> list[str]: return reasons + list(spec.refuse_config(vllm_config)) -def _refuse_model(spec: MonoSpec[Any], model) -> list[str]: - """Every reason the loaded model cannot run ``spec``. A MonoKernel replaces +def _refuse_model(spec: DecodeSpec[Any], model) -> list[str]: + """Every reason the loaded model cannot run ``spec``. Specialized decode replaces the layer loop, which a traced graph has already inlined, so a model carrying ``@support_torch_compile`` is refused.""" from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper @@ -118,15 +119,15 @@ def _refuse_model(spec: MonoSpec[Any], model) -> list[str]: reasons = [] if isinstance(model, TorchCompileWithNoGuardsWrapper): reasons.append( - "model is torch.compile'd: MonoKernels replace the layer loop, " + "model is torch.compile'd: specialized decode replaces the layer loop, " "so the model must not carry @support_torch_compile" ) return reasons + list(spec.refuse_model(model)) def _select_spec( - specs: Sequence[MonoSpec[Any]], vllm_config: "VllmConfig" -) -> tuple[MonoSpec[Any] | None, dict[str, list[str]]]: + specs: Sequence[DecodeSpec[Any]], vllm_config: "VllmConfig" +) -> tuple[DecodeSpec[Any] | None, dict[str, list[str]]]: """The first variant this config can serve, and every variant's refusals. Returns: @@ -150,7 +151,7 @@ def _refusal_message(refusals: dict[str, list[str]]) -> str: ) -def _set_capture_sizes(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> None: +def _set_capture_sizes(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> None: """The capture sizes up to the largest width become exactly the widths, so a replay never pads a request's rows; larger sizes stay as configured (and capture vLLM's path).""" @@ -159,20 +160,22 @@ def _set_capture_sizes(spec: MonoSpec[Any], vllm_config: "VllmConfig") -> None: sizes = cc.cudagraph_capture_sizes or [] new = sorted(set(widths) | {s for s in sizes if s > max(widths)}) if new != sorted(sizes): - logger.info("MonoKernel %s: CUDA graph capture sizes %s", spec.name, new) + logger.info( + "Specialized decode %s: CUDA graph capture sizes %s", spec.name, new + ) cc.cudagraph_capture_sizes = new -def select_mono_spec( +def select_decode_spec( vllm_config: "VllmConfig", - specs: MonoSpec[KernelT] | Sequence[MonoSpec[KernelT]], -) -> MonoSpec[KernelT] | None: + specs: DecodeSpec[KernelT] | Sequence[DecodeSpec[KernelT]], +) -> DecodeSpec[KernelT] | None: """Picks the variant, sets the capture sizes, or raises with every reason. The variant is picked here, once: the model holds what this returns and - hands it to ``MonoRuntime`` after the weights are loaded. A model that prepares - anything of its own at load time reads the returned spec to know which - variant it is preparing for. + hands it to ``SpecializedDecode`` after the weights are loaded. A model + that prepares anything of its own at load time reads the returned spec to + know which variant it is preparing for. Args: vllm_config: The engine's config, before the weights are loaded. @@ -182,20 +185,22 @@ def select_mono_spec( The variant this config runs, or None with ``--monokernel`` off. Raises: - MonoUnsupported: Every variant refused this config. + DecodeUnsupported: Every variant refused this config. """ if not _enabled(vllm_config): return None - ordered = (specs,) if isinstance(specs, MonoSpec) else tuple(specs) + ordered = (specs,) if isinstance(specs, DecodeSpec) else tuple(specs) spec, refusals = _select_spec(ordered, vllm_config) if spec is None: - raise MonoUnsupported("MonoKernel refused: " + _refusal_message(refusals)) + raise DecodeUnsupported( + "Specialized decode refused: " + _refusal_message(refusals) + ) _set_capture_sizes(spec, vllm_config) if vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: logger.warning( - "MonoKernel %s under eager mode: for debugging only, serve with FULL " - "CUDA graphs", + "Specialized decode %s under eager mode: for debugging only, serve " + "with FULL CUDA graphs", spec.name, ) return spec diff --git a/vllm/models/common/mono/runtime.py b/vllm/models/common/specialized_decode/runtime.py similarity index 73% rename from vllm/models/common/mono/runtime.py rename to vllm/models/common/specialized_decode/runtime.py index 085096d73717..48127b99011c 100644 --- a/vllm/models/common/mono/runtime.py +++ b/vllm/models/common/specialized_decode/runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""One ``MonoRuntime`` a rank: the per-step decision and the cache checks, +"""One ``SpecializedDecode`` a rank: the per-step decision and the cache checks, behind the two hooks a model's forward calls (``plan``, then ``run``). """ @@ -13,12 +13,16 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import VllmJitKernel -from vllm.models.common.mono.guards import _refuse_model, tp_agree, tp_group -from vllm.models.common.mono.spec import ( +from vllm.models.common.specialized_decode.guards import ( + _refuse_model, + tp_agree, + tp_group, +) +from vllm.models.common.specialized_decode.spec import ( + DecodeSpec, + DecodeUnsupported, KernelT, ModelKernel, - MonoSpec, - MonoUnsupported, Step, ) @@ -28,7 +32,7 @@ logger = init_logger(__name__) -class MonoWarmup(VllmJitKernel["MonoWarmup.CompileKey"]): +class DecodeWarmup(VllmJitKernel["DecodeWarmup.CompileKey"]): """A spec's kernel builds as vLLM JIT warm-up keys.""" @dataclass(frozen=True) @@ -41,16 +45,16 @@ def __init__(self, spec: str, kernel: ModelKernel): self.spec = spec self.kernel = kernel - def get_warmup_keys(self) -> list["MonoWarmup.CompileKey"]: + def get_warmup_keys(self) -> list["DecodeWarmup.CompileKey"]: return [self.CompileKey(self.spec, b) for b in self.kernel.warmup_keys()] - def compile(self, compile_key: "MonoWarmup.CompileKey") -> None: + def compile(self, compile_key: "DecodeWarmup.CompileKey") -> None: self.kernel.compile(compile_key.build) self._compiled_cache[compile_key] = True -class MonoRuntime(Generic[KernelT]): - """A model's MonoKernel on this rank. +class SpecializedDecode(Generic[KernelT]): + """A model's specialized decode on this rank. The runtime reads the spec and builds the kernel, registers it for warm-up, checks the caches whenever the runner binds them, and decides each step @@ -58,34 +62,34 @@ class MonoRuntime(Generic[KernelT]): KernelT and is held in the ``kernel`` attribute. """ - def __init__(self, model, vllm_config: "VllmConfig", spec: MonoSpec[KernelT]): + def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): """Builds this rank's kernel. Collective over the TP group. Args: model: The loaded model, in ``process_weights_after_loading``. vllm_config: The engine's config. - spec: The variant ``select_mono_spec`` picked. + spec: The variant ``select_decode_spec`` picked. Raises: - MonoUnsupported: This rank, or another, refused the loaded model. + DecodeUnsupported: This rank, or another, refused the loaded model. """ self.spec = spec reasons = _refuse_model(spec, model) agreed = tp_agree(not reasons, tp_group()[2]) - refused = f"MonoKernel {spec.name} refused for the loaded model" + refused = f"Specialized decode {spec.name} refused for the loaded model" if reasons: - raise MonoUnsupported(f"{refused}: " + "; ".join(reasons)) + raise DecodeUnsupported(f"{refused}: " + "; ".join(reasons)) if not agreed: - raise MonoUnsupported(f"{refused}: another TP rank refused") + raise DecodeUnsupported(f"{refused}: another TP rank refused") self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) self.last_reason: str | None = None self._bound = False self.kernel: KernelT = spec.build(vllm_config, model, self.widths) if self.kernel.warmup_keys(): - MonoWarmup(spec.name, self.kernel).register_warmup() - logger.info("MonoKernel %s: widths %s", spec.name, list(self.widths)) + DecodeWarmup(spec.name, self.kernel).register_warmup() + logger.info("Specialized decode %s: widths %s", spec.name, list(self.widths)) def bind_caches(self) -> None: """Checks the caches the kernel will read, as the runner binds them. @@ -94,15 +98,15 @@ def bind_caches(self) -> None: throwaway cache memory profiling allocates and again for the real one. Raises: - MonoUnsupported: The caches are not what this spec serves. + DecodeUnsupported: The caches are not what this spec serves. """ sfc = self.vllm_config.compilation_config.static_forward_context caches = {n: sfc[n].kv_cache for n in self.kernel.cache_layers} reasons = self.spec.refuse_caches(self.vllm_config, self.kernel, caches) if reasons: - raise MonoUnsupported( - f"MonoKernel {self.spec.name} caches: " + "; ".join(reasons) + raise DecodeUnsupported( + f"Specialized decode {self.spec.name} caches: " + "; ".join(reasons) ) self._bound = True @@ -114,7 +118,7 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non it is replayed for mixed batches too), and whatever ``skip_step`` refuses. ``last_reason`` holds why. - Every predicate is the same on every TP rank: a MonoKernel waits + Every predicate is the same on every TP rank: the kernel waits in-kernel on its peers, so the ranks must take the same path. """ width = positions.size(-1) @@ -134,7 +138,7 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non self.last_reason = why if why is not None: logger.info_once( - "MonoKernel %s: step on vLLM's path (%s)", self.spec.name, why + "Specialized decode %s: step on vLLM's path (%s)", self.spec.name, why ) return None return Step(width=width, positions=positions) diff --git a/vllm/models/common/mono/spec.py b/vllm/models/common/specialized_decode/spec.py similarity index 89% rename from vllm/models/common/mono/spec.py rename to vllm/models/common/specialized_decode/spec.py index e7a8665ba08b..66e5777848cf 100644 --- a/vllm/models/common/mono/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""What a model declares to run MonoKernels: ``MonoSpec`` and its callbacks.""" +"""What a model declares to run specialized decode: ``DecodeSpec`` and its +callbacks.""" from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -32,13 +33,13 @@ class Feature(str, Enum): WEIGHT_UPDATES = "weight_updates" -class MonoUnsupported(Exception): +class DecodeUnsupported(Exception): """The configuration or the loaded model is outside what a spec serves.""" @dataclass(frozen=True) class Step: - """One decode step the kernel takes, as ``MonoRuntime.plan`` decided it.""" + """One decode step the kernel takes, as ``SpecializedDecode.plan`` decided it.""" width: int """The step's rows, one of the widths the spec serves.""" @@ -55,7 +56,7 @@ class ModelKernel(Protocol): """ cache_layers: tuple[str, ...] - """Attention layers whose caches ``MonoSpec.refuse_caches`` checks.""" + """Attention layers whose caches ``DecodeSpec.refuse_caches`` checks.""" def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: """Takes and returns what the model's own forward does. @@ -80,12 +81,12 @@ def compile(self, key: Any) -> None: KernelT = TypeVar("KernelT", bound=ModelKernel) """What one spec's ``build`` returns. ``ModelKernel`` is only the lower bound, so a model annotates its runtime with its own kernel class -(``MonoRuntime[MyKernel]``) and reads what that class adds.""" +(``SpecializedDecode[MyKernel]``) and reads what that class adds.""" @dataclass(frozen=True) -class MonoSpec(Generic[KernelT]): - """Definition of when a monokernel can run, and how to build it. +class DecodeSpec(Generic[KernelT]): + """Definition of when specialized decode can run, and how to build it. Attributes: name: Log and metric name. diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 992c25ab960c..c5ff8127fcaf 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -49,18 +49,18 @@ make_layers, maybe_prefix, ) -from vllm.models.common.mono import MonoRuntime, select_mono_spec from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, sp_reduce_scatter, sp_shard, ) +from vllm.models.common.specialized_decode import SpecializedDecode, select_decode_spec from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v41.amd.mono_spec import MONO_SPECS, DSV41Kernel from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention +from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPECS, DSV41Kernel from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors @@ -283,7 +283,7 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - mono_ffn: typing.Any | None = None, + decode_ffn: typing.Any | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # Layer 0's attention seam projects the 2-D embedding with the folded # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn. @@ -384,12 +384,12 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) elif not self.attn.wo_b.reduce_results: - # A partial-layer MonoKernel: the FFN launch takes wo_b's TP partial - # and the rest of the layer. Whole-layer launches never come through - # here, and a new kernel should cover a whole layer instead of - # reaching into one. - if mono_ffn is not None: - out = mono_ffn.launch(x, residual, post_mix, res_mix, attn_pre) + # A partial-layer specialized decode: the FFN launch takes wo_b's + # TP partial and the rest of the layer. Whole-layer launches never + # come through here, and a new kernel should cover a whole layer + # instead of reaching into one. + if decode_ffn is not None: + out = decode_ffn.launch(x, residual, post_mix, res_mix, attn_pre) if out is not None: return out x = tensor_model_parallel_all_reduce(x) @@ -1386,8 +1386,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self.mono_spec = select_mono_spec(vllm_config, MONO_SPECS) - self.mono: MonoRuntime[DSV41Kernel] | None = None + self.specialized_decode_spec = select_decode_spec(vllm_config, DECODE_SPECS) + self.specialized_decode: SpecializedDecode[DSV41Kernel] | None = None expert_dtype = getattr(config, "expert_dtype", "fp4") self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype) @@ -1471,10 +1471,10 @@ def forward( inputs_embeds: torch.Tensor | None = None, lookback_token_ids: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: - if self.mono is not None: - step = self.mono.plan(positions, intermediate_tensors) + if self.specialized_decode is not None: + step = self.specialized_decode.plan(positions, intermediate_tensors) if step is not None: - return self.mono.run( + return self.specialized_decode.run( step, input_ids, inputs_embeds, @@ -1504,8 +1504,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() - if self.mono_spec is not None: - self.mono = MonoRuntime(self.model, self.vllm_config, self.mono_spec) + if self.specialized_decode_spec is not None: + self.specialized_decode = SpecializedDecode( + self.model, self.vllm_config, self.specialized_decode_spec + ) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v41/amd/model_state.py b/vllm/models/deepseek_v41/amd/model_state.py index 84d545158c73..a3e1ef027058 100644 --- a/vllm/models/deepseek_v41/amd/model_state.py +++ b/vllm/models/deepseek_v41/amd/model_state.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek-V4.1's model runner state on ROCm: the shared state plus the -MonoKernel's cache checks.""" +specialized decode's cache checks.""" from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState from vllm.v1.kv_cache_interface import KVCacheConfig @@ -9,12 +9,12 @@ class DeepseekV41ROCMModelState(DeepseekV41ModelState): - """Checks the MonoKernel's caches for every cache the runner binds: the + """Checks the specialized decode's caches for every cache the runner binds: the throwaway one memory profiling allocates, then the real one.""" def initialize_kv_cache( self, kv_cache_config: KVCacheConfig, block_tables: BlockTables ) -> None: super().initialize_kv_cache(kv_cache_config, block_tables) - if self.model.mono is not None: - self.model.mono.bind_caches() + if self.model.specialized_decode is not None: + self.model.specialized_decode.bind_caches() diff --git a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py index afab63471e1c..07a3ba528637 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py +++ b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py @@ -9,8 +9,8 @@ import torch from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap -from vllm.models.common.mono.guards import tp_agree -from vllm.models.common.mono.spec import MonoUnsupported +from vllm.models.common.specialized_decode.guards import tp_agree +from vllm.models.common.specialized_decode.spec import DecodeUnsupported class _DeviceBytes: @@ -34,7 +34,7 @@ class PeerBuffer: The constructor is collective over ``group`` (the handle exchange, then an agreement that every rank mapped every peer): every rank must reach it, so a runner reaches it only after ``tp_agree``. A failed mapping on any rank raises - ``MonoUnsupported`` on all of them. + ``DecodeUnsupported`` on all of them. """ def __init__( @@ -68,7 +68,7 @@ def __init__( failure = err if not tp_agree(failure is None, group): self.close() - raise MonoUnsupported( + raise DecodeUnsupported( "peer memory handshake failed on " f"{'this' if failure else 'another'}" " TP rank" diff --git a/vllm/models/deepseek_v41/amd/mono_spec.py b/vllm/models/deepseek_v41/amd/specialized_decode.py similarity index 92% rename from vllm/models/deepseek_v41/amd/mono_spec.py rename to vllm/models/deepseek_v41/amd/specialized_decode.py index 80372ebd4b7e..024dcc802904 100644 --- a/vllm/models/deepseek_v41/amd/mono_spec.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's MonoKernels on ROCm CDNA4 (``--monokernel``, RFC #60904).""" +"""DeepSeek-V4.1's specialized decode on ROCm CDNA4 (``--monokernel``, +RFC #60904).""" from typing import TYPE_CHECKING, Any, cast @@ -10,12 +11,12 @@ from vllm.config import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.common.mono.spec import Feature, MonoSpec, Step +from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature, Step from vllm.platforms import current_platform if TYPE_CHECKING: from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer - from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterSparseSWAMetadata logger = init_logger(__name__) @@ -73,12 +74,12 @@ class DSV41LayerKernel: ffn_only = False def __init__( - self, mono: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int + self, decode: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int ): - self.mono = mono + self.decode = decode self.layer = layer self.width = width - self.weights = mono.weights[layer.attn.layer_id] + self.weights = decode.weights[layer.attn.layer_id] self.cache_layers: tuple[str, ...] = () def warmup_keys(self) -> list[Any]: @@ -95,9 +96,9 @@ class WholeLayer(DSV41LayerKernel): """The whole backbone layer as K1 and K2, from the layer's own inputs.""" def __init__( - self, mono: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int + self, decode: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int ): - super().__init__(mono, layer, width) + super().__init__(decode, layer, width) attn = layer.attn self.ratio = attn.compress_ratio self.index = attn.indexer is not None @@ -108,7 +109,7 @@ def __init__( ) # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == mono.embed_layer or layer.engram is not None + settled = attn.layer_id == decode.embed_layer or layer.engram is not None self.folds = (True, False) if settled else (True,) def warmup_keys(self) -> list[Any]: @@ -116,8 +117,8 @@ def warmup_keys(self) -> list[Any]: def compile(self, key: Any) -> None: width, ratio, fold, index = key - self.mono.runner.scratch(width) - self.mono.runner.kernels(width, ratio, fold=fold, index=index) + self.decode.runner.scratch(width) + self.decode.runner.kernels(width, ratio, fold=fold, index=index) def launch( self, @@ -131,7 +132,7 @@ def launch( engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: layer, attn = self.layer, self.layer.attn - runner = self.mono.runner + runner = self.decode.runner first = residual is None if first and x.dim() != 2: return None @@ -191,8 +192,8 @@ def warmup_keys(self) -> list[Any]: return [(self.width, "ffn")] def compile(self, key: Any) -> None: - self.mono.runner.scratch(key[0]) - self.mono.runner.ffn_kernel(key[0]) + self.decode.runner.scratch(key[0]) + self.decode.runner.ffn_kernel(key[0]) def launch( self, @@ -205,7 +206,7 @@ def launch( tensors = (part, residual, post_mix, res_mix, pre_mix) if not all(t.is_contiguous() for t in tensors): return None - return self.mono.runner.ffn(self.weights, *tensors) + return self.decode.runner.ffn(self.weights, *tensors) class DSV41Kernel: @@ -222,7 +223,7 @@ class DSV41Kernel: """ def __init__(self, vllm_config, model, widths: tuple[int, ...]): - from vllm.models.common.mono.guards import tp_group + from vllm.models.common.specialized_decode.guards import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS @@ -259,7 +260,7 @@ def _kernel_layers(self, model) -> list["DeepseekV4DecoderLayer"]: layers.append(model.layers[i]) else: logger.info_once( - "DSv4.1 MonoKernel: layer %d on vLLM's path (%s)", i, why + "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", i, why ) return layers @@ -270,7 +271,7 @@ def _build_layer(self, layer, tp_size: int, widths: tuple[int, ...]) -> None: launch: type[DSV41LayerKernel] = WholeLayer if why is not None: logger.info_once( - "DSv4.1 MonoKernel: the FFN launch for layers with %s", why + "DSv4.1 specialized decode: the FFN launch for layers with %s", why ) # A decode step hands wo_b's TP partial to the FFN launch; the layer # reduces it on the steps that keep vLLM's path. @@ -339,7 +340,7 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: residual, engram_hashes, engram_mask, - mono_ffn=ffn, + decode_ffn=ffn, ) hidden, residual, post_mix, res_mix, pre_mix = out aux_recon = model.after_layer( @@ -464,13 +465,13 @@ def _ffn_only(layer) -> str | None: return None -def _skip_step(vllm_config, mono: DSV41Kernel, width: int) -> str | None: +def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP group sees alike.""" fc = get_forward_context() md = fc.attn_metadata - if not mono.kernels: + if not decode.kernels: return "no layer takes a launch" if not 1 <= width <= MAX_ROWS: return f"{width} rows" @@ -479,17 +480,17 @@ def _skip_step(vllm_config, mono: DSV41Kernel, width: int) -> str | None: return f"{fc.cudagraph_runtime_mode} step" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches return "no attention metadata" - if any(md.get(k) is None for k in mono.comp_keys): + if any(md.get(k) is None for k in decode.comp_keys): return "layer metadata missing" - for key in mono.swa_keys: - swa = cast("DeepseekV4ROCMAiterSparseSWAMetadata | None", md.get(key)) + for key in decode.swa_keys: + swa = cast("DeepseekV41ROCMAiterSparseSWAMetadata | None", md.get(key)) if swa is None: return "layer metadata missing" if swa.num_prefills != 0 or swa.num_decodes == 0: return "not a decode-only step" if swa.num_decode_tokens > width: return "more decode tokens than rows" - if not mono.need_swa: + if not decode.need_swa: continue if swa.decode_swa_width != SWA_WIDTH: return f"SWA width {swa.decode_swa_width}" @@ -556,7 +557,7 @@ def _attn_weights(a, world: int): return attn -def _refuse_caches(vllm_config, mono: DSV41Kernel, caches) -> list[str]: +def _refuse_caches(vllm_config, decode: DSV41Kernel, caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as [blocks, block, RECORD] bytes.""" reasons = [] @@ -577,7 +578,7 @@ def _build(vllm_config, model, widths: tuple[int, ...]) -> DSV41Kernel: return DSV41Kernel(vllm_config, model, widths) -MONO_SPEC = MonoSpec[DSV41Kernel]( +DECODE_SPEC = DecodeSpec[DSV41Kernel]( name="dsv41-cdna4", supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), widths=_widths, @@ -588,4 +589,4 @@ def _build(vllm_config, model, widths: tuple[int, ...]) -> DSV41Kernel: build=_build, ) -MONO_SPECS = (MONO_SPEC,) +DECODE_SPECS = (DECODE_SPEC,) From 9892ef0ce00afec4c720554be260ff307aa9ca10 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 17:05:58 -0500 Subject: [PATCH 36/79] fix: remove redundant ffn/whole layer distinction in the abstraction Signed-off-by: simondanielsson --- tests/models/test_dsv41_mono_layers.py | 13 +- tests/models/test_dsv41_specialized_decode.py | 52 +-- vllm/models/deepseek_v41/amd/model.py | 4 +- .../deepseek_v41/amd/specialized_decode.py | 300 +++++++++--------- 4 files changed, 179 insertions(+), 190 deletions(-) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index bb1e573ab28c..dea44b0acc74 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -194,15 +194,12 @@ def save_capture(self, path: str) -> None: def mono_layers(self) -> dict: model = self.model_runner.model # type: ignore[attr-defined] decode = getattr(model, "specialized_decode", None) - kernels = {} if decode is None else decode.kernel.kernels + plans = {} if decode is None else decode.kernel.plans layers = { - i: next( - ( - ["ffn" if k.ffn_only else "whole", k.weights is not None] - for (j, _), k in kernels.items() - if j == i - ), - None, + i: ( + ["ffn" if plans[i].ffn_only else "whole", plans[i].weights is not None] + if i in plans + else None ) for i in self._decoder_layers() } diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index eab098b12e71..03257136382c 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -321,8 +321,8 @@ def _build(monkeypatch, config=None, model=None, **layer_overrides): return kit.make_decode(config, model), model -def _launch(decode, layer_id: int, width: int = 1): - return decode.kernel.kernels[layer_id, width] +def _plan(decode, layer_id: int): + return decode.kernel.plans[layer_id] def _indexer(): @@ -362,7 +362,7 @@ def test_layers_take_whole_or_ffn_launch(monkeypatch, overrides, ffn_only): them keeps vLLM's attention and runs its FFN half as one launch, from wo_b's unreduced output.""" decode, model = _build(monkeypatch, **overrides()) - assert {k.ffn_only for k in decode.kernel.kernels.values()} == {ffn_only} + assert {p.ffn_only for p in decode.kernel.plans.values()} == {ffn_only} assert all(layer.attn.wo_b.reduce_results is not ffn_only for layer in model.layers) @@ -381,7 +381,7 @@ def test_skipped_layers_keep_vllm_path(monkeypatch, skip): for key, value in skip.items(): setattr(layer, key, value) decode, _ = _build(monkeypatch, config=config, model=model) - assert not decode.kernel.kernels + assert not decode.kernel.plans def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): @@ -390,7 +390,7 @@ def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): for layer in model.layers: layer.ffn.n_routed_experts = 256 decode, _ = _build(monkeypatch, config=config, model=model) - assert not decode.kernel.kernels + assert not decode.kernel.plans def _step(decode, model, metadata, width=1): @@ -493,8 +493,8 @@ def test_whole_launch_runs_the_layer(monkeypatch): decode, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(decode, 1) - out = kernel.launch( + out = decode.kernel.launch_whole( + _plan(decode, 1), torch.zeros(1, HIDDEN), positions, torch.zeros(1, 4, HIDDEN), @@ -515,8 +515,9 @@ def test_first_layer_takes_the_embedding(monkeypatch): decode, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(decode, 0) - out = kernel.launch(torch.zeros(1, HIDDEN), positions, None, None, None, None) + out = decode.kernel.launch_whole( + _plan(decode, 0), torch.zeros(1, HIDDEN), positions, None, None, None, None + ) assert out[0] == "whole" and not out[2]["fold"] assert out[1][2].shape == (1, 4, HIDDEN) @@ -524,9 +525,10 @@ def test_first_layer_takes_the_embedding(monkeypatch): def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): decode, model = _build(monkeypatch, compressor=object()) with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(decode, 1) - assert kernel.ffn_only - out = kernel.launch( + plan = _plan(decode, 1) + assert plan.ffn_only + out = decode.kernel.launch_ffn( + plan, torch.zeros(1, HIDDEN), torch.zeros(1, 4, HIDDEN), torch.zeros(1, 4, 1), @@ -545,14 +547,13 @@ def test_run_takes_every_layer_the_way_its_launch_allows(monkeypatch): model.layers[2].attn.compressor = object() model.layers[4].use_sequence_parallel = True decode, _ = _build(monkeypatch, config=config, model=model) - for (layer_id, width), kernel in decode.kernel.kernels.items(): - if width != 1 or kernel.ffn_only: - continue - taken = model.layers[layer_id].taken - declines = layer_id == 1 - kernel.launch = lambda *a, _t=taken, _id=layer_id, _d=declines, **kw: ( - _t.append(("launch", _id)) or (None if _d else _stream()) - ) + + def launch_whole(plan, *args, **kwargs): + layer_id = plan.layer.attn.layer_id + model.layers[layer_id].taken.append(("launch", layer_id)) + return None if layer_id == 1 else _stream() + + decode.kernel.launch_whole = launch_whole step = Step(width=1, positions=torch.zeros(1, dtype=torch.long)) out = decode.kernel.run(step, None, torch.zeros(1, HIDDEN)) @@ -574,8 +575,13 @@ def test_launch_declines_a_layout_it_cannot_read(monkeypatch): decode, model = _build(monkeypatch) positions = torch.zeros(1, dtype=torch.long) with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): - kernel = _launch(decode, 1) - out = kernel.launch( - torch.zeros(1, 4, HIDDEN), positions, None, None, None, None + out = decode.kernel.launch_whole( + _plan(decode, 1), + torch.zeros(1, 4, HIDDEN), + positions, + None, + None, + None, + None, ) assert out is None diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index c5ff8127fcaf..7b8b80995c3d 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -283,7 +283,7 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - decode_ffn: typing.Any | None = None, + decode_ffn: Callable | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # Layer 0's attention seam projects the 2-D embedding with the folded # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn. @@ -389,7 +389,7 @@ def forward( # come through here, and a new kernel should cover a whole layer # instead of reaching into one. if decode_ffn is not None: - out = decode_ffn.launch(x, residual, post_mix, res_mix, attn_pre) + out = decode_ffn(x, residual, post_mix, res_mix, attn_pre) if out is not None: return out x = tensor_model_parallel_all_reduce(x) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index 024dcc802904..2f5fced006b2 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -3,6 +3,8 @@ """DeepSeek-V4.1's specialized decode on ROCm CDNA4 (``--monokernel``, RFC #60904).""" +from dataclasses import dataclass +from functools import partial from typing import TYPE_CHECKING, Any, cast import torch @@ -67,61 +69,140 @@ def _select( ) -class DSV41LayerKernel: - """What both launches share: the layer, its width and its weight views, - with the launch shape left to the subclass.""" +@dataclass(frozen=True) +class LayerPlan: + """One decoder layer's launch, settled at load. - ffn_only = False + An ``ffn_only`` layer keeps vLLM's attention and its launch takes wo_b's + unreduced output, so its weights are all it has; the rest is what a + whole-layer launch reads off the layer. + """ - def __init__( - self, decode: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int - ): - self.decode = decode - self.layer = layer - self.width = width - self.weights = decode.weights[layer.attn.layer_id] - self.cache_layers: tuple[str, ...] = () + layer: "DeepseekV4DecoderLayer" + weights: Any + ffn_only: bool + cache_layers: tuple[str, ...] = () + ratio: int = 0 + index: bool = False + swa_key: str | None = None + comp_key: str | None = None + folds: tuple[bool, ...] = () - def warmup_keys(self) -> list[Any]: - raise NotImplementedError - def compile(self, key: Any) -> None: - raise NotImplementedError +class DSV41Kernel: + """DeepSeek-V4.1's decode on this rank: the runner (one a rank, its peer + memory collective over the TP group) and a launch plan per kernel layer. - def launch(self, *args, **kwargs) -> Any: - raise NotImplementedError + ``run`` owns the backbone's layer loop, between the model's own prologue + and epilogue: those keep the engram hashes, the sequence-parallel shard and + the final collapse on one code path. + Built once at load and never moved: the launches and the captured graphs + hold its addresses. + """ -class WholeLayer(DSV41LayerKernel): - """The whole backbone layer as K1 and K2, from the layer's own inputs.""" + def __init__(self, vllm_config, model, widths: tuple[int, ...]): + from vllm.models.common.specialized_decode.guards import tp_group + from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer - def __init__( - self, decode: "DSV41Kernel", layer: "DeepseekV4DecoderLayer", width: int - ): - super().__init__(decode, layer, width) - attn = layer.attn - self.ratio = attn.compress_ratio - self.index = attn.indexer is not None - self.swa_key = attn.swa_cache_layer.prefix - self.comp_key = attn.compressed_cache_prefix - self.cache_layers = (self.swa_key,) + ( - (self.comp_key,) if self.comp_key is not None else () + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + rank, world, cpu_group = tp_group() + device = next(model.parameters()).device + self.model = model + self.runner = DSV41MonoLayer(world, rank, cpu_group, device) + self.embed_layer = getattr(model, "start_layer", 0) + # The replay layers run on their own trimmed batch, outside the loop + # ``run`` owns, so no launch is built for them. + end = getattr(model, "end_layer", len(model.layers)) + self.layer_end = getattr(model, "decoder_replay_start", end) + self.words = torch.zeros(world, dtype=torch.int32, device=device) + self.widths = widths + self.swa_keys: list[str] = [] + self.comp_keys: list[str] = [] + self.need_swa = False + self.plans: dict[int, LayerPlan] = {} + + tp_size = vllm_config.parallel_config.tensor_parallel_size + for layer in self._kernel_layers(model): + plan = self._plan_layer(layer, tp_size) + self.plans[layer.attn.layer_id] = plan + if plan.ffn_only: + continue + self.need_swa = True + for key, keys in ( + (plan.swa_key, self.swa_keys), + (plan.comp_key, self.comp_keys), + ): + if key is not None and key not in keys: + keys.append(key) + self.cache_layers = tuple( + sorted({n for p in self.plans.values() for n in p.cache_layers}) ) + + def _kernel_layers(self, model) -> list["DeepseekV4DecoderLayer"]: + end = self.layer_end + layers = [] + for i in range(self.embed_layer, end): + why = _skip_layer(model.layers[i]) + if why is None: + layers.append(model.layers[i]) + else: + logger.info_once( + "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", i, why + ) + return layers + + def _plan_layer(self, layer, tp_size: int) -> LayerPlan: + why = _ffn_only(layer) + weights = _weights(layer, why is not None, tp_size) + if why is not None: + logger.info_once( + "DSv4.1 specialized decode: the FFN launch for layers with %s", why + ) + # A decode step hands wo_b's TP partial to the FFN launch; the layer + # reduces it on the steps that keep vLLM's path. + layer.attn.wo_b.reduce_results = False + return LayerPlan(layer, weights, ffn_only=True) + attn = layer.attn + comp_key = attn.compressed_cache_prefix # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == decode.embed_layer or layer.engram is not None - self.folds = (True, False) if settled else (True,) + settled = attn.layer_id == self.embed_layer or layer.engram is not None + return LayerPlan( + layer, + weights, + ffn_only=False, + cache_layers=(attn.swa_cache_layer.prefix,) + + ((comp_key,) if comp_key is not None else ()), + ratio=attn.compress_ratio, + index=attn.indexer is not None, + swa_key=attn.swa_cache_layer.prefix, + comp_key=comp_key, + folds=(True, False) if settled else (True,), + ) def warmup_keys(self) -> list[Any]: - return [(self.width, self.ratio, f, self.index) for f in self.folds] + keys: dict[Any, None] = {} + for plan in self.plans.values(): + for width in self.widths: + if plan.ffn_only: + keys[width, "ffn"] = None + else: + for fold in plan.folds: + keys[width, plan.ratio, fold, plan.index] = None + return list(keys) def compile(self, key: Any) -> None: - width, ratio, fold, index = key - self.decode.runner.scratch(width) - self.decode.runner.kernels(width, ratio, fold=fold, index=index) + self.runner.scratch(key[0]) + if key[1] == "ffn": + self.runner.ffn_kernel(key[0]) + else: + width, ratio, fold, index = key + self.runner.kernels(width, ratio, fold=fold, index=index) - def launch( + def launch_whole( self, + plan: LayerPlan, x: torch.Tensor, positions: torch.Tensor, residual: torch.Tensor | None, @@ -131,8 +212,9 @@ def launch( engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: - layer, attn = self.layer, self.layer.attn - runner = self.decode.runner + """The whole backbone layer as K1 and K2, from the layer's own inputs.""" + layer, runner = plan.layer, self.runner + attn = layer.attn first = residual is None if first and x.dim() != 2: return None @@ -151,21 +233,21 @@ def launch( residual = layer.engram(residual, hashes, engram_mask) fold = False md = cast("dict[str, Any]", get_forward_context().attn_metadata) - swa = md[self.swa_key] + swa = md[cast("str", plan.swa_key)] kw: dict[str, Any] = {} - if self.index: + if plan.index: kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) - if self.ratio: - comp = md[cast("str", self.comp_key)] + if plan.ratio: + comp = md[cast("str", plan.comp_key)] kw.update( topk_indices=attn.topk_indices_buffer, comp_cache=_record_view( - attn._compressed_kv_cache(), comp.block_size // self.ratio + attn._compressed_kv_cache(), comp.block_size // plan.ratio ), comp_block_table=comp.block_table, ) return runner.forward( - self.weights, + plan.weights, x, residual, post_mix, @@ -181,121 +263,21 @@ def launch( **kw, ) - -class FfnSeam(DSV41LayerKernel): - """The FFN seam, MoE and both reductions as one launch, from wo_b's - unreduced output after vLLM's attention.""" - - ffn_only = True - - def warmup_keys(self) -> list[Any]: - return [(self.width, "ffn")] - - def compile(self, key: Any) -> None: - self.decode.runner.scratch(key[0]) - self.decode.runner.ffn_kernel(key[0]) - - def launch( + def launch_ffn( self, + plan: LayerPlan, part: torch.Tensor, residual: torch.Tensor, post_mix: torch.Tensor, res_mix: torch.Tensor, pre_mix: torch.Tensor, ) -> tuple[torch.Tensor, ...] | None: + """The FFN seam, MoE and both reductions as one launch, from wo_b's + unreduced output after vLLM's attention.""" tensors = (part, residual, post_mix, res_mix, pre_mix) if not all(t.is_contiguous() for t in tensors): return None - return self.decode.runner.ffn(self.weights, *tensors) - - -class DSV41Kernel: - """DeepSeek-V4.1's decode on this rank: the runner (one a rank, its peer - memory collective over the TP group), each kernel layer's weight views, and - a launch per (layer, width). - - ``run`` owns the backbone's layer loop, between the model's own prologue - and epilogue: those keep the engram hashes, the sequence-parallel shard and - the final collapse on one code path. - - Built once at load and never moved: the launches and the captured graphs - hold its addresses. - """ - - def __init__(self, vllm_config, model, widths: tuple[int, ...]): - from vllm.models.common.specialized_decode.guards import tp_group - from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer - - assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS - rank, world, cpu_group = tp_group() - device = next(model.parameters()).device - self.model = model - self.runner = DSV41MonoLayer(world, rank, cpu_group, device) - self.embed_layer = getattr(model, "start_layer", 0) - # The replay layers run on their own trimmed batch, outside the loop - # ``run`` owns, so no launch is built for them. - end = getattr(model, "end_layer", len(model.layers)) - self.layer_end = getattr(model, "decoder_replay_start", end) - self.words = torch.zeros(world, dtype=torch.int32, device=device) - self.weights: dict[int, Any] = {} - self.swa_keys: list[str] = [] - self.comp_keys: list[str] = [] - self.need_swa = False - self.kernels: dict[tuple[int, int], DSV41LayerKernel] = {} - self._owners: dict[Any, DSV41LayerKernel] = {} - - tp_size = vllm_config.parallel_config.tensor_parallel_size - for layer in self._kernel_layers(model): - self._build_layer(layer, tp_size, widths) - self.cache_layers = tuple( - sorted({n for k in self.kernels.values() for n in k.cache_layers}) - ) - - def _kernel_layers(self, model) -> list["DeepseekV4DecoderLayer"]: - end = self.layer_end - layers = [] - for i in range(self.embed_layer, end): - why = _skip_layer(model.layers[i]) - if why is None: - layers.append(model.layers[i]) - else: - logger.info_once( - "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", i, why - ) - return layers - - def _build_layer(self, layer, tp_size: int, widths: tuple[int, ...]) -> None: - why = _ffn_only(layer) - layer_id = layer.attn.layer_id - self.weights[layer_id] = _weights(layer, why is not None, tp_size) - launch: type[DSV41LayerKernel] = WholeLayer - if why is not None: - logger.info_once( - "DSv4.1 specialized decode: the FFN launch for layers with %s", why - ) - # A decode step hands wo_b's TP partial to the FFN launch; the layer - # reduces it on the steps that keep vLLM's path. - layer.attn.wo_b.reduce_results = False - launch = FfnSeam - else: - self.need_swa = True - for key, keys in ( - (layer.attn.swa_cache_layer.prefix, self.swa_keys), - (layer.attn.compressed_cache_prefix, self.comp_keys), - ): - if key is not None and key not in keys: - keys.append(key) - for width in widths: - kernel = launch(self, layer, width) - self.kernels[layer_id, width] = kernel - for key in kernel.warmup_keys(): - self._owners.setdefault(key, kernel) - - def warmup_keys(self) -> list[Any]: - return list(self._owners) - - def compile(self, key: Any) -> None: - self._owners[key].compile(key) + return self.runner.ffn(plan.weights, *tensors) def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: model, positions = self.model, step.positions @@ -314,12 +296,16 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: final_aux_recon = None for idx in range(self.embed_layer, self.layer_end): layer = model.layers[idx] - launch = self.kernels.get((layer.attn.layer_id, step.width)) - ffn = launch if launch is not None and launch.ffn_only else None - whole = launch if launch is not None and ffn is None else None + plan = self.plans.get(layer.attn.layer_id) + ffn = ( + partial(self.launch_ffn, plan) + if plan is not None and plan.ffn_only + else None + ) out = None - if whole is not None: - out = whole.launch( + if plan is not None and not plan.ffn_only: + out = self.launch_whole( + plan, hidden, positions, residual, @@ -471,7 +457,7 @@ def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: group sees alike.""" fc = get_forward_context() md = fc.attn_metadata - if not decode.kernels: + if not decode.plans: return "no layer takes a launch" if not 1 <= width <= MAX_ROWS: return f"{width} rows" From 9d4467dca0a895200dbd3a2d182c91f43e985225 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 17:12:19 -0500 Subject: [PATCH 37/79] chore: rename modelkernel to decodeimpl Signed-off-by: simondanielsson --- tests/models/common/specialized_decode/reference.py | 2 +- vllm/models/common/specialized_decode/runtime.py | 4 ++-- vllm/models/common/specialized_decode/spec.py | 13 +++++++------ 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py index 23d5eb87b76d..89bb62637849 100644 --- a/tests/models/common/specialized_decode/reference.py +++ b/tests/models/common/specialized_decode/reference.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The reference kernel: a ``ModelKernel`` over the layers it replaces.""" +"""The reference kernel: a ``DecodeImpl`` over the layers it replaces.""" from dataclasses import replace from types import SimpleNamespace diff --git a/vllm/models/common/specialized_decode/runtime.py b/vllm/models/common/specialized_decode/runtime.py index 48127b99011c..6733376e99be 100644 --- a/vllm/models/common/specialized_decode/runtime.py +++ b/vllm/models/common/specialized_decode/runtime.py @@ -19,10 +19,10 @@ tp_group, ) from vllm.models.common.specialized_decode.spec import ( + DecodeImpl, DecodeSpec, DecodeUnsupported, KernelT, - ModelKernel, Step, ) @@ -40,7 +40,7 @@ class CompileKey: spec: str build: Any - def __init__(self, spec: str, kernel: ModelKernel): + def __init__(self, spec: str, kernel: DecodeImpl): super().__init__() self.spec = spec self.kernel = kernel diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index 66e5777848cf..250cde046eb5 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -48,11 +48,12 @@ class Step: """The step's positions, as the model's forward was handed them.""" -class ModelKernel(Protocol): - """One model's decode on this rank, for every width the spec serves. +class DecodeImpl(Protocol): + """A spec's half of the contract: one model's decode on this rank, for + every width the spec serves. - The kernel owns its own layer loop, so the library has no opinion about - layers: a kernel may replace one layer, two, or the whole backbone. + The implementation owns its own layer loop, so the library has no opinion + about layers: it may replace one layer, two, or the whole backbone. """ cache_layers: tuple[str, ...] @@ -78,8 +79,8 @@ def compile(self, key: Any) -> None: ... -KernelT = TypeVar("KernelT", bound=ModelKernel) -"""What one spec's ``build`` returns. ``ModelKernel`` is only the lower bound, +KernelT = TypeVar("KernelT", bound=DecodeImpl) +"""What one spec's ``build`` returns. ``DecodeImpl`` is only the lower bound, so a model annotates its runtime with its own kernel class (``SpecializedDecode[MyKernel]``) and reads what that class adds.""" From 2db600e0de9864de454ad189b32cd2074688121b Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 17:15:21 -0500 Subject: [PATCH 38/79] fix: remove the kernel_layesr check Signed-off-by: simondanielsson --- .../deepseek_v41/amd/specialized_decode.py | 33 +++++++++---------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index 2f5fced006b2..ed666aed3372 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -123,8 +123,11 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): self.plans: dict[int, LayerPlan] = {} tp_size = vllm_config.parallel_config.tensor_parallel_size - for layer in self._kernel_layers(model): + for idx in range(self.embed_layer, self.layer_end): + layer = model.layers[idx] plan = self._plan_layer(layer, tp_size) + if plan is None: + continue self.plans[layer.attn.layer_id] = plan if plan.ffn_only: continue @@ -139,20 +142,17 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): sorted({n for p in self.plans.values() for n in p.cache_layers}) ) - def _kernel_layers(self, model) -> list["DeepseekV4DecoderLayer"]: - end = self.layer_end - layers = [] - for i in range(self.embed_layer, end): - why = _skip_layer(model.layers[i]) - if why is None: - layers.append(model.layers[i]) - else: - logger.info_once( - "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", i, why - ) - return layers - - def _plan_layer(self, layer, tp_size: int) -> LayerPlan: + def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: + """This layer's launch, or None when it keeps vLLM's path throughout.""" + attn = layer.attn + why = _skip_layer(layer) + if why is not None: + logger.info_once( + "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", + attn.layer_id, + why, + ) + return None why = _ffn_only(layer) weights = _weights(layer, why is not None, tp_size) if why is not None: @@ -161,9 +161,8 @@ def _plan_layer(self, layer, tp_size: int) -> LayerPlan: ) # A decode step hands wo_b's TP partial to the FFN launch; the layer # reduces it on the steps that keep vLLM's path. - layer.attn.wo_b.reduce_results = False + attn.wo_b.reduce_results = False return LayerPlan(layer, weights, ffn_only=True) - attn = layer.attn comp_key = attn.compressed_cache_prefix # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. From 09b5ce2bcda9525109bdff455366606a2f815475 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 17:15:41 -0500 Subject: [PATCH 39/79] remove already made checks from skip_step Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/specialized_decode.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index ed666aed3372..0f9a17f9d6d2 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -10,7 +10,6 @@ import torch import vllm.envs as envs -from vllm.config import CUDAGraphMode from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature, Step @@ -458,11 +457,6 @@ def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: md = fc.attn_metadata if not decode.plans: return "no layer takes a launch" - if not 1 <= width <= MAX_ROWS: - return f"{width} rows" - if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): - # a PIECEWISE (breakable) capture is replayed for mixed batches too - return f"{fc.cudagraph_runtime_mode} step" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches return "no attention metadata" if any(md.get(k) is None for k in decode.comp_keys): From 9eb416a2cec5ba241d52eb6b3a2ddc03dde16d47 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 17:20:24 -0500 Subject: [PATCH 40/79] chore: cleanup Signed-off-by: simondanielsson --- tests/models/test_dsv41_specialized_decode.py | 2 +- vllm/models/deepseek_v41/amd/model.py | 4 +- .../deepseek_v41/amd/specialized_decode.py | 38 +++++-------------- 3 files changed, 13 insertions(+), 31 deletions(-) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 03257136382c..cb49bf78c8ef 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -18,7 +18,7 @@ from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Step from vllm.models.deepseek_v41.amd import specialized_decode as sd -SPECS = sd.DECODE_SPECS +SPECS = (sd.DECODE_SPEC,) LAYERS = 6 HIDDEN = 16 BLOCKS = 256 diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 7b8b80995c3d..d9ccfff53e6c 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -60,7 +60,7 @@ DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention -from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPECS, DSV41Kernel +from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPEC, DSV41Kernel from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors @@ -1386,7 +1386,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self.specialized_decode_spec = select_decode_spec(vllm_config, DECODE_SPECS) + self.specialized_decode_spec = select_decode_spec(vllm_config, DECODE_SPEC) self.specialized_decode: SpecializedDecode[DSV41Kernel] | None = None expert_dtype = getattr(config, "expert_dtype", "fp4") self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index 0f9a17f9d6d2..b28379b64804 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -80,7 +80,6 @@ class LayerPlan: layer: "DeepseekV4DecoderLayer" weights: Any ffn_only: bool - cache_layers: tuple[str, ...] = () ratio: int = 0 index: bool = False swa_key: str | None = None @@ -114,11 +113,9 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): # ``run`` owns, so no launch is built for them. end = getattr(model, "end_layer", len(model.layers)) self.layer_end = getattr(model, "decoder_replay_start", end) - self.words = torch.zeros(world, dtype=torch.int32, device=device) self.widths = widths self.swa_keys: list[str] = [] self.comp_keys: list[str] = [] - self.need_swa = False self.plans: dict[int, LayerPlan] = {} tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -130,16 +127,13 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): self.plans[layer.attn.layer_id] = plan if plan.ffn_only: continue - self.need_swa = True for key, keys in ( (plan.swa_key, self.swa_keys), (plan.comp_key, self.comp_keys), ): if key is not None and key not in keys: keys.append(key) - self.cache_layers = tuple( - sorted({n for p in self.plans.values() for n in p.cache_layers}) - ) + self.cache_layers = tuple(sorted(set(self.swa_keys) | set(self.comp_keys))) def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: """This layer's launch, or None when it keeps vLLM's path throughout.""" @@ -170,8 +164,6 @@ def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: layer, weights, ffn_only=False, - cache_layers=(attn.swa_cache_layer.prefix,) - + ((comp_key,) if comp_key is not None else ()), ratio=attn.compress_ratio, index=attn.indexer is not None, swa_key=attn.swa_cache_layer.prefix, @@ -295,13 +287,12 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: for idx in range(self.embed_layer, self.layer_end): layer = model.layers[idx] plan = self.plans.get(layer.attn.layer_id) - ffn = ( - partial(self.launch_ffn, plan) - if plan is not None and plan.ffn_only - else None - ) - out = None - if plan is not None and not plan.ffn_only: + ffn = out = None + if plan is None: + pass + elif plan.ffn_only: + ffn = partial(self.launch_ffn, plan) + else: out = self.launch_whole( plan, hidden, @@ -410,7 +401,7 @@ def _refuse_model(model) -> list[str]: def _skip_layer(layer) -> str | None: """Why this layer keeps vLLM's path for every step, or None when it takes a - launch (whole or FFN only: ``_build`` picks which).""" + launch (whole or FFN only: ``_plan_layer`` picks which).""" ffn = layer.ffn if layer.use_sequence_parallel: return "sequence-parallel layer" @@ -453,8 +444,7 @@ def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP group sees alike.""" - fc = get_forward_context() - md = fc.attn_metadata + md = get_forward_context().attn_metadata if not decode.plans: return "no layer takes a launch" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches @@ -469,8 +459,6 @@ def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: return "not a decode-only step" if swa.num_decode_tokens > width: return "more decode tokens than rows" - if not decode.need_swa: - continue if swa.decode_swa_width != SWA_WIDTH: return f"SWA width {swa.decode_swa_width}" if ( @@ -553,10 +541,6 @@ def _refuse_caches(vllm_config, decode: DSV41Kernel, caches) -> list[str]: return reasons -def _build(vllm_config, model, widths: tuple[int, ...]) -> DSV41Kernel: - return DSV41Kernel(vllm_config, model, widths) - - DECODE_SPEC = DecodeSpec[DSV41Kernel]( name="dsv41-cdna4", supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), @@ -565,7 +549,5 @@ def _build(vllm_config, model, widths: tuple[int, ...]) -> DSV41Kernel: refuse_model=_refuse_model, refuse_caches=_refuse_caches, skip_step=_skip_step, - build=_build, + build=DSV41Kernel, ) - -DECODE_SPECS = (DECODE_SPEC,) From 103bff2e434956303d51ce390677ce5f3687d49c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 17:27:10 -0500 Subject: [PATCH 41/79] chore: cleanup more Signed-off-by: simondanielsson --- .../deepseek_v41/amd/specialized_decode.py | 61 ++++++++----------- 1 file changed, 24 insertions(+), 37 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index b28379b64804..56ca868ab735 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's specialized decode on ROCm CDNA4 (``--monokernel``, -RFC #60904).""" +"""DeepSeek-V4.1's specialized decode.""" from dataclasses import dataclass from functools import partial @@ -88,17 +87,6 @@ class LayerPlan: class DSV41Kernel: - """DeepSeek-V4.1's decode on this rank: the runner (one a rank, its peer - memory collective over the TP group) and a launch plan per kernel layer. - - ``run`` owns the backbone's layer loop, between the model's own prologue - and epilogue: those keep the engram hashes, the sequence-parallel shard and - the final collapse on one code path. - - Built once at load and never moved: the launches and the captured graphs - hold its addresses. - """ - def __init__(self, vllm_config, model, widths: tuple[int, ...]): from vllm.models.common.specialized_decode.guards import tp_group from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer @@ -136,7 +124,7 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): self.cache_layers = tuple(sorted(set(self.swa_keys) | set(self.comp_keys))) def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: - """This layer's launch, or None when it keeps vLLM's path throughout.""" + """Info needed to run the specialized decode kernel for this layer.""" attn = layer.attn why = _skip_layer(layer) if why is not None: @@ -202,7 +190,7 @@ def launch_whole( engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: - """The whole backbone layer as K1 and K2, from the layer's own inputs.""" + """Run a decode layer as one logical kernel.""" layer, runner = plan.layer, self.runner attn = layer.attn first = residual is None @@ -289,10 +277,12 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: plan = self.plans.get(layer.attn.layer_id) ffn = out = None if plan is None: + # TODO: check if we should crash here pass elif plan.ffn_only: ffn = partial(self.launch_ffn, plan) else: + # Run entire layer out = self.launch_whole( plan, hidden, @@ -305,6 +295,7 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: engram_mask, ) if out is None: + # Run standard model implementation, with injected specialized ffn out = layer( hidden, positions, @@ -344,8 +335,23 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: ) -def _compute_units() -> int: - return current_platform.num_compute_units(torch.accelerator.current_device_index()) +def _ffn_only(layer) -> str | None: + """Why this layer's attention seam and attention stay on vLLM's path, which + leaves it the FFN launch, or None when it runs whole.""" + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + attn = layer.attn + if attn.compressor is not None and attn.indexer is None: + return "a compressor without an indexer" + if attn.indexer is not None and not isinstance( + attn.indexer, DeepseekV41RocmMxfp4Indexer + ): + return f"a {type(attn.indexer).__name__}" + if attn.compress_ratio not in (0, 1, 2): + return f"compress ratio {attn.compress_ratio}" + if attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": + return f"a {attn.kv_cache_dtype} KV cache" + return None def _widths(vllm_config) -> tuple[int, ...]: @@ -384,7 +390,7 @@ def _refuse_model(model) -> list[str]: reasons = [] if get_cdna_version() not in CDNA_VERSIONS: reasons.append(f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}") - cus = _compute_units() + cus = current_platform.num_compute_units(torch.accelerator.current_device_index()) if cus < BLOCKS: reasons.append(f"needs {BLOCKS} compute units, the GPU has {cus}") backends = set() @@ -421,25 +427,6 @@ def _skip_layer(layer) -> str | None: return None -def _ffn_only(layer) -> str | None: - """Why this layer's attention seam and attention stay on vLLM's path, which - leaves it the FFN launch, or None when it runs whole.""" - from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer - - attn = layer.attn - if attn.compressor is not None and attn.indexer is None: - return "a compressor without an indexer" - if attn.indexer is not None and not isinstance( - attn.indexer, DeepseekV41RocmMxfp4Indexer - ): - return f"a {type(attn.indexer).__name__}" - if attn.compress_ratio not in (0, 1, 2): - return f"compress ratio {attn.compress_ratio}" - if attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": - return f"a {attn.kv_cache_dtype} KV cache" - return None - - def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: """Why this step keeps vLLM's path at every layer, or None when the kernels take it. One verdict for every kernel layer, from the batch the whole TP From 91228ba72880b4d969cdc8b7d8431899550976d3 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 21:13:38 -0500 Subject: [PATCH 42/79] fix: remove generic from the runtime Signed-off-by: simondanielsson --- vllm/models/common/specialized_decode/runtime.py | 13 ++++++------- vllm/models/common/specialized_decode/spec.py | 6 +++--- vllm/models/deepseek_v41/amd/model.py | 4 ++-- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/vllm/models/common/specialized_decode/runtime.py b/vllm/models/common/specialized_decode/runtime.py index 6733376e99be..57f5883eeec8 100644 --- a/vllm/models/common/specialized_decode/runtime.py +++ b/vllm/models/common/specialized_decode/runtime.py @@ -5,7 +5,7 @@ """ from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Generic +from typing import TYPE_CHECKING, Any import torch @@ -22,7 +22,6 @@ DecodeImpl, DecodeSpec, DecodeUnsupported, - KernelT, Step, ) @@ -53,16 +52,16 @@ def compile(self, compile_key: "DecodeWarmup.CompileKey") -> None: self._compiled_cache[compile_key] = True -class SpecializedDecode(Generic[KernelT]): +class SpecializedDecode: """A model's specialized decode on this rank. The runtime reads the spec and builds the kernel, registers it for warm-up, checks the caches whenever the runner binds them, and decides each step - between the kernel and vLLM's own forward. The kernel it builds is of type - KernelT and is held in the ``kernel`` attribute. + between the kernel and vLLM's own forward. The kernel it builds is held in + the ``kernel`` attribute. """ - def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): + def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[Any]): """Builds this rank's kernel. Collective over the TP group. Args: @@ -86,7 +85,7 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): self.widths = spec.widths(vllm_config) self.last_reason: str | None = None self._bound = False - self.kernel: KernelT = spec.build(vllm_config, model, self.widths) + self.kernel: DecodeImpl = spec.build(vllm_config, model, self.widths) if self.kernel.warmup_keys(): DecodeWarmup(spec.name, self.kernel).register_warmup() logger.info("Specialized decode %s: widths %s", spec.name, list(self.widths)) diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index 250cde046eb5..53501fea70c5 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -80,9 +80,9 @@ def compile(self, key: Any) -> None: KernelT = TypeVar("KernelT", bound=DecodeImpl) -"""What one spec's ``build`` returns. ``DecodeImpl`` is only the lower bound, -so a model annotates its runtime with its own kernel class -(``SpecializedDecode[MyKernel]``) and reads what that class adds.""" +"""What one spec's ``build`` returns, and what its ``refuse_caches`` and +``skip_step`` are handed: a spec declared ``DecodeSpec[MyKernel]`` ties the +three together.""" @dataclass(frozen=True) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index d9ccfff53e6c..c176090fa41e 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -60,7 +60,7 @@ DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention -from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPEC, DSV41Kernel +from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPEC from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors @@ -1387,7 +1387,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.vllm_config = vllm_config self.specialized_decode_spec = select_decode_spec(vllm_config, DECODE_SPEC) - self.specialized_decode: SpecializedDecode[DSV41Kernel] | None = None + self.specialized_decode: SpecializedDecode | None = None expert_dtype = getattr(config, "expert_dtype", "fp4") self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype) From b3c84484bf794e731c33c850eadee6bd9679ff74 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 21:22:35 -0500 Subject: [PATCH 43/79] fix: remove last_reason attribute Signed-off-by: simondanielsson --- tests/models/common/specialized_decode/kit.py | 5 ----- tests/models/test_dsv41_specialized_decode.py | 16 +++++----------- vllm/models/common/specialized_decode/runtime.py | 6 ++---- 3 files changed, 7 insertions(+), 20 deletions(-) diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py index c39bc0760433..b377685bdff0 100644 --- a/tests/models/common/specialized_decode/kit.py +++ b/tests/models/common/specialized_decode/kit.py @@ -175,29 +175,24 @@ def test_every_rank_decides_alike(self, config, model): def test_skipped_step_runs_vllm_layers(self, config, model, decode): with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): assert decode.plan(self._positions(decode)) is None - assert decode.last_reason is not None def test_pipeline_stage_input_runs_vllm_layers(self, config, model, decode): with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): assert decode.plan(self._positions(decode), object()) is None - assert decode.last_reason == "pipeline stage input" def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, decode): off = torch.zeros(max(decode.widths) + 1, dtype=torch.long) with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): assert decode.plan(off) is None - assert decode.last_reason == "width" def test_piecewise_runs_vllm_layers(self, config, model, decode): meta = self.decode_metadata(config, model) with forward(config, CUDAGraphMode.PIECEWISE, meta): assert decode.plan(self._positions(decode)) is None - assert decode.last_reason == "piecewise" def test_unbound_caches_run_vllm_layers(self, config, model): decode = self.make_decode(config, model, bind=False) assert self._plan(config, model, decode) is None - assert decode.last_reason == "caches not bound" def test_caches_checked_on_every_bind(self, config, model, decode): checks: list = [] diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index cb49bf78c8ef..95667ed371a9 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -393,10 +393,9 @@ def test_moe_routing_outside_the_kernels_keeps_vllm_path(monkeypatch): assert not decode.kernel.plans -def _step(decode, model, metadata, width=1): +def _step(decode, metadata, width=1): with forward(decode.vllm_config, CUDAGraphMode.FULL, metadata): - decode.plan(torch.zeros(width, dtype=torch.long)) - return decode.last_reason + return sd._skip_step(decode.vllm_config, decode.kernel, width) @pytest.mark.parametrize( @@ -412,22 +411,17 @@ def _step(decode, model, metadata, width=1): ) def test_other_steps_take_vllm_path(monkeypatch, case, why): decode, model = _build(monkeypatch) - assert why in _step(decode, model, _metadata(model, **case)) + assert why in _step(decode, _metadata(model, **case)) def test_decode_step_takes_the_kernels(monkeypatch): decode, model = _build(monkeypatch) - assert _step(decode, model, _metadata(model)) is None + assert _step(decode, _metadata(model)) is None def test_profile_run_takes_vllm_path(monkeypatch): decode, model = _build(monkeypatch) - assert _step(decode, model, None) == "no attention metadata" - - -def test_rows_past_the_kernels_take_vllm_path(monkeypatch): - decode, model = _build(monkeypatch) - assert _step(decode, model, _metadata(model), width=sd.MAX_ROWS + 1) == "width" + assert _step(decode, None) == "no attention metadata" def test_widths_follow_the_capture_sizes(monkeypatch): diff --git a/vllm/models/common/specialized_decode/runtime.py b/vllm/models/common/specialized_decode/runtime.py index 57f5883eeec8..cda46a15acbc 100644 --- a/vllm/models/common/specialized_decode/runtime.py +++ b/vllm/models/common/specialized_decode/runtime.py @@ -83,7 +83,6 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[Any]): raise DecodeUnsupported(f"{refused}: another TP rank refused") self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) - self.last_reason: str | None = None self._bound = False self.kernel: DecodeImpl = spec.build(vllm_config, model, self.widths) if self.kernel.warmup_keys(): @@ -114,8 +113,8 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non Steps that run the model's forward: a pipeline stage's inputs, a width off the spec's widths, a PIECEWISE step (keyed by token count alone, so - it is replayed for mixed batches too), and whatever ``skip_step`` - refuses. ``last_reason`` holds why. + it is replayed for mixed batches too), caches the runner has not bound, + and whatever ``skip_step`` refuses. Every predicate is the same on every TP rank: the kernel waits in-kernel on its peers, so the ranks must take the same path. @@ -134,7 +133,6 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non why = "caches not bound" else: why = self.spec.skip_step(self.vllm_config, self.kernel, width) - self.last_reason = why if why is not None: logger.info_once( "Specialized decode %s: step on vLLM's path (%s)", self.spec.name, why From 5865b069dcc550c3a6a64f4f1c8c4af5c0b73da4 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 21:25:26 -0500 Subject: [PATCH 44/79] chore: rename runtime to dispatch Signed-off-by: simondanielsson --- tests/models/common/specialized_decode/kit.py | 8 ++++---- tests/models/common/specialized_decode/test_reference.py | 2 +- tests/models/test_dsv41_specialized_decode.py | 2 +- vllm/models/common/specialized_decode/__init__.py | 2 +- .../common/specialized_decode/{runtime.py => dispatch.py} | 2 +- 5 files changed, 8 insertions(+), 8 deletions(-) rename vllm/models/common/specialized_decode/{runtime.py => dispatch.py} (98%) diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py index b377685bdff0..c51f4baac290 100644 --- a/tests/models/common/specialized_decode/kit.py +++ b/tests/models/common/specialized_decode/kit.py @@ -16,13 +16,13 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry +from vllm.models.common.specialized_decode.dispatch import SpecializedDecode from vllm.models.common.specialized_decode.guards import ( _refuse_config, _select_spec, _set_capture_sizes, select_decode_spec, ) -from vllm.models.common.specialized_decode.runtime import SpecializedDecode from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Feature KERNEL_BACKENDS = ("aiter", "flydsl") @@ -98,9 +98,9 @@ def decode(self, config, model): return self.make_decode(config, model) def make_decode(self, config, model=None, bind: bool = True): - """This rank's runtime, built the way a model builds it: the variant - picked once by ``select_decode_spec``, then the runtime over the loaded model, - then the caches the runner binds.""" + """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 diff --git a/tests/models/common/specialized_decode/test_reference.py b/tests/models/common/specialized_decode/test_reference.py index e05a1d7a27b9..e2534d76c59b 100644 --- a/tests/models/common/specialized_decode/test_reference.py +++ b/tests/models/common/specialized_decode/test_reference.py @@ -8,11 +8,11 @@ from tests.models.common.specialized_decode import reference from tests.models.common.specialized_decode.kit import DecodeConformance, forward from vllm.config.compilation import CUDAGraphMode +from vllm.models.common.specialized_decode.dispatch import SpecializedDecode from vllm.models.common.specialized_decode.guards import ( _select_spec, select_decode_spec, ) -from vllm.models.common.specialized_decode.runtime import SpecializedDecode from vllm.models.common.specialized_decode.spec import DecodeUnsupported SPECS = reference.DECODE_SPECS diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 95667ed371a9..98d3c239b2c1 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -13,8 +13,8 @@ from tests.models.common.specialized_decode.kit import DecodeConformance, forward from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.models.common.specialized_decode.dispatch import SpecializedDecode from vllm.models.common.specialized_decode.guards import select_decode_spec -from vllm.models.common.specialized_decode.runtime import SpecializedDecode from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Step from vllm.models.deepseek_v41.amd import specialized_decode as sd diff --git a/vllm/models/common/specialized_decode/__init__.py b/vllm/models/common/specialized_decode/__init__.py index a8d333824080..6a118ebf1b47 100644 --- a/vllm/models/common/specialized_decode/__init__.py +++ b/vllm/models/common/specialized_decode/__init__.py @@ -3,8 +3,8 @@ """Specialized decode (RFC #60904, ``--monokernel``): a model's decode step in one launch.""" +from vllm.models.common.specialized_decode.dispatch import SpecializedDecode from vllm.models.common.specialized_decode.guards import select_decode_spec -from vllm.models.common.specialized_decode.runtime import SpecializedDecode from vllm.models.common.specialized_decode.spec import DecodeSpec, Step __all__ = [ diff --git a/vllm/models/common/specialized_decode/runtime.py b/vllm/models/common/specialized_decode/dispatch.py similarity index 98% rename from vllm/models/common/specialized_decode/runtime.py rename to vllm/models/common/specialized_decode/dispatch.py index cda46a15acbc..f28fe95ef7b3 100644 --- a/vllm/models/common/specialized_decode/runtime.py +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -55,7 +55,7 @@ def compile(self, compile_key: "DecodeWarmup.CompileKey") -> None: class SpecializedDecode: """A model's specialized decode on this rank. - The runtime reads the spec and builds the kernel, registers it for warm-up, + It reads the spec and builds the kernel, registers it for warm-up, checks the caches whenever the runner binds them, and decides each step between the kernel and vLLM's own forward. The kernel it builds is held in the ``kernel`` attribute. From 4bd6edf354a0b0524efe9382e98a679f6749f058 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 21:41:16 -0500 Subject: [PATCH 45/79] feat: add --specialized-decode and deprecate --monokernel Signed-off-by: simondanielsson --- .../common/specialized_decode/reference.py | 2 +- .../specialized_decode/test_reference.py | 2 +- tests/models/test_dsv41_mono_layers.py | 2 +- tests/models/test_dsv41_specialized_decode.py | 4 +-- vllm/config/kernel.py | 9 +++---- vllm/engine/arg_utils.py | 25 ++++++++++++++++--- .../common/specialized_decode/__init__.py | 2 +- .../common/specialized_decode/guards.py | 6 ++--- 8 files changed, 34 insertions(+), 18 deletions(-) diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py index 89bb62637849..c4471e0e5b48 100644 --- a/tests/models/common/specialized_decode/reference.py +++ b/tests/models/common/specialized_decode/reference.py @@ -159,7 +159,7 @@ def make_config(**overrides) -> Any: """ config = SimpleNamespace( - kernel_config=SimpleNamespace(monokernel=True), + kernel_config=SimpleNamespace(specialized_decode=True), model_config=SimpleNamespace(dtype=torch.bfloat16, enable_sleep_mode=False), scheduler_config=SimpleNamespace(max_num_seqs=4), parallel_config=SimpleNamespace( diff --git a/tests/models/common/specialized_decode/test_reference.py b/tests/models/common/specialized_decode/test_reference.py index e2534d76c59b..e5292318ad52 100644 --- a/tests/models/common/specialized_decode/test_reference.py +++ b/tests/models/common/specialized_decode/test_reference.py @@ -42,7 +42,7 @@ def bad_caches(self, config, model): def test_switch_off_means_no_spec(): config = reference.make_config() - config.kernel_config.monokernel = False + config.kernel_config.specialized_decode = False assert select_decode_spec(config, SPECS) is None diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index dea44b0acc74..09bb18321850 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -227,7 +227,7 @@ def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): hf_overrides=TRUNCATED, worker_extension_cls=f"{__name__}.MonoProbe", moe_backend="aiter", - monokernel=mono, + specialized_decode=mono, language_model_only=True, tokenizer_mode="deepseek_v41", attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 98d3c239b2c1..9f05f175eee4 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -195,7 +195,7 @@ def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: def make_config(**overrides) -> Any: config = SimpleNamespace( - kernel_config=SimpleNamespace(monokernel=True), + kernel_config=SimpleNamespace(specialized_decode=True), model_config=SimpleNamespace( enable_sleep_mode=False, hf_config=SimpleNamespace( @@ -479,7 +479,7 @@ def test_other_moe_backend_refuses(monkeypatch): def test_off_without_the_flag(monkeypatch): _gpu(monkeypatch) config = make_config() - config.kernel_config.monokernel = False + config.kernel_config.specialized_decode = False assert select_decode_spec(config, SPECS) is None diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 5c4b877fd795..842126014102 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -291,11 +291,10 @@ class KernelConfig: enable_jit_warmup: bool = True """If True, run JIT compile warmup during kernel warmup.""" - monokernel: bool = False - """If True, run the eligible layers of decode steps on the model's - MonoKernel: a persistent kernel that runs most or all of one decoder layer - in a single launch (DeepSeek-V4.1 on MI355X). A deployment that cannot run - it is refused at startup.""" + specialized_decode: bool = False + """If True, run decode steps on the model's specialized decode: a kernel + that runs most or all of a decoder layer in a single launch (DeepSeek-V4.1 + on MI355X). A deployment that cannot run it is refused at startup.""" moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ccf66f15b9dc..82503370dccd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -470,6 +470,12 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) +_MONOKERNEL_DEPRECATION_MESSAGE = ( + "--monokernel is deprecated and will be removed in v0.33.0. " + "Use --specialized-decode instead." +) + + class DeprecatedLogConfigFileAction(argparse.Action): """Warn when the legacy ``--log-config-file`` option is used.""" @@ -769,7 +775,8 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) - monokernel: bool = get_field(KernelConfig, "monokernel") + specialized_decode: bool = get_field(KernelConfig, "specialized_decode") + monokernel: bool | None = None worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1823,7 +1830,15 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) - kernel_group.add_argument("--monokernel", **kernel_kwargs["monokernel"]) + kernel_group.add_argument( + "--specialized-decode", **kernel_kwargs["specialized_decode"] + ) + kernel_group.add_argument( + "--monokernel", + action=argparse.BooleanOptionalAction, + default=None, + help=_MONOKERNEL_DEPRECATION_MESSAGE, + ) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -2777,8 +2792,10 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune - if self.monokernel: - kernel_config.monokernel = True + if self.monokernel is not None: + warnings.warn(_MONOKERNEL_DEPRECATION_MESSAGE, stacklevel=1) + if self.specialized_decode or self.monokernel: + kernel_config.specialized_decode = True if self.moe_backend != "auto": kernel_config.moe_backend = self.moe_backend if self.linear_backend != "auto": diff --git a/vllm/models/common/specialized_decode/__init__.py b/vllm/models/common/specialized_decode/__init__.py index 6a118ebf1b47..4b64dea3ecec 100644 --- a/vllm/models/common/specialized_decode/__init__.py +++ b/vllm/models/common/specialized_decode/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Specialized decode (RFC #60904, ``--monokernel``): a model's decode step +"""Specialized decode (RFC #60904, ``--specialized-decode``): a model's decode step in one launch.""" from vllm.models.common.specialized_decode.dispatch import SpecializedDecode diff --git a/vllm/models/common/specialized_decode/guards.py b/vllm/models/common/specialized_decode/guards.py index 6b20d4e54acc..d34af18e3232 100644 --- a/vllm/models/common/specialized_decode/guards.py +++ b/vllm/models/common/specialized_decode/guards.py @@ -65,8 +65,8 @@ def tp_agree(ok: bool, group) -> bool: def _enabled(vllm_config: "VllmConfig") -> bool: - """Whether ``--monokernel`` asked for the model's specialized decode.""" - return bool(vllm_config.kernel_config.monokernel) + """Whether ``--specialized-decode`` asked for the model's specialized decode.""" + return bool(vllm_config.kernel_config.specialized_decode) def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: @@ -182,7 +182,7 @@ def select_decode_spec( specs: The model's variants, in the order they are preferred. Returns: - The variant this config runs, or None with ``--monokernel`` off. + The variant this config runs, or None with ``--specialized-decode`` off. Raises: DecodeUnsupported: Every variant refused this config. From 9a80c3a25fbabd315239ba19ec3d4cb02b33f178 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 21:52:52 -0500 Subject: [PATCH 46/79] fix: add bind_caches hook Signed-off-by: simondanielsson --- tests/models/common/specialized_decode/kit.py | 23 ++++++++++++++++--- .../common/specialized_decode/reference.py | 4 ++++ .../common/specialized_decode/dispatch.py | 5 +++- vllm/models/common/specialized_decode/spec.py | 12 ++++++++++ .../deepseek_v41/amd/specialized_decode.py | 4 ++++ 5 files changed, 44 insertions(+), 4 deletions(-) diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py index c51f4baac290..59d614051aaf 100644 --- a/tests/models/common/specialized_decode/kit.py +++ b/tests/models/common/specialized_decode/kit.py @@ -54,6 +54,19 @@ def recorded(*args): return replace(spec, **{field: recorded}) +def _record(obj, method: str) -> list: + """``obj.method``'s calls from here on, as they are made.""" + calls: list = [] + inner = getattr(obj, method) + + def recorded(*args, **kwargs): + calls.append(args) + return inner(*args, **kwargs) + + setattr(obj, method, recorded) + return calls + + class DecodeConformance: """What every ``DecodeSpec`` must satisfy. Subclass it in the model's test module and fill in the hooks below.""" @@ -194,16 +207,18 @@ def test_unbound_caches_run_vllm_layers(self, config, model): decode = self.make_decode(config, model, bind=False) assert self._plan(config, model, decode) is None - def test_caches_checked_on_every_bind(self, config, model, decode): + def test_caches_checked_and_handed_over_on_every_bind(self, config, model, decode): checks: list = [] + taken = _record(decode.kernel, "bind_caches") decode.spec = _spy(decode.spec, "refuse_caches", checks) decode.bind_caches() - assert len(checks) == 1 + assert len(checks) == len(taken) == 1 + assert set(taken[0][0]) == set(decode.kernel.cache_layers) sfc = config.compilation_config.static_forward_context for name in decode.kernel.cache_layers: sfc[name].kv_cache = sfc[name].kv_cache.clone() decode.bind_caches() - assert len(checks) == 2 + assert len(checks) == len(taken) == 2 def test_refuse_caches_reports_every_reason(self, config, model, decode): reasons = decode.spec.refuse_caches( @@ -213,11 +228,13 @@ def test_refuse_caches_reports_every_reason(self, config, model, decode): def test_mismatched_caches_refuse_the_bind(self, config, model, decode): sfc = config.compilation_config.static_forward_context + taken = _record(decode.kernel, "bind_caches") bad = self.bad_caches(config, model) for name, cache in zip(decode.kernel.cache_layers, bad.values(), strict=False): sfc[name].kv_cache = cache with pytest.raises(DecodeUnsupported): decode.bind_caches() + assert not taken def test_build_does_not_move_weights(self, config): model = self.make_model(config) diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py index c4471e0e5b48..1181a0536317 100644 --- a/tests/models/common/specialized_decode/reference.py +++ b/tests/models/common/specialized_decode/reference.py @@ -71,6 +71,7 @@ def __init__(self, vllm_config, model: ReferenceModel, widths: tuple[int, ...]): self.weights = {layer.index: layer.weight.data for layer in served} self.cache_layers = tuple(layer.attn_name for layer in served) self.launched: list[int] = [] + self.bound: dict[str, Any] = {} def run(self, step: Step, input_ids, inputs_embeds, **kwargs): self.launched.append(step.width) @@ -83,6 +84,9 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs): hidden = hidden + weight return hidden + def bind_caches(self, caches) -> None: + self.bound = dict(caches) + def warmup_keys(self) -> list[Any]: return [(i, w) for i in self.weights for w in self.widths] diff --git a/vllm/models/common/specialized_decode/dispatch.py b/vllm/models/common/specialized_decode/dispatch.py index f28fe95ef7b3..8abe6bea1255 100644 --- a/vllm/models/common/specialized_decode/dispatch.py +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -90,10 +90,12 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[Any]): logger.info("Specialized decode %s: widths %s", spec.name, list(self.widths)) def bind_caches(self) -> None: - """Checks the caches the kernel will read, as the runner binds them. + """Checks the caches the kernel will read and hands them to it, as the + runner binds them. Called from ``ModelState.initialize_kv_cache``, which runs once for the throwaway cache memory profiling allocates and again for the real one. + The kernel only ever sees caches this spec has accepted. Raises: DecodeUnsupported: The caches are not what this spec serves. @@ -106,6 +108,7 @@ def bind_caches(self) -> None: raise DecodeUnsupported( f"Specialized decode {self.spec.name} caches: " + "; ".join(reasons) ) + self.kernel.bind_caches(caches) self._bound = True def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | None: diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index 53501fea70c5..3a165b1f137e 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -69,6 +69,18 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: """ ... + def bind_caches(self, caches: Mapping[str, Any]) -> None: + """Takes the ``cache_layers`` caches the runner has bound, once + ``DecodeSpec.refuse_caches`` has passed them. + + For a kernel whose launches close over the cache layout (a compile + against its strides, a cached view), this is where that happens. Called + again whenever the runner rebinds, so it must replace what the last call + built, not add to it: a normal startup binds the throwaway cache memory + profiling allocates before the real one. + """ + ... + def warmup_keys(self) -> list[Any]: """Hashable build keys of this kernel's launches, compiled before CUDA graph capture by vLLM's JIT warm-up.""" diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index 56ca868ab735..5738d3d9ff67 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek-V4.1's specialized decode.""" +from collections.abc import Mapping from dataclasses import dataclass from functools import partial from typing import TYPE_CHECKING, Any, cast @@ -159,6 +160,9 @@ def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: folds=(True, False) if settled else (True,), ) + def bind_caches(self, caches: Mapping[str, Any]) -> None: + """Nothing to take: a launch reads the cache off its attention layer.""" + def warmup_keys(self) -> list[Any]: keys: dict[Any, None] = {} for plan in self.plans.values(): From 5bbeea49ccd96f183b935a876b9b01561db4de7e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 22:07:05 -0500 Subject: [PATCH 47/79] fix: remove the common code tests for now Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 6 +- tests/models/common/__init__.py | 2 - .../common/specialized_decode/__init__.py | 2 - tests/models/common/specialized_decode/kit.py | 267 ------------------ .../common/specialized_decode/reference.py | 220 --------------- .../specialized_decode/test_reference.py | 103 ------- 6 files changed, 2 insertions(+), 598 deletions(-) delete mode 100644 tests/models/common/__init__.py delete mode 100644 tests/models/common/specialized_decode/__init__.py delete mode 100644 tests/models/common/specialized_decode/kit.py delete mode 100644 tests/models/common/specialized_decode/reference.py delete mode 100644 tests/models/common/specialized_decode/test_reference.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index cd417098bfad..04ca2e1bbc49 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -71,9 +71,8 @@ steps: - tests/models/test_deepseek_v41_replay_batch.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py - - 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/common/specialized_decode + - 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 mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -97,9 +96,8 @@ steps: - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_specialized_decode.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_specialized_decode.py models/common/specialized_decode + - 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 - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/common/__init__.py b/tests/models/common/__init__.py deleted file mode 100644 index 208f01a7cb5e..000000000000 --- a/tests/models/common/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/common/specialized_decode/__init__.py b/tests/models/common/specialized_decode/__init__.py deleted file mode 100644 index 208f01a7cb5e..000000000000 --- a/tests/models/common/specialized_decode/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py deleted file mode 100644 index 59d614051aaf..000000000000 --- a/tests/models/common/specialized_decode/kit.py +++ /dev/null @@ -1,267 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The specialized decode conformance kit (RFC #60904): what a ``DecodeSpec`` -must satisfy.""" - -import subprocess -import sys -from collections.abc import Iterator -from contextlib import contextmanager -from dataclasses import replace -from typing import Any - -import pytest -import torch - -from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import ForwardContext, override_forward_context -from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry -from vllm.models.common.specialized_decode.dispatch import SpecializedDecode -from vllm.models.common.specialized_decode.guards import ( - _refuse_config, - _select_spec, - _set_capture_sizes, - select_decode_spec, -) -from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Feature - -KERNEL_BACKENDS = ("aiter", "flydsl") -"""Backends a spec module must not pull in: it is imported at config time, on -every platform, before anything has decided the spec can run.""" - - -@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 - - -def _spy(spec, field: str, calls: list): - """``spec`` with one callback recording its calls.""" - inner = getattr(spec, field) - - def recorded(*args): - calls.append(args) - return inner(*args) - - return replace(spec, **{field: recorded}) - - -def _record(obj, method: str) -> list: - """``obj.method``'s calls from here on, as they are made.""" - calls: list = [] - inner = getattr(obj, method) - - def recorded(*args, **kwargs): - calls.append(args) - return inner(*args, **kwargs) - - setattr(obj, method, recorded) - return calls - - -class DecodeConformance: - """What every ``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 decode_metadata(self, config, model) -> Any: - """Attention metadata for a decode step the kernels take.""" - raise NotImplementedError - - def skipped_metadata(self, config, model) -> Any: - """Attention metadata for a step ``skip_step`` sends to vLLM's path.""" - 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 spec(self, config): - spec, refusals = _select_spec(self.specs, config) - assert spec is not None, refusals - return spec - - def test_refusals_report_every_reason(self, config): - config.lora_config = object() - config.speculative_config = object() - config.use_v2_model_runner = False - for spec in self.specs: - reasons = _refuse_config(spec, config) - assert sum("Model Runner V2" in r for r in reasons) == 1, reasons - for feature in (Feature.LORA, Feature.SPECULATIVE_DECODING): - want = 0 if feature in spec.supports else 1 - assert sum(feature.value in r for r in reasons) == want, reasons - - def test_refusal_names_every_variant(self, config): - config.use_v2_model_runner = False - with pytest.raises(DecodeUnsupported) as excinfo: - select_decode_spec(config, self.specs) - for spec in self.specs: - assert spec.name in str(excinfo.value) - - def test_spec_module_imports_no_kernel_backend(self): - module = self.specs_module - source = ( - "import sys; sys.path.insert(0, '.'); import importlib;" - f"importlib.import_module({module!r});" - "print('\\n'.join(sys.modules))" - ) - 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_capture_sizes_are_the_widths(self, config): - spec = self.spec(config) - widths = spec.widths(config) - config.compilation_config.cudagraph_capture_sizes = list( - range(1, max(widths) * 2 + 1) - ) - _set_capture_sizes(spec, config) - sizes = config.compilation_config.cudagraph_capture_sizes - assert [s for s in sizes if s <= max(widths)] == sorted(widths) - - def test_decided_once_a_step(self, config, model, decode): - asked: list = [] - decode.spec = _spy(decode.spec, "skip_step", asked) - step = self._plan(config, model, decode) - assert step is not None and step.width == min(decode.widths) - assert len(asked) == 1 - - def test_every_rank_decides_alike(self, config, model): - decisions = [] - for _ in range(2): - other = self.make_decode(config) - decisions.append(self._plan(config, model, other)) - assert decisions[0] is not None and decisions[1] is not None - - def test_skipped_step_runs_vllm_layers(self, config, model, decode): - with forward(config, CUDAGraphMode.NONE, self.skipped_metadata(config, model)): - assert decode.plan(self._positions(decode)) is None - - def test_pipeline_stage_input_runs_vllm_layers(self, config, model, decode): - with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - assert decode.plan(self._positions(decode), object()) is None - - def test_width_off_the_capture_sizes_runs_vllm_layers(self, config, model, decode): - off = torch.zeros(max(decode.widths) + 1, dtype=torch.long) - with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - assert decode.plan(off) is None - - def test_piecewise_runs_vllm_layers(self, config, model, decode): - meta = self.decode_metadata(config, model) - with forward(config, CUDAGraphMode.PIECEWISE, meta): - assert decode.plan(self._positions(decode)) is None - - def test_unbound_caches_run_vllm_layers(self, config, model): - decode = self.make_decode(config, model, bind=False) - assert self._plan(config, model, decode) is None - - def test_caches_checked_and_handed_over_on_every_bind(self, config, model, decode): - checks: list = [] - taken = _record(decode.kernel, "bind_caches") - decode.spec = _spy(decode.spec, "refuse_caches", checks) - decode.bind_caches() - assert len(checks) == len(taken) == 1 - assert set(taken[0][0]) == set(decode.kernel.cache_layers) - sfc = config.compilation_config.static_forward_context - for name in decode.kernel.cache_layers: - sfc[name].kv_cache = sfc[name].kv_cache.clone() - decode.bind_caches() - assert len(checks) == len(taken) == 2 - - 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_mismatched_caches_refuse_the_bind(self, config, model, decode): - sfc = config.compilation_config.static_forward_context - taken = _record(decode.kernel, "bind_caches") - bad = self.bad_caches(config, model) - for name, cache in zip(decode.kernel.cache_layers, bad.values(), strict=False): - sfc[name].kv_cache = cache - with pytest.raises(DecodeUnsupported): - decode.bind_caches() - assert not taken - - 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 - - def test_build_runs_once_over_every_width(self, config, model): - calls: list = [] - spec = _spy(select_decode_spec(config, self.specs), "build", calls) - decode = SpecializedDecode(model, config, spec) - assert [args[1] for args in calls] == [model] - assert set(calls[0][2]) == set(decode.widths) - - def test_warmup_registered_for_every_key(self, config): - registry = JitWarmupRegistry(config) - with registry.activate(): - decode = self.make_decode(config) - keys = decode.kernel.warmup_keys() - assert len(keys) == len(set(keys)) >= 1 - assert len(registry) == len(keys) - registry.warmup() - - def _positions(self, decode) -> torch.Tensor: - return torch.zeros(min(decode.widths), dtype=torch.long) - - def _plan(self, config, model, decode): - with forward(config, CUDAGraphMode.NONE, self.decode_metadata(config, model)): - return decode.plan(self._positions(decode)) diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py deleted file mode 100644 index 1181a0536317..000000000000 --- a/tests/models/common/specialized_decode/reference.py +++ /dev/null @@ -1,220 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The reference kernel: a ``DecodeImpl`` over the layers it replaces.""" - -from dataclasses import replace -from types import SimpleNamespace -from typing import Any - -import torch - -from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import get_forward_context -from vllm.models.common.specialized_decode.spec import DecodeSpec, Step - -CACHE_DTYPE = torch.bfloat16 -CACHE_NDIM = 4 - - -class ReferenceLayer(torch.nn.Module): - """A decoder layer the reference kernel can replace: one weight, one - attention layer name, and ``dense`` for the layers a spec skips.""" - - def __init__(self, index: int, hidden: int, dense: bool = False): - super().__init__() - self.index = index - self.dense = dense - self.attn_name = f"layers.{index}.attn" - self.weight = torch.nn.Parameter(torch.ones(hidden) * index) - - def forward(self, positions: torch.Tensor, hidden: torch.Tensor): - return hidden + self.weight - - -class ReferenceModel(torch.nn.Module): - """A model shaped like the layer loop a kernel hooks into.""" - - def __init__(self, layers: int = 4, hidden: int = 8): - super().__init__() - self.config = SimpleNamespace(num_hidden_layers=layers, hidden_size=hidden) - self.layers = torch.nn.ModuleList( - [ReferenceLayer(i, hidden, dense=i == 0) for i in range(layers)] - ) - self.start_layer = 0 - self.end_layer = layers - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors=None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: - hidden = inputs_embeds - for layer in self.layers: - hidden = layer(positions, hidden) - return hidden - - -class ReferenceKernel: - """The model's decode in one launch: the work of every layer it replaces, - off views of the model's weights, with the dense layers left to vLLM.""" - - def __init__(self, vllm_config, model: ReferenceModel, widths: tuple[int, ...]): - self.model = model - self.widths = widths - served = [ - model.layers[i] - for i in range(model.start_layer, model.end_layer) - if not model.layers[i].dense - ] - self.weights = {layer.index: layer.weight.data for layer in served} - self.cache_layers = tuple(layer.attn_name for layer in served) - self.launched: list[int] = [] - self.bound: dict[str, Any] = {} - - def run(self, step: Step, input_ids, inputs_embeds, **kwargs): - self.launched.append(step.width) - hidden = inputs_embeds - for layer in self.model.layers: - weight = self.weights.get(layer.index) - if weight is None: - hidden = layer(step.positions, hidden) - else: - hidden = hidden + weight - return hidden - - def bind_caches(self, caches) -> None: - self.bound = dict(caches) - - def warmup_keys(self) -> list[Any]: - return [(i, w) for i in self.weights for w in self.widths] - - def compile(self, key: Any) -> None: - pass - - -def widths(vllm_config) -> tuple[int, ...]: - most = vllm_config.scheduler_config.max_num_seqs - return tuple(w for w in (1, 2, 4, 8) if w <= most) - - -def refuse_config(vllm_config) -> list[str]: - reasons = [] - if not widths(vllm_config): - reasons.append("max_num_seqs below the smallest width") - if vllm_config.model_config.dtype is not torch.bfloat16: - reasons.append(f"dtype {vllm_config.model_config.dtype}: needs bfloat16") - return reasons - - -def refuse_model(model) -> list[str]: - return [] if isinstance(model, ReferenceModel) else ["not a ReferenceModel"] - - -def skip_step(vllm_config, kernel: ReferenceKernel, width: int) -> str | None: - md = get_forward_context().attn_metadata - return "prefill" if md.get("prefill") else None - - -def refuse_caches(vllm_config, kernel: ReferenceKernel, caches) -> list[str]: - reasons = [] - for name, cache in caches.items(): - if cache.dtype is not CACHE_DTYPE: - reasons.append(f"{name}: dtype {cache.dtype}, needs {CACHE_DTYPE}") - if cache.ndim != CACHE_NDIM: - reasons.append(f"{name}: {cache.ndim}d, needs {CACHE_NDIM}d") - return reasons - - -def build(vllm_config, model, widths: tuple[int, ...]) -> ReferenceKernel: - return ReferenceKernel(vllm_config, model, widths) - - -REFERENCE = DecodeSpec[ReferenceKernel]( - name="reference", - supports=frozenset(), - widths=widths, - refuse_config=refuse_config, - refuse_model=refuse_model, - refuse_caches=refuse_caches, - skip_step=skip_step, - build=build, -) - -REFERENCE_WIDE = replace( - REFERENCE, - name="reference-wide", - refuse_config=lambda c: ( - refuse_config(c) - + (["needs max_num_seqs >= 8"] if c.scheduler_config.max_num_seqs < 8 else []) - ), -) - -DECODE_SPECS = (REFERENCE_WIDE, REFERENCE) - - -def make_config(**overrides) -> Any: - """A config shaped like the parts of ``VllmConfig`` the library reads. - - Args: - **overrides: Top-level attributes to replace. Nested ones are set on - the returned object (``config.parallel_config.enable_eplb = True``). - - """ - config = SimpleNamespace( - kernel_config=SimpleNamespace(specialized_decode=True), - model_config=SimpleNamespace(dtype=torch.bfloat16, enable_sleep_mode=False), - scheduler_config=SimpleNamespace(max_num_seqs=4), - parallel_config=SimpleNamespace( - pipeline_parallel_size=1, - data_parallel_size=1, - enable_expert_parallel=False, - enable_eplb=False, - decode_context_parallel_size=1, - prefill_context_parallel_size=1, - ), - offload_config=SimpleNamespace( - uva=SimpleNamespace(cpu_offload_gb=0), - prefetch=SimpleNamespace(offload_group_size=0), - ), - aux_output_config=SimpleNamespace(enable_return_routed_experts=False), - compilation_config=SimpleNamespace( - cudagraph_mode=CUDAGraphMode.FULL, - cudagraph_capture_sizes=[1, 2, 3, 4, 16], - static_forward_context={}, - ), - lora_config=None, - kv_transfer_config=None, - speculative_config=None, - weight_transfer_config=None, - use_v2_model_runner=True, - ) - for key, value in overrides.items(): - setattr(config, key, value) - return config - - -def make_model(config, layers: int = 4) -> ReferenceModel: - """A loaded reference model, with its KV caches already in the config's - static forward context.""" - model = ReferenceModel(layers=layers) - set_caches(config, model) - return model - - -def set_caches(config, model, **kwargs) -> dict[str, torch.Tensor]: - """Puts one cache a layer in the static forward context, and returns them.""" - caches = { - layer.attn_name: make_cache(**kwargs) - for layer in model.layers - if not layer.dense - } - config.compilation_config.static_forward_context = { - name: SimpleNamespace(kv_cache=cache) for name, cache in caches.items() - } - return caches - - -def make_cache(dtype=torch.bfloat16, ndim: int = 4) -> torch.Tensor: - return torch.zeros((2,) + (1,) * (ndim - 1), dtype=dtype) diff --git a/tests/models/common/specialized_decode/test_reference.py b/tests/models/common/specialized_decode/test_reference.py deleted file mode 100644 index e5292318ad52..000000000000 --- a/tests/models/common/specialized_decode/test_reference.py +++ /dev/null @@ -1,103 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The conformance kit against the reference spec: the library's own test.""" - -import pytest -import torch - -from tests.models.common.specialized_decode import reference -from tests.models.common.specialized_decode.kit import DecodeConformance, forward -from vllm.config.compilation import CUDAGraphMode -from vllm.models.common.specialized_decode.dispatch import SpecializedDecode -from vllm.models.common.specialized_decode.guards import ( - _select_spec, - select_decode_spec, -) -from vllm.models.common.specialized_decode.spec import DecodeUnsupported - -SPECS = reference.DECODE_SPECS - - -class TestReferenceSpec(DecodeConformance): - specs = SPECS - specs_module = "tests.models.common.specialized_decode.reference" - - def make_config(self, **overrides): - return reference.make_config(**overrides) - - def make_model(self, config): - return reference.make_model(config) - - def decode_metadata(self, config, model): - return {"prefill": False} - - def skipped_metadata(self, config, model): - return {"prefill": True} - - def bad_caches(self, config, model): - return { - "wrong": reference.make_cache(dtype=torch.float16, ndim=2), - } - - -def test_switch_off_means_no_spec(): - config = reference.make_config() - config.kernel_config.specialized_decode = False - assert select_decode_spec(config, SPECS) is None - - -def test_the_first_servable_variant_wins(): - config = reference.make_config() - config.scheduler_config.max_num_seqs = 4 - spec, refusals = _select_spec(reference.DECODE_SPECS, config) - assert spec is reference.REFERENCE - assert refusals["reference-wide"] == ["needs max_num_seqs >= 8"] - - config.scheduler_config.max_num_seqs = 8 - spec, refusals = _select_spec(reference.DECODE_SPECS, config) - assert spec is reference.REFERENCE_WIDE - assert refusals == {} - - -def _make_decode(config, model): - decode = SpecializedDecode(model, config, select_decode_spec(config, SPECS)) - decode.bind_caches() - return decode - - -def test_skipped_layers_keep_vllm_path(): - config = reference.make_config() - model = reference.make_model(config) - decode = _make_decode(config, model) - assert set(decode.kernel.weights) == {1, 2, 3} - - -def test_kernel_layers_follow_the_pipeline_slice(): - config = reference.make_config() - model = reference.make_model(config) - model.start_layer, model.end_layer = 2, 4 - decode = _make_decode(config, model) - assert set(decode.kernel.weights) == {2, 3} - - -def test_the_kernel_runs_the_layers_it_replaces(): - config = reference.make_config() - model = reference.make_model(config) - decode = _make_decode(config, model) - width = min(decode.widths) - positions = torch.zeros(width, dtype=torch.long) - embeds = torch.zeros(width, model.config.hidden_size) - with forward(config, CUDAGraphMode.NONE, {"prefill": False}): - step = decode.plan(positions) - assert step is not None - hidden = decode.run(step, None, embeds) - assert torch.equal(hidden, model(None, positions, None, embeds)) - assert decode.kernel.launched == [width] - - -def test_model_refusal_is_raised(): - config = reference.make_config() - with pytest.raises(DecodeUnsupported, match="not a ReferenceModel"): - SpecializedDecode( - torch.nn.Linear(2, 2), config, select_decode_spec(config, SPECS) - ) From 6b3d53e10ac7c3b49b3d9798ce792f312c2f3d5e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 22:15:44 -0500 Subject: [PATCH 48/79] feat: refactor kimi-k3 on nvidia Signed-off-by: simondanielsson --- .../common/specialized_decode/dispatch.py | 9 +- .../common/specialized_decode/guards.py | 1 + vllm/models/common/specialized_decode/spec.py | 1 + vllm/models/kimi_k3/nvidia/model.py | 28 +-- vllm/models/kimi_k3/nvidia/model_state.py | 6 +- .../specialized_decode/runtime/__init__.py | 136 +------------ .../specialized_decode/runtime/decode.py | 119 +++++------ .../kimi_k3/nvidia/specialized_decode/spec.py | 189 ++++++++++++++++++ 8 files changed, 280 insertions(+), 209 deletions(-) create mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/spec.py diff --git a/vllm/models/common/specialized_decode/dispatch.py b/vllm/models/common/specialized_decode/dispatch.py index 8abe6bea1255..5e87dd49c12c 100644 --- a/vllm/models/common/specialized_decode/dispatch.py +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -5,7 +5,7 @@ """ from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Generic import torch @@ -22,6 +22,7 @@ DecodeImpl, DecodeSpec, DecodeUnsupported, + KernelT, Step, ) @@ -52,7 +53,7 @@ def compile(self, compile_key: "DecodeWarmup.CompileKey") -> None: self._compiled_cache[compile_key] = True -class SpecializedDecode: +class SpecializedDecode(Generic[KernelT]): """A model's specialized decode on this rank. It reads the spec and builds the kernel, registers it for warm-up, @@ -61,7 +62,7 @@ class SpecializedDecode: the ``kernel`` attribute. """ - def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[Any]): + def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): """Builds this rank's kernel. Collective over the TP group. Args: @@ -84,7 +85,7 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[Any]): self.vllm_config = vllm_config self.widths = spec.widths(vllm_config) self._bound = False - self.kernel: DecodeImpl = spec.build(vllm_config, model, self.widths) + self.kernel: KernelT = spec.build(vllm_config, model, self.widths) if self.kernel.warmup_keys(): DecodeWarmup(spec.name, self.kernel).register_warmup() logger.info("Specialized decode %s: widths %s", spec.name, list(self.widths)) diff --git a/vllm/models/common/specialized_decode/guards.py b/vllm/models/common/specialized_decode/guards.py index d34af18e3232..8206b982c09e 100644 --- a/vllm/models/common/specialized_decode/guards.py +++ b/vllm/models/common/specialized_decode/guards.py @@ -79,6 +79,7 @@ def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: Feature.EXPERT_PARALLEL: pc.enable_expert_parallel, Feature.EPLB: pc.enable_eplb, Feature.DECODE_CONTEXT_PARALLEL: pc.decode_context_parallel_size > 1, + Feature.UBATCHING: pc.use_ubatching, Feature.PREFILL_CONTEXT_PARALLEL: pc.prefill_context_parallel_size > 1, Feature.LORA: vllm_config.lora_config is not None, Feature.KV_CONNECTOR: vllm_config.kv_transfer_config is not None, diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index 3a165b1f137e..c73adba87018 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -23,6 +23,7 @@ class Feature(str, Enum): EXPERT_PARALLEL = "expert_parallel" EPLB = "eplb" DECODE_CONTEXT_PARALLEL = "decode_context_parallel" + UBATCHING = "ubatching" PREFILL_CONTEXT_PARALLEL = "prefill_context_parallel" LORA = "lora" KV_CONNECTOR = "kv_connector" diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index fca74b4624b7..74468c3e8878 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -97,6 +97,10 @@ sp_reduce_scatter, sp_shard, ) +from vllm.models.common.specialized_decode import ( + SpecializedDecode, + select_decode_spec, +) from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.nvidia.latent_moe_runner import ( LatentMoERunner, @@ -121,12 +125,12 @@ KimiK3MultiModalProcessor, KimiK3ProcessingInfo, ) -from .specialized_decode.runtime import specialized_decode_config +from .specialized_decode.spec import DECODE_SPEC if TYPE_CHECKING: from vllm.v1.worker.gpu.model_states.interface import ModelState - from .specialized_decode.runtime.decode import SpecializedDecode + from .specialized_decode.runtime.decode import KimiK3Kernel logger = init_logger(__name__) @@ -441,6 +445,7 @@ def __init__( torch.empty(num_experts, dtype=torch.float32) ) + self.shared_experts: KimiMLP | None if self.num_shared_experts is not None: shared_intermediate_size = moe_intermediate_size * self.num_shared_experts self.shared_experts = KimiMLP( @@ -1427,19 +1432,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.vocab_size, scale=logit_scale ) - # decode_backend=latency_specialized: a specialized decode executor over - # self.model, built by the model state once the weights are loaded. - self.specialized_decode_config = specialized_decode_config(vllm_config) - self.specialized_decode: SpecializedDecode | None = None + # --specialized-decode: a specialized decode executor over self.model, + # built by the model state once the weights are loaded. + self.decode_spec = select_decode_spec(vllm_config, DECODE_SPEC) + self.specialized_decode: SpecializedDecode[KimiK3Kernel] | None = None # The final norm is applied in compute_logits unless the specialized decode # applies it in forward. self.needs_norm_in_logits = True def build_specialized_decode(self) -> None: - from .specialized_decode.runtime.decode import SpecializedDecode - + assert self.decode_spec is not None self.specialized_decode = SpecializedDecode( - self.model, self.specialized_decode_config + self.model, self.vllm_config, self.decode_spec ) # From here on forward() returns normalized hidden states on every step: the # specialized decode's last AttnRes applies the norm, and forward() applies it @@ -1454,7 +1458,7 @@ def get_model_state_cls(self) -> "type[ModelState]": def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: if self.specialized_decode is not None: - return self.specialized_decode.embed_input_ids(input_ids) + return self.specialized_decode.kernel.embed_input_ids(input_ids) return self.model.embed_input_ids(input_ids) def make_empty_intermediate_tensors( @@ -1477,13 +1481,13 @@ def forward( # type: ignore[override] if decode is not None: step = decode.plan(positions, intermediate_tensors) if step is not None: - return decode.run(step, input_ids, inputs_embeds) + return decode.run(step, input_ids, inputs_embeds, **kwargs) hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ) if not self.needs_norm_in_logits: assert decode is not None - hidden_states = decode.normalize(hidden_states) + hidden_states = decode.kernel.normalize(hidden_states) return hidden_states @classmethod diff --git a/vllm/models/kimi_k3/nvidia/model_state.py b/vllm/models/kimi_k3/nvidia/model_state.py index d1d0e489b529..cbb3b2bba29d 100644 --- a/vllm/models/kimi_k3/nvidia/model_state.py +++ b/vllm/models/kimi_k3/nvidia/model_state.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kimi K3's model runner state: the hybrid (KDA + MLA) state, plus the lifecycle of -the specialized decode when ``decode_backend=latency_specialized`` selects it.""" +the specialized decode when ``--specialized-decode`` selects it.""" from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState @@ -16,11 +16,11 @@ def __init__(self, vllm_config, model, encoder_cache, device): if hasattr(model, "get_language_model"): model = model.get_language_model() self.language_model = model - if model.specialized_decode_config is not None: + if model.decode_spec is not None: model.build_specialized_decode() def initialize_kv_cache(self, kv_cache_config, block_tables): super().initialize_kv_cache(kv_cache_config, block_tables) decode = self.language_model.specialized_decode if decode is not None: - decode.prepare_kv() + decode.bind_caches() diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py index dccfb9a5cbe0..a5e3b8dc95cc 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py @@ -1,135 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, enabled -by ``additional_config={"decode_backend": "latency_specialized"}``.""" - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class DecodeConfig: - max_context: int - mla_core: str - dcp: int - # Tokens per request in a verify step (1 + draft tokens; 1 without speculation). - q: int - capture_sizes: tuple[int, ...] - - -def specialized_decode_config(vllm_config) -> DecodeConfig | None: - """The deployment's specialized decode settings, None for the stock decode. - Raises when the specialized decode is requested for an unsupported deployment.""" - extra = vllm_config.additional_config - backend = extra.get("decode_backend", "auto") if isinstance(extra, dict) else "auto" - if backend == "auto": - return None - if backend != "latency_specialized": - raise ValueError("Supported decode_backend values are auto/latency_specialized") - reason = unsupported(vllm_config) - if reason is not None: - raise ValueError(f"Kimi K3 specialized decode does not support {reason}") - spec = vllm_config.speculative_config - return DecodeConfig( - max_context=vllm_config.model_config.max_model_len, - mla_core=extra.get("mla_core", "auto"), - dcp=vllm_config.parallel_config.decode_context_parallel_size, - q=1 + spec.num_speculative_tokens if spec is not None else 1, - capture_sizes=tuple( - vllm_config.compilation_config.cudagraph_capture_sizes or () - ), - ) - - -def kda_dtypes(vllm_config): - """(conv, recurrent) KDA state dtypes the deployment's cache config selects.""" - from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator - - cache = vllm_config.cache_config - return MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, - cache.mamba_cache_dtype, - cache.mamba_ssm_cache_dtype, - ) - - -def unsupported(vllm_config) -> str | None: - import torch - - from vllm import envs - from vllm.config import CompilationMode, CUDAGraphMode - from vllm.platforms import current_platform - - parallel, spec = vllm_config.parallel_config, vllm_config.speculative_config - text = vllm_config.model_config.hf_text_config - moe = tuple( - getattr(text, name, None) - for name in ( - "hidden_act", - "activation_situ_beta", - "activation_situ_linear_beta", - "num_experts", - "num_experts_per_token", - "num_expert_group", - "topk_group", - "moe_router_activation_func", - ) - ) - compilation = vllm_config.compilation_config - graphs = compilation.cudagraph_mode - checks = ( - (not current_platform.is_device_capability_family(100), "this GPU (SM100/103)"), - ( - moe != ("situ", 4.0, 25.0, 896, 16, 1, 1, "sigmoid"), - "MoE configs but the one-group sigmoid top-16 SiTU MoE of Kimi K3", - ), - (parallel.tensor_parallel_size not in (8, 16), "TP other than 8/16"), - (parallel.pipeline_parallel_size != 1, "pipeline parallelism"), - (parallel.enable_expert_parallel, "expert parallelism"), - (parallel.use_ubatching, "DBO"), - ( - parallel.decode_context_parallel_size - not in (1, parallel.tensor_parallel_size), - "decode context parallel other than 1 or TP", - ), - (vllm_config.model_config.enable_sleep_mode, "sleep mode"), - (vllm_config.lora_config is not None, "LoRA"), - # The runner's ModelState hooks build the executor and prepare the KV layouts. - (not vllm_config.use_v2_model_runner, "the V1 model runner"), - # The decode loop runs in forward; a compiled forward would trace it. - (compilation.mode == CompilationMode.STOCK_TORCH_COMPILE, "torch.compile"), - # The loop runs only in FULL decode graphs; eager and PIECEWISE steps run stock. - ( - graphs is None or graphs.decode_mode() != CUDAGraphMode.FULL, - "deployments without FULL decode CUDA graphs (e.g. enforce-eager)", - ), - ( - spec is not None and spec.method != "dspark", - "speculative methods but DSpark", - ), - ( - spec is not None and spec.enable_adaptive_verification, - "adaptive verification", - ), - # KdaSplit keeps a BF16 conv state and an FP32 recurrent state in place; it - # maintains neither other cache dtypes nor ReplaySSM / RecoverSSM records. - ( - kda_dtypes(vllm_config) != (torch.bfloat16, torch.float32), - "KDA cache dtypes other than a BF16 conv state and an FP32 recurrent " - "state (mamba cache dtype)", - ), - ( - vllm_config.cache_config.use_replayssm - or vllm_config.cache_config.use_kda_recoverssm, - "ReplaySSM / RecoverSSM (--use-replayssm)", - ), - # The verify kernels take at most 8 tokens per request (KdaSplit's barriers). - ( - spec is not None and spec.num_speculative_tokens > 7, - "more than 7 speculative tokens", - ), - ( - spec is not None and envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM, - "the AttnRes aux stream", - ), - ) - return next((reason for failed, reason in checks if failed), None) +"""Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, +enabled by ``--specialized-decode``. What the deployment has to satisfy to run it +is declared in ``..spec``.""" diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py index e1563d38cbbe..794d48e9698c 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py @@ -3,17 +3,17 @@ """The specialized decode executor over the stock Kimi K3 backbone: one routing decision per step, then every layer's prepared providers in one loop.""" +from collections.abc import Mapping from functools import partial from typing import Any, NamedTuple import torch -from vllm.config import CUDAGraphMode from vllm.distributed import get_tp_group from vllm.distributed import tensor_model_parallel_all_reduce as all_reduce -from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.forward_context import get_forward_context +from vllm.models.common.specialized_decode import Step from vllm.platforms import current_platform -from vllm.utils.math_utils import round_up from ..cutedsl.attn_res.prepared import PreparedAttnRes from ..cutedsl.common.embedding import PreparedEmbedding @@ -25,11 +25,12 @@ from ..cutedsl.moe.small_batch import chain as small_batch_chain from ..cutedsl.moe.tail_runtime import MoeTailBox, PreparedTail from ..cutedsl.oproj.prepared import OprojMailbox +from ..spec import PREPARED_MAX, tokens_per_request from .prefetch import PrefetchController -from .providers import BATCHES, PREPARED_MAX, Attention, pack_front +from .providers import Attention, pack_front -class Step(NamedTuple): +class StepPlan(NamedTuple): q: int # tokens per request metadata: list # per layer slot_mapping: Any @@ -65,28 +66,26 @@ def is_uniform_verify(md, q, b): return full and not md.num_decodes and tokens == md.num_spec_decodes * q -def caches_of(attention): - cache = getattr(attention, "kv_cache", ()) - return [cache] if isinstance(cache, torch.Tensor) else list(cache) - - -class SpecializedDecode: +class KimiK3Kernel: """Built once the weights are loaded; owns every prepared provider, the mailboxes and the decode buffers. Reads the backbone, never modifies it.""" - def __init__(self, backbone, config): - self.backbone, self.buffers = backbone, {} - q = config.q - # Verify steps decode q tokens per request at M = requests * q tokens; FULL - # verify graphs round each capture size up to whole requests. - rounded = {round_up(s, q) for s in config.capture_sizes} - verify = {m for m in rounded if m <= PREPARED_MAX} - self.batches, self.q = tuple(sorted({*BATCHES, *verify})), q - self.queries = (1, q) if q > 1 else (1,) + def __init__(self, vllm_config, backbone, widths: tuple[int, ...]): + self.backbone = backbone + self.buffers: dict[str, torch.Tensor] = {} + extra = vllm_config.additional_config + # Verify steps decode q tokens per request at M = requests * q tokens; the + # widths already round each capture size up to whole requests. + self.batches, self.q = tuple(widths), tokens_per_request(vllm_config) + self.queries = (1, self.q) if self.q > 1 else (1,) # The longest sequence the deployment serves bounds the MLA split count. - self.max_context, self.mla_core = config.max_context, config.mla_core + self.max_context = vllm_config.model_config.max_model_len + self.mla_core = ( + extra.get("mla_core", "auto") if isinstance(extra, dict) else ("auto") + ) # Decode context parallel over the TP ranks: MLA KV is sequence-sharded. - self.dcp, self.dcp_box, self.mla_layers = config.dcp, None, 0 + self.dcp = vllm_config.parallel_config.decode_context_parallel_size + self.dcp_box, self.mla_layers = None, 0 group = get_tp_group() self.tp, self.rank = group.world_size, group.rank_in_group self.group = group.device_group @@ -95,6 +94,12 @@ def __init__(self, backbone, config): layers = backbone.layers[backbone.start_layer : backbone.end_layer] self.layers = layers = tuple(layers) self.build_providers() + self.cache_layers = tuple(layer_name(d.self_attn) for d in layers) + self.cache_kinds = { + name: self.attention[d.layer_idx].kind + for name, d in zip(self.cache_layers, layers) + } + self.caches: dict[int, list] = {} self.embedding = PreparedEmbedding(backbone.embed_tokens, self.tp, self.rank) self.before = { d.layer_idx: PreparedAttnRes.for_target(d, has_delta=d is not layers[0]) @@ -176,12 +181,21 @@ def build_providers(self): ) self.feed_forward[index] = moe - def prepare_kv(self): - """Compile each layer's cores for its bound KV cache layout (the profiling - cache, then the real one).""" - for layer in self.layers: - caches = caches_of(layer.self_attn) - self.attention[layer.layer_idx].prepare_layout(caches) + def bind_caches(self, caches: Mapping[str, Any]) -> None: + """Compile each layer's cores for its bound KV cache layout and keep the + layer's tensors for every launch (the profiling cache, then the real one).""" + for name, layer in zip(self.cache_layers, self.layers): + cache = caches[name] + tensors = [cache] if isinstance(cache, torch.Tensor) else list(cache) + self.caches[layer.layer_idx] = tensors + self.attention[layer.layer_idx].prepare_layout(tensors) + + def warmup_keys(self) -> list[Any]: + """None: every compile finishes in ``build_providers`` and ``bind_caches``.""" + return [] + + def compile(self, key: Any) -> None: + raise AssertionError("Kimi K3 registers no warm-up keys") def embed_input_ids(self, ids): rows = ids.shape[0] if ids.ndim == 1 and ids.is_contiguous() else 0 @@ -190,44 +204,34 @@ def embed_input_ids(self, ids): table = self.backbone.embed_tokens.weight return self.embedding.run(ids, table.new_empty((ids.shape[0], table.shape[1]))) - def plan(self, positions, intermediate_tensors): - """This step's decode plan, or None to run the stock forward instead. + def step_plan(self, width: int) -> StepPlan | None: + """Every layer's metadata for a step of ``width`` rows, or None when the + prepared providers do not serve it. - Steps that run the stock forward: - - pipeline-parallel stage inputs (``intermediate_tensors``); - - token counts without prepared providers (above 64 or not in ``batches``); - - steps outside a CUDA graph capture (profiling, warmup and the warmup run - before each capture), so every stock kernel runs before its capture; - - PIECEWISE steps, keyed by token count alone and so captured as stock; - - steps with prompt rows (prefills, chunked prefills, fresh prompts); - - speculative steps whose requests do not all verify q tokens (partial - drafts, verify rows mixed with plain decode rows). + None for steps with prompt rows (prefills, chunked prefills, fresh + prompts), and for speculative steps whose requests do not all verify q + tokens (partial drafts, verify rows mixed with plain decode rows). - Every predicate is the same on all TP ranks: the mailboxes of the specialized - decode need every rank to take the same path. + Every predicate is the same on all TP ranks: the mailboxes of the + specialized decode need every rank to take the same path. """ - b = positions.shape[0] - if intermediate_tensors is not None or b not in self.batches: - return None - if not torch.cuda.is_current_stream_capturing(): - return None - if not is_forward_context_available(): - return None context = get_forward_context() - if context.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: - return None by_layer = context.attn_metadata rows: list[Any] = list(by_layer.values()) if isinstance(by_layer, dict) else [] q = self.q if any(getattr(md, "num_spec_decodes", 0) for md in rows) else 1 - if b % q: + if width % q: return None - metadata = [decode_metadata(d.self_attn, context, q, b) for d in self.layers] + metadata = [ + decode_metadata(d.self_attn, context, q, width) for d in self.layers + ] if any(md is None for md in metadata): return None - return Step(q, metadata, context.slot_mapping) + return StepPlan(q, metadata, context.slot_mapping) - def run(self, step, input_ids, inputs_embeds): + def run(self, step: Step, input_ids, inputs_embeds, **kwargs): """Normalized hidden states (and the aux streams) of a planned step.""" + plan = self.step_plan(step.width) + assert plan is not None x = inputs_embeds if x is None: x = self.embed_input_ids(input_ids) @@ -235,12 +239,13 @@ def run(self, step, input_ids, inputs_embeds): # The stock aux streams: the AttnRes prefix entering each aux layer. aux_out = [torch.empty_like(x) for _ in self.aux_layers] aux = dict(zip(self.aux_layers, aux_out)) - b, q, slots = x.shape[0], step.q, step.slot_mapping + b, q, slots = x.shape[0], plan.q, plan.slot_mapping prefix, blocks = self.prefix[:b], self.blocks[:b] first, second = self.hidden[0][:b], self.hidden[1][:b] hidden, pre_ready = x, False - for layer, md in zip(self.layers, step.metadata): - index, caches = layer.layer_idx, caches_of(layer.self_attn) + for layer, md in zip(self.layers, plan.metadata): + index = layer.layer_idx + caches = self.caches[index] if not pre_ready: self.before[index].run(hidden, prefix, blocks, first, aux.get(index)) self.attention[index].run( diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py b/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py new file mode 100644 index 000000000000..7aff5a4aae75 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""What Kimi K3 declares to run its specialized decode on SM100/103. + +Imported at config time, on every platform, before anything has decided the +kernel can run: nothing here may pull in CuTe DSL. +""" + +from typing import TYPE_CHECKING + +import torch + +from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature +from vllm.utils.math_utils import round_up + +if TYPE_CHECKING: + from .runtime.decode import KimiK3Kernel + +BATCHES = (*range(1, 9), 16, 32, 64) +"""Decode rows the prepared providers serve whatever the deployment captures.""" + +PREPARED_MAX = BATCHES[-1] +"""The most rows a launch serves; larger steps run vLLM's path.""" + +KDA_STATE_DTYPES = (torch.bfloat16, torch.float32) +"""The (conv, recurrent) KDA states the prepared KDA cores keep in place.""" + +MOE = ("situ", 4.0, 25.0, 896, 16, 1, 1, "sigmoid") +"""The one-group sigmoid top-16 SiTU MoE the prepared experts are packed for.""" + +HIDDEN_SIZE = 7168 +"""The hidden size the decode buffers are shaped for.""" + +MAX_SPECULATIVE_TOKENS = 7 +"""The verify kernels' tokens per request, bounded by KdaSplit's barriers.""" + + +def tokens_per_request(vllm_config) -> int: + """A verify step's tokens per request: 1 + draft tokens, 1 without + speculation.""" + spec = vllm_config.speculative_config + return 1 + spec.num_speculative_tokens if spec is not None else 1 + + +def kda_dtypes(vllm_config): + """(conv, recurrent) KDA state dtypes the deployment's cache config selects.""" + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator + + cache = vllm_config.cache_config + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, + cache.mamba_cache_dtype, + cache.mamba_ssm_cache_dtype, + ) + + +def _widths(vllm_config) -> tuple[int, ...]: + """Every width a launch serves: the prepared batches, plus each captured + size rounded up to whole requests, up to ``PREPARED_MAX``.""" + q = tokens_per_request(vllm_config) + sizes = vllm_config.compilation_config.cudagraph_capture_sizes or () + verify = {round_up(s, q) for s in sizes} + return tuple(sorted({*BATCHES, *(m for m in verify if m <= PREPARED_MAX)})) + + +def _refuse_config(vllm_config) -> list[str]: + """Every reason this config cannot run the kernels, before load.""" + from vllm import envs + from vllm.config import CompilationMode, CUDAGraphMode + from vllm.platforms import current_platform + + parallel, spec = vllm_config.parallel_config, vllm_config.speculative_config + text = vllm_config.model_config.hf_text_config + moe = tuple( + getattr(text, name, None) + for name in ( + "hidden_act", + "activation_situ_beta", + "activation_situ_linear_beta", + "num_experts", + "num_experts_per_token", + "num_expert_group", + "topk_group", + "moe_router_activation_func", + ) + ) + graphs = vllm_config.compilation_config.cudagraph_mode + reasons = [] + if not current_platform.is_device_capability_family(100): + reasons.append("needs an SM100/103 GPU") + if moe != MOE: + reasons.append(f"needs Kimi K3's SiTU MoE {MOE}, not {moe}") + if parallel.tensor_parallel_size not in (8, 16): + reasons.append("needs tensor parallel size 8 or 16") + if parallel.decode_context_parallel_size not in ( + 1, + parallel.tensor_parallel_size, + ): + reasons.append("needs decode context parallel size 1 or the TP size") + # The decode loop runs in forward; a compiled forward would trace it. + if vllm_config.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: + reasons.append("needs a forward that is not torch.compile'd") + # A launch only ever runs while a FULL decode graph is captured. + if graphs is None or graphs.decode_mode() != CUDAGraphMode.FULL: + reasons.append("needs FULL decode CUDA graphs (not enforce-eager)") + if spec is not None: + if spec.method != "dspark": + reasons.append(f"needs the DSpark speculative method, not {spec.method}") + if spec.enable_adaptive_verification: + reasons.append("needs adaptive verification off") + if spec.num_speculative_tokens > MAX_SPECULATIVE_TOKENS: + reasons.append(f"needs at most {MAX_SPECULATIVE_TOKENS} draft tokens") + if envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: + reasons.append("needs the AttnRes aux stream off while speculating") + # KdaSplit keeps a BF16 conv state and an FP32 recurrent state in place; it + # maintains neither other cache dtypes nor ReplaySSM / RecoverSSM records. + if kda_dtypes(vllm_config) != KDA_STATE_DTYPES: + reasons.append( + "needs a BF16 KDA conv state and an FP32 recurrent state " + "(--mamba-cache-dtype)" + ) + if vllm_config.cache_config.use_replayssm or ( + vllm_config.cache_config.use_kda_recoverssm + ): + reasons.append("needs ReplaySSM / RecoverSSM off (--use-replayssm)") + return reasons + + +def _refuse_model(backbone) -> list[str]: + """Every reason the loaded backbone cannot run the kernels.""" + reasons = [] + hidden = getattr(backbone.config, "hidden_size", None) + if hidden != HIDDEN_SIZE: + reasons.append(f"needs hidden size {HIDDEN_SIZE}, the model has {hidden}") + if not getattr(backbone, "num_attn_res_blocks", 0): + reasons.append("needs the backbone's AttnRes blocks") + return reasons + + +def _refuse_caches(vllm_config, kernel: "KimiK3Kernel", caches) -> list[str]: + """Every reason the prepared cores cannot compile against these caches: a + KDA layer's states are kept in place, and an MLA layer's KV is indexed as + one contiguous tensor.""" + reasons = [] + for name, kind in kernel.cache_kinds.items(): + cache = caches.get(name) + tensors = [cache] if isinstance(cache, torch.Tensor) else list(cache or ()) + if kind == "kda": + if len(tensors) < 2: + reasons.append(f"{name}: {len(tensors)} KDA states, needs 2") + continue + dtypes = tuple(t.dtype for t in tensors[:2]) + if dtypes != KDA_STATE_DTYPES: + reasons.append(f"{name}: KDA states {dtypes}, needs {KDA_STATE_DTYPES}") + elif not tensors: + reasons.append(f"{name}: no MLA KV cache") + elif not tensors[0].is_contiguous(): + reasons.append(f"{name}: MLA KV cache is not contiguous") + return reasons + + +def _skip_step(vllm_config, kernel: "KimiK3Kernel", width: int) -> str | None: + """Why this step keeps vLLM's path, or None when the kernels take it.""" + # Every stock kernel has to run before its graph is captured, so a launch + # only ever replaces one inside a capture: profiling, warm-up and the + # warm-up run before each capture all run vLLM's path. + if not torch.cuda.is_current_stream_capturing(): + return "not capturing a CUDA graph" + if kernel.step_plan(width) is None: + return "not a decode or uniform verify step the providers serve" + return None + + +def _build(vllm_config, model, widths: tuple[int, ...]) -> "KimiK3Kernel": + from .runtime.decode import KimiK3Kernel + + return KimiK3Kernel(vllm_config, model, widths) + + +DECODE_SPEC = DecodeSpec["KimiK3Kernel"]( + name="kimi-k3-sm100", + supports=frozenset({Feature.DECODE_CONTEXT_PARALLEL, Feature.SPECULATIVE_DECODING}), + widths=_widths, + refuse_config=_refuse_config, + refuse_model=_refuse_model, + refuse_caches=_refuse_caches, + skip_step=_skip_step, + build=_build, +) From 8256f687ef0e395cd0151ec6979d8125ffa5340e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 22:16:13 -0500 Subject: [PATCH 49/79] fix: add generic param to the dsv41 implementation Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/model.py | 4 ++-- .../kimi_k3/nvidia/specialized_decode/runtime/providers.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index c176090fa41e..d9ccfff53e6c 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -60,7 +60,7 @@ DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention -from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPEC +from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPEC, DSV41Kernel from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors @@ -1387,7 +1387,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.vllm_config = vllm_config self.specialized_decode_spec = select_decode_spec(vllm_config, DECODE_SPEC) - self.specialized_decode: SpecializedDecode | None = None + self.specialized_decode: SpecializedDecode[DSV41Kernel] | None = None expert_dtype = getattr(config, "expert_dtype", "fp4") self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py index 63e96bfb79b4..8bbfc1ada0c4 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Per-layer attention providers (input projection, KDA/MLA core, o-projection) and -the helpers ``SpecializedDecode.build_providers`` uses.""" +the helpers ``KimiK3Kernel.build_providers`` uses.""" from typing import Any From 36feceaef8e90ced0d2d8264bbb50dc7c520a176 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 22:29:26 -0500 Subject: [PATCH 50/79] test: add unit tests Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 6 +- tests/models/common/__init__.py | 2 + .../common/specialized_decode/__init__.py | 2 + tests/models/common/specialized_decode/kit.py | 116 +++++++ .../common/specialized_decode/reference.py | 222 ++++++++++++++ .../test_specialized_decode.py | 282 ++++++++++++++++++ tests/models/test_dsv41_specialized_decode.py | 7 +- 7 files changed, 629 insertions(+), 8 deletions(-) create mode 100644 tests/models/common/__init__.py create mode 100644 tests/models/common/specialized_decode/__init__.py create mode 100644 tests/models/common/specialized_decode/kit.py create mode 100644 tests/models/common/specialized_decode/reference.py create mode 100644 tests/models/common/specialized_decode/test_specialized_decode.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 04ca2e1bbc49..cd417098bfad 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -71,8 +71,9 @@ steps: - tests/models/test_deepseek_v41_replay_batch.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py + - 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 + - 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/common/specialized_decode mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -96,8 +97,9 @@ steps: - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_specialized_decode.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_specialized_decode.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/common/specialized_decode - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/common/__init__.py b/tests/models/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/models/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/common/specialized_decode/__init__.py b/tests/models/common/specialized_decode/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/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 diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py new file mode 100644 index 000000000000..72e307308563 --- /dev/null +++ b/tests/models/common/specialized_decode/kit.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The specialized decode conformance kit (RFC #60904): what a model's own +``DecodeSpec`` must satisfy. + +The library's behaviour is tested once, in +``tests/models/common/specialized_decode/test_specialized_decode.py``. Only what +a model answers for itself belongs here, so a model pays for it in CI. +""" + +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: it is imported at config time, on +every platform, before anything has decided the spec can run.""" + + +@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 + source = ( + "import sys; sys.path.insert(0, '.'); import importlib;" + f"importlib.import_module({module!r});" + "print('\\n'.join(sys.modules))" + ) + 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 diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py new file mode 100644 index 000000000000..5eea922602b8 --- /dev/null +++ b/tests/models/common/specialized_decode/reference.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The reference kernel: a ``DecodeImpl`` over the layers it replaces.""" + +from dataclasses import replace +from types import SimpleNamespace +from typing import Any + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import get_forward_context +from vllm.models.common.specialized_decode.spec import DecodeSpec, Step + +CACHE_DTYPE = torch.bfloat16 +CACHE_NDIM = 4 + + +class ReferenceLayer(torch.nn.Module): + """A decoder layer the reference kernel can replace: one weight, one + attention layer name, and ``dense`` for the layers a spec skips.""" + + def __init__(self, index: int, hidden: int, dense: bool = False): + super().__init__() + self.index = index + self.dense = dense + self.attn_name = f"layers.{index}.attn" + self.weight = torch.nn.Parameter(torch.ones(hidden) * index) + + def forward(self, positions: torch.Tensor, hidden: torch.Tensor): + return hidden + self.weight + + +class ReferenceModel(torch.nn.Module): + """A model shaped like the layer loop a kernel hooks into.""" + + def __init__(self, layers: int = 4, hidden: int = 8): + super().__init__() + self.config = SimpleNamespace(num_hidden_layers=layers, hidden_size=hidden) + self.layers = torch.nn.ModuleList( + [ReferenceLayer(i, hidden, dense=i == 0) for i in range(layers)] + ) + self.start_layer = 0 + self.end_layer = layers + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden = inputs_embeds + for layer in self.layers: + hidden = layer(positions, hidden) + return hidden + + +class ReferenceKernel: + """The model's decode in one launch: the work of every layer it replaces, + off views of the model's weights, with the dense layers left to vLLM.""" + + def __init__(self, vllm_config, model: ReferenceModel, widths: tuple[int, ...]): + self.model = model + self.widths = widths + served = [ + model.layers[i] + for i in range(model.start_layer, model.end_layer) + if not model.layers[i].dense + ] + self.weights = {layer.index: layer.weight.data for layer in served} + self.cache_layers = tuple(layer.attn_name for layer in served) + self.launched: list[int] = [] + self.bound: dict[str, Any] = {} + self.compiled: set[Any] = set() + + def run(self, step: Step, input_ids, inputs_embeds, **kwargs): + self.launched.append(step.width) + hidden = inputs_embeds + for layer in self.model.layers: + weight = self.weights.get(layer.index) + if weight is None: + hidden = layer(step.positions, hidden) + else: + hidden = hidden + weight + return hidden + + def bind_caches(self, caches) -> None: + self.bound = dict(caches) + + def warmup_keys(self) -> list[Any]: + return [(i, w) for i in self.weights for w in self.widths] + + def compile(self, key: Any) -> None: + self.compiled.add(key) + + +def widths(vllm_config) -> tuple[int, ...]: + most = vllm_config.scheduler_config.max_num_seqs + return tuple(w for w in (1, 2, 4, 8) if w <= most) + + +def refuse_config(vllm_config) -> list[str]: + reasons = [] + if not widths(vllm_config): + reasons.append("max_num_seqs below the smallest width") + if vllm_config.model_config.dtype is not torch.bfloat16: + reasons.append(f"dtype {vllm_config.model_config.dtype}: needs bfloat16") + return reasons + + +def refuse_model(model) -> list[str]: + return [] if isinstance(model, ReferenceModel) else ["not a ReferenceModel"] + + +def skip_step(vllm_config, kernel: ReferenceKernel, width: int) -> str | None: + md = get_forward_context().attn_metadata + return "prefill" if md.get("prefill") else None + + +def refuse_caches(vllm_config, kernel: ReferenceKernel, caches) -> list[str]: + reasons = [] + for name, cache in caches.items(): + if cache.dtype is not CACHE_DTYPE: + reasons.append(f"{name}: dtype {cache.dtype}, needs {CACHE_DTYPE}") + if cache.ndim != CACHE_NDIM: + reasons.append(f"{name}: {cache.ndim}d, needs {CACHE_NDIM}d") + return reasons + + +def build(vllm_config, model, widths: tuple[int, ...]) -> ReferenceKernel: + return ReferenceKernel(vllm_config, model, widths) + + +REFERENCE = DecodeSpec[ReferenceKernel]( + name="reference", + supports=frozenset(), + widths=widths, + refuse_config=refuse_config, + refuse_model=refuse_model, + refuse_caches=refuse_caches, + skip_step=skip_step, + build=build, +) + +REFERENCE_WIDE = replace( + REFERENCE, + name="reference-wide", + refuse_config=lambda c: ( + refuse_config(c) + + (["needs max_num_seqs >= 8"] if c.scheduler_config.max_num_seqs < 8 else []) + ), +) + +DECODE_SPECS = (REFERENCE_WIDE, REFERENCE) + + +def make_config(**overrides) -> Any: + """A config shaped like the parts of ``VllmConfig`` the library reads. + + Args: + **overrides: Top-level attributes to replace. Nested ones are set on + the returned object (``config.parallel_config.enable_eplb = True``). + + """ + config = SimpleNamespace( + kernel_config=SimpleNamespace(specialized_decode=True), + model_config=SimpleNamespace(dtype=torch.bfloat16, enable_sleep_mode=False), + scheduler_config=SimpleNamespace(max_num_seqs=4), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=False, + enable_eplb=False, + decode_context_parallel_size=1, + use_ubatching=False, + prefill_context_parallel_size=1, + ), + offload_config=SimpleNamespace( + uva=SimpleNamespace(cpu_offload_gb=0), + prefetch=SimpleNamespace(offload_group_size=0), + ), + aux_output_config=SimpleNamespace(enable_return_routed_experts=False), + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.FULL, + cudagraph_capture_sizes=[1, 2, 3, 4, 16], + static_forward_context={}, + ), + lora_config=None, + kv_transfer_config=None, + speculative_config=None, + weight_transfer_config=None, + use_v2_model_runner=True, + ) + for key, value in overrides.items(): + setattr(config, key, value) + return config + + +def make_model(config, layers: int = 4) -> ReferenceModel: + """A loaded reference model, with its KV caches already in the config's + static forward context.""" + model = ReferenceModel(layers=layers) + set_caches(config, model) + return model + + +def set_caches(config, model, **kwargs) -> dict[str, torch.Tensor]: + """Puts one cache a layer in the static forward context, and returns them.""" + caches = { + layer.attn_name: make_cache(**kwargs) + for layer in model.layers + if not layer.dense + } + config.compilation_config.static_forward_context = { + name: SimpleNamespace(kv_cache=cache) for name, cache in caches.items() + } + return caches + + +def make_cache(dtype=torch.bfloat16, ndim: int = 4) -> torch.Tensor: + return torch.zeros((2,) + (1,) * (ndim - 1), dtype=dtype) diff --git a/tests/models/common/specialized_decode/test_specialized_decode.py b/tests/models/common/specialized_decode/test_specialized_decode.py new file mode 100644 index 000000000000..cc6f77da95fe --- /dev/null +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The specialized decode library against the reference spec: which variant a +config gets, what the build and the cache binding hand the kernel, and which +steps reach it.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from tests.models.common.specialized_decode import reference +from tests.models.common.specialized_decode.kit import forward +from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry +from vllm.models.common.specialized_decode.dispatch import SpecializedDecode +from vllm.models.common.specialized_decode.guards import ( + _features_in_use, + _refuse_config, + _select_spec, + select_decode_spec, +) +from vllm.models.common.specialized_decode.spec import ( + DecodeUnsupported, + Feature, + Step, +) + +SPECS = reference.DECODE_SPECS + + +@pytest.fixture +def config(): + return reference.make_config() + + +@pytest.fixture +def model(config): + return reference.make_model(config) + + +def make_decode(config, model, bind: bool = True): + """This rank's dispatch, built the way a model builds it.""" + spec = select_decode_spec(config, SPECS) + assert spec is not None + decode = SpecializedDecode(model, config, spec) + if bind: + decode.bind_caches() + return decode + + +@pytest.fixture +def decode(config, model): + return make_decode(config, model) + + +def positions(width: int) -> torch.Tensor: + return torch.zeros(width, dtype=torch.long) + + +# Selection + + +def test_off_without_the_flag(config): + config.kernel_config.specialized_decode = False + sizes = list(config.compilation_config.cudagraph_capture_sizes) + assert select_decode_spec(config, SPECS) is None + assert config.compilation_config.cudagraph_capture_sizes == sizes + + +def test_the_first_servable_variant_wins(config): + config.scheduler_config.max_num_seqs = 4 + spec, refusals = _select_spec(SPECS, config) + assert spec is reference.REFERENCE + assert refusals == {"reference-wide": ["needs max_num_seqs >= 8"]} + + config.scheduler_config.max_num_seqs = 8 + spec, refusals = _select_spec(SPECS, config) + assert spec is reference.REFERENCE_WIDE + assert refusals == {} + + +def test_every_variant_refused_names_every_variant(config): + config.model_config.dtype = torch.float16 + with pytest.raises(DecodeUnsupported) as excinfo: + select_decode_spec(config, SPECS) + message = str(excinfo.value) + for spec in SPECS: + assert spec.name in message + assert message.count("needs bfloat16") == len(SPECS) + + +def test_refusals_report_every_reason(config): + config.parallel_config.enable_eplb = True + config.use_v2_model_runner = False + config.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + config.model_config.dtype = torch.float16 + assert _refuse_config(reference.REFERENCE, config) == [ + "eplb is not supported", + "needs Model Runner V2", + "CUDA graph mode PIECEWISE: decode needs FULL graphs", + f"dtype {torch.float16}: needs bfloat16", + ] + + +IN_USE = [ + ( + Feature.PIPELINE_PARALLEL, + lambda c: setattr(c.parallel_config, "pipeline_parallel_size", 2), + ), + ( + Feature.DATA_PARALLEL, + lambda c: setattr(c.parallel_config, "data_parallel_size", 2), + ), + ( + Feature.EXPERT_PARALLEL, + lambda c: setattr(c.parallel_config, "enable_expert_parallel", True), + ), + (Feature.EPLB, lambda c: setattr(c.parallel_config, "enable_eplb", True)), + ( + Feature.DECODE_CONTEXT_PARALLEL, + lambda c: setattr(c.parallel_config, "decode_context_parallel_size", 2), + ), + (Feature.UBATCHING, lambda c: setattr(c.parallel_config, "use_ubatching", True)), + ( + Feature.PREFILL_CONTEXT_PARALLEL, + lambda c: setattr(c.parallel_config, "prefill_context_parallel_size", 2), + ), + (Feature.LORA, lambda c: setattr(c, "lora_config", object())), + (Feature.KV_CONNECTOR, lambda c: setattr(c, "kv_transfer_config", object())), + ( + Feature.ROUTED_EXPERT_CAPTURE, + lambda c: setattr(c.aux_output_config, "enable_return_routed_experts", True), + ), + ( + Feature.SPECULATIVE_DECODING, + lambda c: setattr(c, "speculative_config", object()), + ), + (Feature.SLEEP_MODE, lambda c: setattr(c.model_config, "enable_sleep_mode", True)), + ( + Feature.CPU_WEIGHT_OFFLOAD, + lambda c: setattr(c.offload_config.uva, "cpu_offload_gb", 1), + ), + (Feature.WEIGHT_UPDATES, lambda c: setattr(c, "weight_transfer_config", object())), +] + + +def test_every_feature_can_be_turned_on(): + assert {feature for feature, _ in IN_USE} == set(Feature) + + +@pytest.mark.parametrize( + "feature, turn_on", IN_USE, ids=[feature.value for feature, _ in IN_USE] +) +def test_a_feature_the_spec_does_not_list_refuses(config, feature, turn_on): + turn_on(config) + assert _features_in_use(config) == {feature} + assert _refuse_config(reference.REFERENCE, config) == [ + f"{feature.value} is not supported" + ] + serving = replace(reference.REFERENCE, supports=frozenset({feature})) + assert _refuse_config(serving, config) == [] + + +def test_the_reference_config_reads_like_a_real_one(): + """The fake config the tests pass has every field ``_features_in_use`` + reads off a real ``VllmConfig``, so a new feature cannot drift past it.""" + real = VllmConfig() + real.model_config = SimpleNamespace(enable_sleep_mode=False) + assert _features_in_use(real) == _features_in_use(reference.make_config()) == set() + + +def test_capture_sizes_become_the_widths(config): + config.compilation_config.cudagraph_capture_sizes = [1, 3, 4, 16, 32] + assert select_decode_spec(config, SPECS) is reference.REFERENCE + assert config.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 16, 32] + + +# Build + + +def test_build_takes_the_loaded_model_and_every_width(config, model, decode): + assert decode.widths == reference.widths(config) + assert decode.kernel.model is model + assert decode.kernel.widths == decode.widths + + +def test_the_loaded_model_can_refuse(config): + spec = select_decode_spec(config, SPECS) + with pytest.raises(DecodeUnsupported, match="not a ReferenceModel"): + SpecializedDecode(torch.nn.Linear(2, 2), config, spec) + + +def test_a_torch_compiled_model_is_refused(config): + spec = select_decode_spec(config, SPECS) + compiled = object.__new__(TorchCompileWithNoGuardsWrapper) + with pytest.raises(DecodeUnsupported, match="support_torch_compile"): + SpecializedDecode(compiled, config, spec) + + +def test_warmup_compiles_every_key_the_kernel_declares(config, model): + registry = JitWarmupRegistry(config) + with registry.activate(): + decode = make_decode(config, model, bind=False) + assert len(registry) == 1 + registry.warmup() + assert decode.kernel.compiled == set(decode.kernel.warmup_keys()) + + +def test_a_kernel_without_keys_registers_no_warmup(config, model): + spec = replace( + select_decode_spec(config, SPECS), + build=lambda c, m, w: SimpleNamespace( + cache_layers=(), warmup_keys=lambda: [], bind_caches=lambda caches: None + ), + ) + registry = JitWarmupRegistry(config) + with registry.activate(): + SpecializedDecode(model, config, spec) + assert len(registry) == 0 + + +# Cache binding + + +def test_binding_hands_the_kernel_its_cache_layers(config, model, decode): + assert set(decode.kernel.bound) == set(decode.kernel.cache_layers) + + rebound = reference.set_caches(config, model) + decode.bind_caches() + assert set(decode.kernel.bound) == set(decode.kernel.cache_layers) + assert all(decode.kernel.bound[name] is rebound[name] for name in rebound) + + +def test_caches_the_spec_refuses_never_reach_the_kernel(config, model, decode): + before = decode.kernel.bound + reference.set_caches(config, model, dtype=torch.float16) + with pytest.raises(DecodeUnsupported, match="needs torch.bfloat16"): + decode.bind_caches() + assert decode.kernel.bound is before + + +# The per-step decision + + +def test_a_served_decode_step_takes_the_kernel(config, decode): + rows = positions(min(decode.widths)) + with forward(config, CUDAGraphMode.FULL, {"prefill": False}): + step = decode.plan(rows) + assert isinstance(step, Step) + assert step.width == min(decode.widths) + assert step.positions is rows + + +@pytest.mark.parametrize( + "case", ["pipeline stage", "width", "piecewise", "unbound caches", "skip_step"] +) +def test_steps_the_library_sends_to_vllms_path(config, model, case): + decode = make_decode(config, model, bind=case != "unbound caches") + rows = positions(max(decode.widths) + 1 if case == "width" else min(decode.widths)) + mode = CUDAGraphMode.PIECEWISE if case == "piecewise" else CUDAGraphMode.FULL + stage = object() if case == "pipeline stage" else None + with forward(config, mode, {"prefill": case == "skip_step"}): + assert decode.plan(rows, stage) is None + + +def test_a_step_without_a_forward_context_takes_vllms_path(decode): + assert decode.plan(positions(min(decode.widths))) is None + + +def test_run_takes_the_planned_step_to_the_kernel(config, model, decode): + width = min(decode.widths) + rows = positions(width) + embeds = torch.zeros(width, model.config.hidden_size) + with forward(config, CUDAGraphMode.FULL, {"prefill": False}): + step = decode.plan(rows) + hidden = decode.run(step, None, embeds) + assert decode.kernel.launched == [width] + assert torch.equal(hidden, model(None, rows, None, embeds)) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 9f05f175eee4..585d020f68aa 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -210,6 +210,7 @@ def make_config(**overrides) -> Any: enable_expert_parallel=False, enable_eplb=False, decode_context_parallel_size=1, + use_ubatching=False, prefill_context_parallel_size=1, ), offload_config=SimpleNamespace( @@ -280,12 +281,6 @@ def make_model(self, config): ) return model - def decode_metadata(self, config, model): - return _metadata(model) - - def skipped_metadata(self, config, model): - return _metadata(model, num_prefills=1) - def bad_caches(self, config, model): return {"swa": torch.zeros(4, dtype=torch.bfloat16)} From 781963ebca38c92c53cc7f705e46c4cfaf5dca5d Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 23:08:51 -0500 Subject: [PATCH 51/79] chore: remove comment Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/model.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index d9ccfff53e6c..42b4c1520957 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -384,10 +384,6 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) elif not self.attn.wo_b.reduce_results: - # A partial-layer specialized decode: the FFN launch takes wo_b's - # TP partial and the rest of the layer. Whole-layer launches never - # come through here, and a new kernel should cover a whole layer - # instead of reaching into one. if decode_ffn is not None: out = decode_ffn(x, residual, post_mix, res_mix, attn_pre) if out is not None: From 90a01c657aa3bfc1e5b8fbe9ee5b547e3c76bc4d Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 15 Oct 2026 23:15:01 -0500 Subject: [PATCH 52/79] fix: finalize dsv41 weights after loading Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/model.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 42b4c1520957..e5c150291f8e 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -1494,13 +1494,18 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - self.process_weights_after_loading() + self.finalize_weights() return loaded_params - def process_weights_after_loading(self) -> None: + def finalize_weights(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() - if self.specialized_decode_spec is not None: + + def process_weights_after_loading(self) -> None: + # The loader calls this once every module's own post-load processing has + # run, which is what the specialized decode reads its weights in. + self.finalize_weights() + if self.specialized_decode_spec is not None and self.specialized_decode is None: self.specialized_decode = SpecializedDecode( self.model, self.vllm_config, self.specialized_decode_spec ) From 0bc115fc2442d932ebe837b462d57c34b55e0878 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 16 Oct 2026 10:15:07 -0500 Subject: [PATCH 53/79] fix: only create specializeddecode in global process_weights_after_loading and dont run replay Signed-off-by: simondanielsson --- tests/models/common/specialized_decode/kit.py | 11 ++++-- tests/models/test_dsv41_mono_layers.py | 39 ++++++++++++++----- tests/models/test_dsv41_specialized_decode.py | 12 ++++-- vllm/models/deepseek_v41/amd/model.py | 15 ++++--- .../deepseek_v41/amd/specialized_decode.py | 12 ++++-- vllm/models/deepseek_v41/amd/vl_model.py | 21 +++++----- 6 files changed, 73 insertions(+), 37 deletions(-) diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py index 72e307308563..274657471bae 100644 --- a/tests/models/common/specialized_decode/kit.py +++ b/tests/models/common/specialized_decode/kit.py @@ -22,8 +22,10 @@ from vllm.models.common.specialized_decode.guards import select_decode_spec KERNEL_BACKENDS = ("aiter", "flydsl") -"""Backends a spec module must not pull in: it is imported at config time, on -every platform, before anything has decided the spec can run.""" +"""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 @@ -88,10 +90,13 @@ def make_decode(self, config, model=None, bind: bool = True): 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(sys.modules))" + "print('\\n'.join(set(sys.modules) - loaded))" ) loaded = subprocess.run( [sys.executable, "-c", source], diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index 09bb18321850..98b20f4ab78b 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -136,19 +136,22 @@ def capture_steps(self, lens: list[int]) -> None: self._captured: dict = {} lens = sorted(lens) + def record(layer, layer_id, pos, out) -> None: + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(layer.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + kept = [t[row].float().cpu() for t in out[:2]] + self._captured.setdefault((layer_id, n, p), kept) + def hook(layer_id): def keep(module, args, kwargs, out): - md = get_forward_context().attn_metadata - if not isinstance(md, dict): - return - swa = md.get(module.attn.swa_cache_layer.prefix) - if swa is None or swa.num_prefills != 0: - return pos = kwargs.get("positions", args[1] if len(args) > 1 else None) - for row, p in enumerate(pos.tolist()): - n = max(n for n in lens if n <= p) - kept = [t[row].float().cpu() for t in out[:2]] - self._captured.setdefault((layer_id, n, p), kept) + record(module, layer_id, pos, out) return keep @@ -187,6 +190,22 @@ def pos_hook(module, args, kwargs): gap = gap_hook(i, layer) layer.ffn.register_forward_pre_hook(gap, with_kwargs=True) + # A whole-layer launch replaces the layer's forward, so no module hook + # sees it: wrap the launch itself to keep the same outputs. + decode = getattr(self.model_runner.model, "specialized_decode", None) # type: ignore[attr-defined] + if decode is None: + return + kernel = decode.kernel + launch_whole = kernel.launch_whole + + def launch(plan, x, positions, *args, **kwargs): + out = launch_whole(plan, x, positions, *args, **kwargs) + if out is not None: + record(plan.layer, plan.layer.attn.layer_id, positions, out) + return out + + kernel.launch_whole = launch + def save_capture(self, path: str) -> None: rank = self.rank # type: ignore[attr-defined] torch.save((self._captured, self._gaps), f"{path}.rank{rank}") diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 585d020f68aa..034dc9206619 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -178,14 +178,17 @@ def __init__(self, config): self.layers = torch.nn.ModuleList(_Layer(i) for i in range(LAYERS)) self.start_layer = 0 self.end_layer = LAYERS + self.decoder_replay_layers = None def prologue(self, input_ids, positions, intermediate_tensors, embeds, lookback): return embeds, input_ids, None, None, None, positions.shape[0] - def after_layer(self, idx, hidden, residual, post_mix, res_mix, tokens, aux): + def after_layer( + self, idx, hidden, residual, post_mix, res_mix, tokens, aux, replay=True + ): return None - def epilogue(self, hidden, *rest): + def epilogue(self, hidden, *rest, replay=True): return hidden @@ -301,7 +304,10 @@ def _rocm_platform() -> ModuleType: def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS) -> None: _stub_modules(monkeypatch) - monkeypatch.setattr(sd, "_compute_units", lambda: cus) + # The compute-unit probe reads the current accelerator, which a CPU run has + # none of. + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + monkeypatch.setattr(sd.current_platform, "num_compute_units", lambda _: cus) monkeypatch.setattr(_rocm_platform(), "get_cdna_version", lambda: cdna) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index e5c150291f8e..e47b1bd6d3a3 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -818,13 +818,15 @@ def epilogue( final_aux_recon: torch.Tensor | None, ran_layers: bool, full_num_tokens: int, + replay: bool = True, ) -> torch.Tensor | IntermediateTensors: """Everything after the last decoder layer: the replay batch, the final - post and collapse, and this rank's return.""" + post and collapse, and this rank's return. + """ # Replay collapses inside its own batch. The early loop still sets # ran_layers, so that flag alone must not post again on the full batch. collapsed = False - if self.decoder_replay_layers is not None: + if replay and self.decoder_replay_layers is not None: hidden_states, pre_mix, *late_aux = self.decoder_replay_layers( hidden_states, positions, @@ -941,15 +943,18 @@ def after_layer( res_mix: torch.Tensor | None, full_num_tokens: int, aux_hidden_states: list[torch.Tensor], + replay: bool = True, ) -> torch.Tensor | None: """This layer's aux hidden state, appended to ``aux_hidden_states``, or None when layer ``idx + 1`` does not want one. The aux id of the first replay layer is the post of the last full-batch - layer; that one is captured on the trimmed rows instead. + layer; that one is captured on the trimmed rows instead, unless + ``replay`` is False and the caller runs every layer on the full batch. """ if idx + 1 not in self.aux_hidden_state_layers or ( - self.decoder_replay_layers is not None + replay + and self.decoder_replay_layers is not None and idx + 1 == self.decoder_replay_start ): return None @@ -1502,8 +1507,6 @@ def finalize_weights(self) -> None: self.model.finalize_mhc_broadcast_weights() def process_weights_after_loading(self) -> None: - # The loader calls this once every module's own post-load processing has - # run, which is what the specialized decode reads its weights in. self.finalize_weights() if self.specialized_decode_spec is not None and self.specialized_decode is None: self.specialized_decode = SpecializedDecode( diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode.py index 5738d3d9ff67..c8235ef02cd3 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode.py @@ -98,10 +98,9 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): self.model = model self.runner = DSV41MonoLayer(world, rank, cpu_group, device) self.embed_layer = getattr(model, "start_layer", 0) - # The replay layers run on their own trimmed batch, outside the loop - # ``run`` owns, so no launch is built for them. - end = getattr(model, "end_layer", len(model.layers)) - self.layer_end = getattr(model, "decoder_replay_start", end) + # A decode step is never trimmed, so ``run`` owns every layer, the ones + # past the replay cut included. + self.layer_end = getattr(model, "end_layer", len(model.layers)) self.widths = widths self.swa_keys: list[str] = [] self.comp_keys: list[str] = [] @@ -321,6 +320,7 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: res_mix, full_num_tokens, aux_hidden_states, + replay=False, ) if aux_recon is not None: final_aux_recon = aux_recon @@ -336,6 +336,7 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: final_aux_recon, self.layer_end > self.embed_layer, full_num_tokens, + replay=False, ) @@ -440,6 +441,9 @@ def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: return "no layer takes a launch" if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches return "no attention metadata" + replay = decode.model.decoder_replay_layers + if replay is not None and replay.uses_replay_batch(): + return "replay batch step" if any(md.get(k) is None for k in decode.comp_keys): return "layer metadata missing" for key in decode.swa_keys: diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 5b8f7b826205..7fcad26dc25e 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -292,9 +292,13 @@ def embed_input_ids( @staticmethod def get_model_state_cls(): - from ..nvidia.model_state import DeepseekV41ModelState + from .model_state import DeepseekV41ROCMModelState - return DeepseekV41ModelState + return DeepseekV41ROCMModelState + + @property + def specialized_decode(self): + return self.language_model.specialized_decode @property def token_lookback_depth(self) -> int: @@ -342,15 +346,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # must not run on a partially loaded model). mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0]) loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(mapped) - # The child's load_weights already ran its post-load finalization. - self._weights_finalized = True - return loaded_params + return loader.load_weights(mapped) def process_weights_after_loading(self) -> None: - # Model-level post-load hook (called by the loader after any load - # format). Under DummyModelLoader the child's load_weights — and - # hence its finalize step — is bypassed, so run it here instead. - if getattr(self, "_weights_finalized", False): - return + # Model-level post-load hook, called by the loader once every module's + # own post-load processing has run. The child finalizes idempotently + # and builds its specialized decode here, so it always runs. self.language_model.process_weights_after_loading() From f745d67ff23dcb5053c231d2fee75c8104477fb5 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 16 Oct 2026 11:22:12 -0500 Subject: [PATCH 54/79] chore: cleanup docstrings Signed-off-by: simondanielsson --- vllm/config/kernel.py | 4 +- .../common/specialized_decode/__init__.py | 3 +- .../common/specialized_decode/dispatch.py | 37 ++++-------- .../common/specialized_decode/guards.py | 58 ++++--------------- vllm/models/common/specialized_decode/spec.py | 56 ++++++------------ 5 files changed, 43 insertions(+), 115 deletions(-) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 357d8622f1e9..7090fb8e24c7 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -294,8 +294,8 @@ class KernelConfig: specialized_decode: bool = False """If True, run decode steps on the model's specialized decode: a kernel - that runs most or all of a decoder layer in a single launch (DeepSeek-V4.1 - on MI355X). A deployment that cannot run it is refused at startup.""" + that runs most or all of a decoder layer in a single launch. A deployment + that cannot run it is refused at startup.""" moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/models/common/specialized_decode/__init__.py b/vllm/models/common/specialized_decode/__init__.py index 4b64dea3ecec..76a7b2d5855d 100644 --- a/vllm/models/common/specialized_decode/__init__.py +++ b/vllm/models/common/specialized_decode/__init__.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Specialized decode (RFC #60904, ``--specialized-decode``): a model's decode step -in one launch.""" +"""Specialized decode: a model's decode step in one launch.""" from vllm.models.common.specialized_decode.dispatch import SpecializedDecode from vllm.models.common.specialized_decode.guards import select_decode_spec diff --git a/vllm/models/common/specialized_decode/dispatch.py b/vllm/models/common/specialized_decode/dispatch.py index 5e87dd49c12c..3470dd90fb4c 100644 --- a/vllm/models/common/specialized_decode/dispatch.py +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""One ``SpecializedDecode`` a rank: the per-step decision and the cache checks, -behind the two hooks a model's forward calls (``plan``, then ``run``). +"""Dispatcher for specialized decoder kernels. + +Performs validation, jit registration, planning, and dispatching of specialized decode +kernels for a model based on a given specification. """ from dataclasses import dataclass @@ -63,15 +65,12 @@ class SpecializedDecode(Generic[KernelT]): """ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): - """Builds this rank's kernel. Collective over the TP group. + """Builds this rank's kernel. Args: model: The loaded model, in ``process_weights_after_loading``. vllm_config: The engine's config. - spec: The variant ``select_decode_spec`` picked. - - Raises: - DecodeUnsupported: This rank, or another, refused the loaded model. + spec: The spec to build and run. """ self.spec = spec @@ -91,16 +90,9 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): logger.info("Specialized decode %s: widths %s", spec.name, list(self.widths)) def bind_caches(self) -> None: - """Checks the caches the kernel will read and hands them to it, as the - runner binds them. - - Called from ``ModelState.initialize_kv_cache``, which runs once for the - throwaway cache memory profiling allocates and again for the real one. - The kernel only ever sees caches this spec has accepted. - - Raises: - DecodeUnsupported: The caches are not what this spec serves. + """Validate and bind caches to the kernel. + The underlying kernel only ever sees caches this spec has accepted. """ sfc = self.vllm_config.compilation_config.static_forward_context caches = {n: sfc[n].kv_cache for n in self.kernel.cache_layers} @@ -115,13 +107,8 @@ def bind_caches(self) -> None: def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | None: """This step's plan, or None to run the model's own forward instead. - Steps that run the model's forward: a pipeline stage's inputs, a width - off the spec's widths, a PIECEWISE step (keyed by token count alone, so - it is replayed for mixed batches too), caches the runner has not bound, - and whatever ``skip_step`` refuses. - - Every predicate is the same on every TP rank: the kernel waits - in-kernel on its peers, so the ranks must take the same path. + A step can fall back to running the model's own forward when the kernel cannot + serve it, as described by the spec. """ width = positions.size(-1) why: str | None @@ -145,6 +132,6 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non return Step(width=width, positions=positions) def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - """A planned step on the kernel, returning what the model's own forward - would have returned.""" + """Run the specialized kernel, returning what the model's own forward would + have returned.""" return self.kernel.run(step, input_ids, inputs_embeds, **kwargs) diff --git a/vllm/models/common/specialized_decode/guards.py b/vllm/models/common/specialized_decode/guards.py index 8206b982c09e..822053f50b8a 100644 --- a/vllm/models/common/specialized_decode/guards.py +++ b/vllm/models/common/specialized_decode/guards.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Refusals: the variant a config can serve, and TP agreement on every stage.""" +"""Common validation logic for specialized decode kernels.""" from collections.abc import Sequence from typing import TYPE_CHECKING, Any @@ -42,20 +42,12 @@ def tp_group() -> tuple[int, int, object | None]: def tp_agree(ok: bool, group) -> bool: - """Agrees a verdict every TP rank must take alike. - - The kernel waits in-kernel on every peer rank, so a rank that turned - specialized decode off alone, or failed to build what the others launch, - leaves the others spinning on the GPU for ever. Every rank must call this - at the same point. + """Agree on a boolean across the TP group. Args: ok: This rank's local verdict. group: The TP group's CPU (gloo) process group, or None for one rank. - Returns: - True when every rank of ``group`` passed ``ok`` true. - """ if group is None: return ok @@ -65,12 +57,12 @@ def tp_agree(ok: bool, group) -> bool: def _enabled(vllm_config: "VllmConfig") -> bool: - """Whether ``--specialized-decode`` asked for the model's specialized decode.""" + """Whether the engine is configured to run specialized decode.""" return bool(vllm_config.kernel_config.specialized_decode) def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: - """The deployment's features a spec has to list in ``supports``.""" + """Features in use by this config.""" pc = vllm_config.parallel_config oc = vllm_config.offload_config on = { @@ -97,8 +89,7 @@ def _features_in_use(vllm_config: "VllmConfig") -> set[Feature]: def _refuse_config(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> list[str]: - """Every reason this config cannot run ``spec``: default-deny on features, - then the runner and the graph mode, then the spec's own.""" + """The reasons the config cannot run the specialized decode kernel, if any.""" reasons = [ f"{f.value} is not supported" for f in sorted(_features_in_use(vllm_config) - spec.supports) @@ -112,9 +103,7 @@ def _refuse_config(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> list[str def _refuse_model(spec: DecodeSpec[Any], model) -> list[str]: - """Every reason the loaded model cannot run ``spec``. Specialized decode replaces - the layer loop, which a traced graph has already inlined, so a model - carrying ``@support_torch_compile`` is refused.""" + """The reasons the mode cannot run the specialized decode kernel, if any.""" from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper reasons = [] @@ -129,13 +118,7 @@ def _refuse_model(spec: DecodeSpec[Any], model) -> list[str]: def _select_spec( specs: Sequence[DecodeSpec[Any]], vllm_config: "VllmConfig" ) -> tuple[DecodeSpec[Any] | None, dict[str, list[str]]]: - """The first variant this config can serve, and every variant's refusals. - - Returns: - (spec, refusals). ``spec`` is None when every variant refused, and - ``refusals`` holds each variant's reasons, by spec name. - - """ + """Get the first compatible spec.""" refusals: dict[str, list[str]] = {} for spec in specs: reasons = _refuse_config(spec, vllm_config) @@ -145,17 +128,8 @@ def _select_spec( return None, refusals -def _refusal_message(refusals: dict[str, list[str]]) -> str: - """Every variant's reasons as one message.""" - return "; ".join( - f"{name}: {', '.join(reasons)}" for name, reasons in refusals.items() - ) - - def _set_capture_sizes(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> None: - """The capture sizes up to the largest width become exactly the widths, so - a replay never pads a request's rows; larger sizes stay as configured (and - capture vLLM's path).""" + """Amend the capture sizes with the supported widths to avoid padding.""" cc = vllm_config.compilation_config widths = spec.widths(vllm_config) sizes = cc.cudagraph_capture_sizes or [] @@ -171,23 +145,12 @@ def select_decode_spec( vllm_config: "VllmConfig", specs: DecodeSpec[KernelT] | Sequence[DecodeSpec[KernelT]], ) -> DecodeSpec[KernelT] | None: - """Picks the variant, sets the capture sizes, or raises with every reason. - - The variant is picked here, once: the model holds what this returns and - hands it to ``SpecializedDecode`` after the weights are loaded. A model - that prepares anything of its own at load time reads the returned spec to - know which variant it is preparing for. + """Select the first compatible spec and set the capture sizes, if enabled. Args: vllm_config: The engine's config, before the weights are loaded. specs: The model's variants, in the order they are preferred. - Returns: - The variant this config runs, or None with ``--specialized-decode`` off. - - Raises: - DecodeUnsupported: Every variant refused this config. - """ if not _enabled(vllm_config): return None @@ -195,7 +158,8 @@ def select_decode_spec( spec, refusals = _select_spec(ordered, vllm_config) if spec is None: raise DecodeUnsupported( - "Specialized decode refused: " + _refusal_message(refusals) + "Specialized decode refused: " + + "; ".join(f"{n}: {', '.join(rs)}" for n, rs in refusals.items()) ) _set_capture_sizes(spec, vllm_config) if vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index c73adba87018..fd73efe4aa20 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""What a model declares to run specialized decode: ``DecodeSpec`` and its -callbacks.""" +"""Interfaces for when and how to run a specialized decode kernel.""" from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -15,8 +14,7 @@ class Feature(str, Enum): - """What a deployment can turn on. A feature in use that a spec's - ``supports`` does not list is refused.""" + """Features a specialized decode kernel may require.""" PIPELINE_PARALLEL = "pipeline_parallel" DATA_PARALLEL = "data_parallel" @@ -50,8 +48,7 @@ class Step: class DecodeImpl(Protocol): - """A spec's half of the contract: one model's decode on this rank, for - every width the spec serves. + """A specialized decode kernel, substituting a model's forward. The implementation owns its own layer loop, so the library has no opinion about layers: it may replace one layer, two, or the whole backbone. @@ -61,30 +58,15 @@ class DecodeImpl(Protocol): """Attention layers whose caches ``DecodeSpec.refuse_caches`` checks.""" def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - """Takes and returns what the model's own forward does. - - A launch replaces a *whole* decoder layer: it returns what that layer's - forward returns, or declines, and the kernel runs the layer on vLLM's - path instead. Covering part of a layer is the model's own problem, and - is discouraged: a new kernel should cover a whole layer or fall back. - """ + """Takes and returns what the model's own forward does.""" ... def bind_caches(self, caches: Mapping[str, Any]) -> None: - """Takes the ``cache_layers`` caches the runner has bound, once - ``DecodeSpec.refuse_caches`` has passed them. - - For a kernel whose launches close over the cache layout (a compile - against its strides, a cached view), this is where that happens. Called - again whenever the runner rebinds, so it must replace what the last call - built, not add to it: a normal startup binds the throwaway cache memory - profiling allocates before the real one. - """ + """Bind the KV caches the kernel will read.""" ... def warmup_keys(self) -> list[Any]: - """Hashable build keys of this kernel's launches, compiled before CUDA - graph capture by vLLM's JIT warm-up.""" + """Build keys of this kernel's launches.""" ... def compile(self, key: Any) -> None: @@ -93,14 +75,12 @@ def compile(self, key: Any) -> None: KernelT = TypeVar("KernelT", bound=DecodeImpl) -"""What one spec's ``build`` returns, and what its ``refuse_caches`` and -``skip_step`` are handed: a spec declared ``DecodeSpec[MyKernel]`` ties the -three together.""" @dataclass(frozen=True) class DecodeSpec(Generic[KernelT]): - """Definition of when specialized decode can run, and how to build it. + """Definition of when specialized decode can run, and how to build its + implementation. Attributes: name: Log and metric name. @@ -108,17 +88,15 @@ class DecodeSpec(Generic[KernelT]): that is absent from it refuses the deployment. widths: The rows a launch serves, from the engine's config. These become the CUDA graph capture sizes. - refuse_config: Deployment refusals the engine's config answers, before - the weights are loaded. - refuse_model: Deployment refusals only the loaded model answers. - refuse_caches: Deployment refusals the KV caches the kernel will read - answer, re-run whenever the runner binds caches. - skip_step: Why one step takes vLLM's path, None when it takes the - kernel's. Takes the step's width; anything else it needs it reads - from the forward context, which every rank sees alike. - build: This rank's kernel, built once off the loaded model (whose - parameters carry the device it launches on) and every width the - spec serves. The only place a weight copy is ever made. + refuse_config: Conditions under which the specialized decode kernel cannot run, + based on the config. + refuse_model: Conditions under which the specialized decode kernel cannot run, + based on the model. + refuse_caches: Conditions under which the specialized decode kernel cannot run, + based on the cache. + skip_step: Whether to run the specialized decode kernel for this forward call, + or to fall back to the model's own forward. + build: This rank's kernel. """ From 4cf46c8c956c59935664eb9390018b780970ca04 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 16 Oct 2026 12:02:19 -0500 Subject: [PATCH 55/79] fix: don't require automatic jit Signed-off-by: simondanielsson --- .../common/specialized_decode/dispatch.py | 14 +++++++------ vllm/models/common/specialized_decode/spec.py | 20 ++++++++++++++++++- .../specialized_decode/runtime/decode.py | 7 ------- .../kimi_k3/nvidia/specialized_decode/spec.py | 6 +----- 4 files changed, 28 insertions(+), 19 deletions(-) diff --git a/vllm/models/common/specialized_decode/dispatch.py b/vllm/models/common/specialized_decode/dispatch.py index 3470dd90fb4c..82968a1bded8 100644 --- a/vllm/models/common/specialized_decode/dispatch.py +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -21,9 +21,9 @@ tp_group, ) from vllm.models.common.specialized_decode.spec import ( - DecodeImpl, DecodeSpec, DecodeUnsupported, + DecodeWarmable, KernelT, Step, ) @@ -42,7 +42,7 @@ class CompileKey: spec: str build: Any - def __init__(self, spec: str, kernel: DecodeImpl): + def __init__(self, spec: str, kernel: DecodeWarmable): super().__init__() self.spec = spec self.kernel = kernel @@ -58,8 +58,9 @@ def compile(self, compile_key: "DecodeWarmup.CompileKey") -> None: class SpecializedDecode(Generic[KernelT]): """A model's specialized decode on this rank. - It reads the spec and builds the kernel, registers it for warm-up, - checks the caches whenever the runner binds them, and decides each step + It reads the spec and builds the kernel, registers it for warm-up if it + compiles from keys, checks the caches whenever the runner binds them, and + decides each step between the kernel and vLLM's own forward. The kernel it builds is held in the ``kernel`` attribute. """ @@ -85,8 +86,9 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): self.widths = spec.widths(vllm_config) self._bound = False self.kernel: KernelT = spec.build(vllm_config, model, self.widths) - if self.kernel.warmup_keys(): - DecodeWarmup(spec.name, self.kernel).register_warmup() + kernel = self.kernel + if isinstance(kernel, DecodeWarmable) and kernel.warmup_keys(): + DecodeWarmup(spec.name, kernel).register_warmup() logger.info("Specialized decode %s: widths %s", spec.name, list(self.widths)) def bind_caches(self) -> None: diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index fd73efe4aa20..2deafe490705 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -5,7 +5,14 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Protocol, + TypeVar, + runtime_checkable, +) if TYPE_CHECKING: import torch @@ -65,6 +72,17 @@ def bind_caches(self, caches: Mapping[str, Any]) -> None: """Bind the KV caches the kernel will read.""" ... + +@runtime_checkable +class DecodeWarmable(Protocol): + """A kernel whose launches compile from keys, registered for vLLM's JIT + warm-up. + + SpecializedDecodes created inside process_weights_after_loading can implement + this for automatic warm-up of their kernels, so that the first decode step does not + pay the compilation cost. + """ + def warmup_keys(self) -> list[Any]: """Build keys of this kernel's launches.""" ... diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py index 794d48e9698c..d62b1505d0dd 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py @@ -190,13 +190,6 @@ def bind_caches(self, caches: Mapping[str, Any]) -> None: self.caches[layer.layer_idx] = tensors self.attention[layer.layer_idx].prepare_layout(tensors) - def warmup_keys(self) -> list[Any]: - """None: every compile finishes in ``build_providers`` and ``bind_caches``.""" - return [] - - def compile(self, key: Any) -> None: - raise AssertionError("Kimi K3 registers no warm-up keys") - def embed_input_ids(self, ids): rows = ids.shape[0] if ids.ndim == 1 and ids.is_contiguous() else 0 if not 1 <= rows <= PREPARED_MAX: diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py b/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py index 7aff5a4aae75..009d8236e4b0 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""What Kimi K3 declares to run its specialized decode on SM100/103. - -Imported at config time, on every platform, before anything has decided the -kernel can run: nothing here may pull in CuTe DSL. -""" +"""Kimi K3's specialized decode.""" from typing import TYPE_CHECKING From 2b4eaee8447c57748f5a825d5f1e74a638cff561 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 16 Oct 2026 17:48:48 -0500 Subject: [PATCH 56/79] fix: update directory structure Signed-off-by: simondanielsson --- .../test_specialized_decode.py | 17 +++++++++++ tests/models/test_dsv41_mono_numerics.py | 12 ++++---- tests/models/test_dsv41_specialized_decode.py | 11 ++++--- vllm/engine/arg_utils.py | 17 +---------- .../common/specialized_decode/amd/__init__.py | 7 +++++ .../specialized_decode/amd/mono/__init__.py | 14 +++++++++ .../specialized_decode/amd/mono}/mx.py | 2 +- .../specialized_decode/amd/mono}/ops.py | 2 +- .../amd/mono}/peer_memory.py | 0 .../specialized_decode/amd/mono}/plan.py | 27 ++++++++++++----- .../specialized_decode/amd/mono}/sync.py | 2 +- .../deepseek_v41/amd/mono/common/__init__.py | 9 ------ vllm/models/deepseek_v41/amd/mono/sources.py | 8 ----- .../__init__.py} | 17 +++++++---- .../{ => specialized_decode}/mono/README.md | 0 .../{ => specialized_decode}/mono/__init__.py | 0 .../mono/attention/__init__.py | 0 .../mono/attention/back.py | 0 .../mono/attention/device.py | 0 .../mono/attention/front.py | 0 .../mono/attention/gemv.py | 0 .../mono/attention/plan.py | 0 .../{ => specialized_decode}/mono/layer.py | 23 ++++++++++++--- .../{ => specialized_decode}/mono/runner.py | 7 +++-- .../amd/specialized_decode/mono/sources.py | 11 +++++++ .../mono/stages/__init__.py | 0 .../mono/stages/dims.py | 2 +- .../mono/stages/gemv.py | 8 ++--- .../mono/stages/moe.py | 29 ++++++++++++------- .../mono/stages/moe_plan.py | 2 +- .../mono/stages/moe_shape.py | 21 ++++++++++---- .../mono/stages/router.py | 4 +-- .../mono/stages/seam.py | 12 +++++--- 33 files changed, 171 insertions(+), 93 deletions(-) create mode 100644 vllm/models/common/specialized_decode/amd/__init__.py create mode 100644 vllm/models/common/specialized_decode/amd/mono/__init__.py rename vllm/models/{deepseek_v41/amd/mono/common => common/specialized_decode/amd/mono}/mx.py (97%) rename vllm/models/{deepseek_v41/amd/mono/common => common/specialized_decode/amd/mono}/ops.py (99%) rename vllm/models/{deepseek_v41/amd/mono/common => common/specialized_decode/amd/mono}/peer_memory.py (100%) rename vllm/models/{deepseek_v41/amd/mono/common => common/specialized_decode/amd/mono}/plan.py (85%) rename vllm/models/{deepseek_v41/amd/mono/common => common/specialized_decode/amd/mono}/sync.py (98%) delete mode 100644 vllm/models/deepseek_v41/amd/mono/common/__init__.py delete mode 100644 vllm/models/deepseek_v41/amd/mono/sources.py rename vllm/models/deepseek_v41/amd/{specialized_decode.py => specialized_decode/__init__.py} (97%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/README.md (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/attention/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/attention/back.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/attention/device.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/attention/front.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/attention/gemv.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/attention/plan.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/layer.py (98%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/runner.py (98%) create mode 100644 vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/dims.py (96%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/gemv.py (97%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/moe.py (98%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/moe_plan.py (96%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/moe_shape.py (93%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/router.py (95%) rename vllm/models/deepseek_v41/amd/{ => specialized_decode}/mono/stages/seam.py (97%) diff --git a/tests/models/common/specialized_decode/test_specialized_decode.py b/tests/models/common/specialized_decode/test_specialized_decode.py index cc6f77da95fe..e663825e9a91 100644 --- a/tests/models/common/specialized_decode/test_specialized_decode.py +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -4,6 +4,8 @@ config gets, what the build and the cache binding hand the kernel, and which steps reach it.""" +import importlib +import sys from dataclasses import replace from types import SimpleNamespace @@ -280,3 +282,18 @@ def test_run_takes_the_planned_step_to_the_kernel(config, model, decode): hidden = decode.run(step, None, embeds) assert decode.kernel.launched == [width] assert torch.equal(hidden, model(None, rows, None, embeds)) + + +def test_the_contract_imports_without_a_vendor_runtime(monkeypatch): + """The library loads where neither FlyDSL nor aiter is installed, so a + vendor's shared kernel code must stay out of its import graph.""" + for name in list(sys.modules): + if name.split(".")[0] in ("flydsl", "aiter") or "specialized_decode" in name: + monkeypatch.delitem(sys.modules, name, raising=False) + for name in ("flydsl", "aiter"): + monkeypatch.setitem(sys.modules, name, None) + + module = importlib.import_module("vllm.models.common.specialized_decode") + + assert module.DecodeSpec is not None + assert not [n for n in sys.modules if n.startswith(("flydsl.", "aiter."))] diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py index 83308d46b0c8..575b77d05e5c 100644 --- a/tests/models/test_dsv41_mono_numerics.py +++ b/tests/models/test_dsv41_mono_numerics.py @@ -273,15 +273,17 @@ def _check_layers(rank: int, tp: int, device) -> None: from vllm.model_executor.model_loader.utils import process_weights_after_loading from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE - from vllm.models.deepseek_v41.amd.mono.layer import scratch_layout - from vllm.models.deepseek_v41.amd.mono.runner import ( + from vllm.models.deepseek_v41.amd.rocm import ( + compute_global_topk_ragged_indices_and_indptr, + ) + from vllm.models.deepseek_v41.amd.specialized_decode.mono.layer import ( + scratch_layout, + ) + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( AttnWeights, DSV41MonoLayer, MonoLayerWeights, ) - from vllm.models.deepseek_v41.amd.rocm import ( - compute_global_topk_ragged_indices_and_indptr, - ) from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 034dc9206619..769c720da36f 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -49,19 +49,22 @@ def ffn(self, *args): return ("ffn", args) +_MONO = "vllm.models.deepseek_v41.amd.specialized_decode.mono" + + def _stub_modules(monkeypatch) -> None: """The FlyDSL runner package, as a spec-shaped stub.""" - runner: Any = ModuleType("vllm.models.deepseek_v41.amd.mono.runner") + runner: Any = ModuleType(f"{_MONO}.runner") runner.BLOCKS = BLOCKS runner.MAX_TOKENS = 64 runner.DSV41MonoLayer = _StubRunner runner.MonoLayerWeights = _weights_stub runner.AttnWeights = _weights_stub - plan: Any = ModuleType("vllm.models.deepseek_v41.amd.mono.attention.plan") + plan: Any = ModuleType(f"{_MONO}.attention.plan") plan.Dims = lambda world: world for name, module in ( - ("vllm.models.deepseek_v41.amd.mono.runner", runner), - ("vllm.models.deepseek_v41.amd.mono.attention.plan", plan), + (f"{_MONO}.runner", runner), + (f"{_MONO}.attention.plan", plan), ): monkeypatch.setitem(sys.modules, name, module) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 82503370dccd..7a726f7fd6af 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -470,12 +470,6 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) -_MONOKERNEL_DEPRECATION_MESSAGE = ( - "--monokernel is deprecated and will be removed in v0.33.0. " - "Use --specialized-decode instead." -) - - class DeprecatedLogConfigFileAction(argparse.Action): """Warn when the legacy ``--log-config-file`` option is used.""" @@ -776,7 +770,6 @@ class EngineArgs: KernelConfig, "enable_flashinfer_autotune" ) specialized_decode: bool = get_field(KernelConfig, "specialized_decode") - monokernel: bool | None = None worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1833,12 +1826,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: kernel_group.add_argument( "--specialized-decode", **kernel_kwargs["specialized_decode"] ) - kernel_group.add_argument( - "--monokernel", - action=argparse.BooleanOptionalAction, - default=None, - help=_MONOKERNEL_DEPRECATION_MESSAGE, - ) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -2792,9 +2779,7 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune - if self.monokernel is not None: - warnings.warn(_MONOKERNEL_DEPRECATION_MESSAGE, stacklevel=1) - if self.specialized_decode or self.monokernel: + if self.specialized_decode: kernel_config.specialized_decode = True if self.moe_backend != "auto": kernel_config.moe_backend = self.moe_backend diff --git a/vllm/models/common/specialized_decode/amd/__init__.py b/vllm/models/common/specialized_decode/amd/__init__.py new file mode 100644 index 000000000000..93c43d34f99e --- /dev/null +++ b/vllm/models/common/specialized_decode/amd/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kernel code AMD specialized decodes share, one subpackage a kernel flavor. + +Nothing here is imported by the specialized decode contract itself, which stays +free of vendor runtimes so that it loads on any platform. +""" diff --git a/vllm/models/common/specialized_decode/amd/mono/__init__.py b/vllm/models/common/specialized_decode/amd/mono/__init__.py new file mode 100644 index 000000000000..8eb2e8ce3069 --- /dev/null +++ b/vllm/models/common/specialized_decode/amd/mono/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/__init__.py +"""Mechanisms shared by the monokernels, the FlyDSL persistent-kernel flavor of +specialized decode: ``ops`` (wave / math primitives, the TP group's peer +buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX arithmetic), ``plan`` +(execution model, layouts, build keys) and ``peer_memory`` (symmetric peer +buffers). + +A model's own kernels live under its ``specialized_decode/mono/``, and name this +package's ``plan.SHARED`` in their build key so that an edit here rebuilds them. +""" diff --git a/vllm/models/deepseek_v41/amd/mono/common/mx.py b/vllm/models/common/specialized_decode/amd/mono/mx.py similarity index 97% rename from vllm/models/deepseek_v41/amd/mono/common/mx.py rename to vllm/models/common/specialized_decode/amd/mono/mx.py index 9c228abdbf27..21cc405cb046 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/mx.py +++ b/vllm/models/common/specialized_decode/amd/mono/mx.py @@ -14,7 +14,7 @@ from flydsl.expr import rocdl from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.mono.common.ops import CM_NT +from vllm.models.common.specialized_decode.amd.mono.ops import CM_NT FP8_MAX = 448.0 # the f8f6f4 MFMA's operand formats (cbsz / blgp) diff --git a/vllm/models/deepseek_v41/amd/mono/common/ops.py b/vllm/models/common/specialized_decode/amd/mono/ops.py similarity index 99% rename from vllm/models/deepseek_v41/amd/mono/common/ops.py rename to vllm/models/common/specialized_decode/amd/mono/ops.py index f9f7fb22ec92..16cafbfdb57e 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/ops.py +++ b/vllm/models/common/specialized_decode/amd/mono/ops.py @@ -35,7 +35,7 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T, as_ir_value -from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES +from vllm.models.common.specialized_decode.amd.mono.plan import THREADS, WAVES # ---------------------------------------------------------------- primitives # gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device diff --git a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py b/vllm/models/common/specialized_decode/amd/mono/peer_memory.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/common/peer_memory.py rename to vllm/models/common/specialized_decode/amd/mono/peer_memory.py diff --git a/vllm/models/deepseek_v41/amd/mono/common/plan.py b/vllm/models/common/specialized_decode/amd/mono/plan.py similarity index 85% rename from vllm/models/deepseek_v41/amd/mono/common/plan.py rename to vllm/models/common/specialized_decode/amd/mono/plan.py index 2b05d7a073e9..d37f5b2f3c4d 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/plan.py +++ b/vllm/models/common/specialized_decode/amd/mono/plan.py @@ -99,22 +99,33 @@ def padded(n: int, align: int) -> int: # ---------------------------------------------------------------- build keys _SCALARS = (int, float, bool, str, enum.Enum) -_ROOT = Path(__file__).resolve().parents[1] # the mono package +SHARED = Path(__file__).resolve().parent +"""This package, which every model's digest covers alongside its own.""" @functools.cache -def source_digest(*paths: str) -> str: - """sha256 of every ``.py`` file under ``paths`` (files or directories, - relative to the mono package), each file's path and contents: one model's - kernel sources, so an edit to another model's rebuilds none of these.""" +def source_digest(*roots: Path) -> str: + """sha256 of every ``.py`` file under ``roots`` (files or directories), each + file's path within its root and its contents: one model's kernel sources, so + an edit to another model's rebuilds none of these. + + Args: + roots: Absolute paths. A model passes its own package and ``SHARED``. + + Returns: + The digest's first 16 hex characters. + + Raises: + FileNotFoundError: A root holds no ``.py`` file. + + """ h = hashlib.sha256() - for rel in paths: - root = _ROOT / rel + for root in roots: files = [root] if root.is_file() else sorted(root.rglob("*.py")) if not files: raise FileNotFoundError(f"no kernel sources at {root}") for f in files: - h.update(str(f.relative_to(_ROOT)).encode()) + h.update(str(f.relative_to(root.parent)).encode()) h.update(f.read_bytes()) return h.hexdigest()[:16] diff --git a/vllm/models/deepseek_v41/amd/mono/common/sync.py b/vllm/models/common/specialized_decode/amd/mono/sync.py similarity index 98% rename from vllm/models/deepseek_v41/amd/mono/common/sync.py rename to vllm/models/common/specialized_decode/amd/mono/sync.py index 0220ce12c32a..5fcc355c32fc 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/sync.py +++ b/vllm/models/common/specialized_decode/amd/mono/sync.py @@ -14,7 +14,7 @@ from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.mono.common.ops import ( +from vllm.models.common.specialized_decode.amd.mono.ops import ( CM_DEV, CM_SYS, bf16_pair, diff --git a/vllm/models/deepseek_v41/amd/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/mono/common/__init__.py deleted file mode 100644 index 397642c6b441..000000000000 --- a/vllm/models/deepseek_v41/amd/mono/common/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, -# Copyright (C) 2026, Advanced Micro Devices, Inc.: -# atom/mono/__init__.py -"""The mono kernels' shared mechanisms: ``ops`` (wave / math primitives, the TP -group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX -arithmetic), ``plan`` (execution model, layouts, build keys) and -``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/deepseek_v41/amd/mono/sources.py b/vllm/models/deepseek_v41/amd/mono/sources.py deleted file mode 100644 index 49f34d8fed67..000000000000 --- a/vllm/models/deepseek_v41/amd/mono/sources.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The mono layer's kernel sources, in every build's JIT cache key -(``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" - -from vllm.models.deepseek_v41.amd.mono.common.plan import source_digest - -SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v41/amd/specialized_decode.py b/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py similarity index 97% rename from vllm/models/deepseek_v41/amd/specialized_decode.py rename to vllm/models/deepseek_v41/amd/specialized_decode/__init__.py index c8235ef02cd3..4a999c81180b 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py @@ -90,7 +90,10 @@ class LayerPlan: class DSV41Kernel: def __init__(self, vllm_config, model, widths: tuple[int, ...]): from vllm.models.common.specialized_decode.guards import tp_group - from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + MAX_TOKENS, + DSV41MonoLayer, + ) assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS rank, world, cpu_group = tp_group() @@ -378,7 +381,7 @@ def _refuse_config(vllm_config) -> list[str]: "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" ) try: - import vllm.models.deepseek_v41.amd.mono.runner # noqa: F401 + import vllm.models.deepseek_v41.amd.specialized_decode.mono.runner # noqa: F401 except ImportError as err: reasons.append(f"needs FlyDSL and AITER's FlyDSL helpers ({err})") return reasons @@ -389,7 +392,7 @@ def _refuse_model(model) -> list[str]: kernels read every one of the 384 experts on each rank, TP-sharded, in AITER's A8W4 MXFP4 layout, from CTAs that all stay resident.""" from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend - from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import BLOCKS from vllm.platforms.rocm import get_cdna_version reasons = [] @@ -466,7 +469,9 @@ def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: def _weights(layer, ffn_only: bool, world: int): - from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + MonoLayerWeights, + ) a, f = layer.attn, layer.ffn e, sh = f.experts.routed_experts, f.shared_experts @@ -496,8 +501,8 @@ def _weights(layer, ffn_only: bool, world: int): def _attn_weights(a, world: int): - from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims - from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights + from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import AttnWeights attn = AttnWeights( layer_id=a.layer_id, diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/specialized_decode/mono/README.md similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/README.md rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/README.md diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/__init__.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/__init__.py diff --git a/vllm/models/deepseek_v41/amd/mono/attention/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/attention/__init__.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/__init__.py diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/back.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/attention/back.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/back.py diff --git a/vllm/models/deepseek_v41/amd/mono/attention/device.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/device.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/attention/device.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/device.py diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/front.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/attention/front.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/front.py diff --git a/vllm/models/deepseek_v41/amd/mono/attention/gemv.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/gemv.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/attention/gemv.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/gemv.py diff --git a/vllm/models/deepseek_v41/amd/mono/attention/plan.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/plan.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/attention/plan.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/plan.py diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py similarity index 98% rename from vllm/models/deepseek_v41/amd/mono/layer.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py index dcd39ae2b754..a3f5b6c40eb2 100644 --- a/vllm/models/deepseek_v41/amd/mono/layer.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py @@ -22,14 +22,29 @@ from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import Int32, Int64, T +from vllm.models.common.specialized_decode.amd.mono.mx import clamp_fp8 +from vllm.models.common.specialized_decode.amd.mono.ops import ( + butterfly, + fp8_pack4, + peer_bases, + traced, +) +from vllm.models.common.specialized_decode.amd.mono.plan import ( + WAVES, + first_task, + key_tuple, +) +from vllm.models.common.specialized_decode.amd.mono.sync import ( + Mailbox, + preg, + publish, + sreg, +) + from .attention import back as mb_back from .attention import front as mb_front from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc from .attention.plan import HIDDEN, Dims, back_scratch, front_scratch -from .common.mx import clamp_fp8 -from .common.ops import butterfly, fp8_pack4, peer_bases, traced -from .common.plan import WAVES, first_task, key_tuple -from .common.sync import Mailbox, preg, publish, sreg from .sources import SOURCES from .stages import moe, seam from .stages.dims import Dims as StageDims diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py similarity index 98% rename from vllm/models/deepseek_v41/amd/mono/runner.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py index ef9a1978e422..91bfdf9f0803 100644 --- a/vllm/models/deepseek_v41/amd/mono/runner.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py @@ -22,8 +22,9 @@ import torch +from vllm.models.common.specialized_decode.amd.mono.plan import BLOCKS + from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, Dims -from .common.plan import BLOCKS from .layer import ( EPOCH_WORDS, MAX_TOKENS, @@ -156,7 +157,9 @@ class DSV41MonoLayer: def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"): _ensure_writable_flydsl_cache() - from .common.peer_memory import PeerBuffer + from vllm.models.common.specialized_decode.amd.mono.peer_memory import ( + PeerBuffer, + ) self.tp, self.rank = tp, rank self.d = Dims(tp) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py new file mode 100644 index 000000000000..fbf50ff495ab --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's kernel sources, in every build's JIT cache key +(``plan.key_tuple``): an edit to this package or to the shared AMD primitives +rebuilds.""" + +from pathlib import Path + +from vllm.models.common.specialized_decode.amd.mono.plan import SHARED, source_digest + +SOURCES = source_digest(Path(__file__).resolve().parent, SHARED) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/mono/stages/__init__.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/__init__.py diff --git a/vllm/models/deepseek_v41/amd/mono/stages/dims.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py similarity index 96% rename from vllm/models/deepseek_v41/amd/mono/stages/dims.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py index 6fd119c77206..9b743a7709cc 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/dims.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py @@ -8,7 +8,7 @@ from dataclasses import dataclass -from vllm.models.deepseek_v41.amd.mono.common.plan import Shard, padded +from vllm.models.common.specialized_decode.amd.mono.plan import Shard, padded # the model's full widths MOE_INTER = 2304 # a routed expert's intermediate width diff --git a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py similarity index 97% rename from vllm/models/deepseek_v41/amd/mono/stages/gemv.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py index 31a19ea4c72f..ce9714f0c23b 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py @@ -12,18 +12,18 @@ from flydsl.expr import const_expr, range_constexpr from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.mono.common.mx import ( +from vllm.models.common.specialized_decode.amd.mono.mx import ( UNIT_SCALE, mfma_scaled, ) -from vllm.models.deepseek_v41.amd.mono.common.ops import ( +from vllm.models.common.specialized_decode.amd.mono.ops import ( CM_NT, fresh, rsrc, traced, ) -from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES -from vllm.models.deepseek_v41.amd.mono.common.sync import POLL_MAX +from vllm.models.common.specialized_decode.amd.mono.plan import WAVES +from vllm.models.common.specialized_decode.amd.mono.sync import POLL_MAX ROWS = 16 # GEMV rows a task TILE = 16 # an MFMA's N: the tokens a GEMV pass takes (``token_tiles``) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py similarity index 98% rename from vllm/models/deepseek_v41/amd/mono/stages/moe.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py index 63f8947166d2..958676744cbe 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/moe.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py @@ -49,7 +49,7 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.mono.common.mx import ( +from vllm.models.common.specialized_decode.amd.mono.mx import ( FP4, FP8, FP8_MAX, @@ -61,7 +61,7 @@ rcp_pow2, scale_index, ) -from vllm.models.deepseek_v41.amd.mono.common.ops import ( +from vllm.models.common.specialized_decode.amd.mono.ops import ( CM_DEV, CM_NT, _hw_f32, @@ -84,20 +84,22 @@ sum_partials, traced, ) -from vllm.models.deepseek_v41.amd.mono.common.plan import ( +from vllm.models.common.specialized_decode.amd.mono.plan import ( BLOCKS, THREADS, WAVES, first_task, pair_layout, ) -from vllm.models.deepseek_v41.amd.mono.common.sync import preg, publish, sreg -from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan -from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART +from vllm.models.common.specialized_decode.amd.mono.sync import preg, publish, sreg +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( + moe_plan as mplan, +) +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.dims import UG_PART # helpers the kernel calls are imported by name: FlyDSL keys a build by the # source of the same-directory functions it names, never through a module -from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( ROWS, TILE, gemv_fp8_loads, @@ -110,18 +112,23 @@ tile_rows, token_tiles, ) -from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.moe_shape import ( MASK_WORDS, MoeBuild, route_shape, ) -from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS -from vllm.models.deepseek_v41.amd.mono.stages.router import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.router import ( + PARTS as ROUTER_PARTS, +) +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.router import ( router_tile_mfma, router_weight_loads, router_x_loads, ) -from vllm.models.deepseek_v41.amd.mono.stages.seam import HIDDEN, attn_peer_bytes +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.seam import ( + HIDDEN, + attn_peer_bytes, +) # a pick tile is a token tile's MFMA N: ``gather_x8r`` fills one from the other assert mplan.TILE == TILE diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py similarity index 96% rename from vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py index 5b5aa6fcdded..a46eeefc4948 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py @@ -15,7 +15,7 @@ by the kernel on traced Int32 and by the CPU property test on Python ints. """ -from vllm.models.deepseek_v41.amd.mono.common.plan import cdiv, imin +from vllm.models.common.specialized_decode.amd.mono.plan import cdiv, imin TILE = 16 # picks a pick tile: an MFMA's N diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py similarity index 93% rename from vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py index a07e88cfb064..b2d1eb20785c 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py @@ -12,11 +12,22 @@ import flydsl.expr as fx -from vllm.models.deepseek_v41.amd.mono.common.plan import BLOCKS, WAVES -from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan -from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART, Dims -from vllm.models.deepseek_v41.amd.mono.stages.gemv import ROWS, TILE, token_tiles -from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS +from vllm.models.common.specialized_decode.amd.mono.plan import BLOCKS, WAVES +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( + moe_plan as mplan, +) +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.dims import ( + UG_PART, + Dims, +) +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( + ROWS, + TILE, + token_tiles, +) +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.router import ( + PARTS as ROUTER_PARTS, +) EXPERTS, TOPK = 384, 6 # the target's routing (a build's: ``MoeBuild``) ROUTER0 = 8 # the router's first task: past the xq tasks (a token each) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/router.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py similarity index 95% rename from vllm/models/deepseek_v41/amd/mono/stages/router.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py index a7d6f411a816..967259733099 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/router.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py @@ -20,8 +20,8 @@ from flydsl.expr import const_expr, range_constexpr from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.mono.common.ops import mfma_bf16, rsrc -from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES +from vllm.models.common.specialized_decode.amd.mono.ops import mfma_bf16, rsrc +from vllm.models.common.specialized_decode.amd.mono.plan import WAVES HIDDEN = 5120 ROWS = 16 # a tile's rows: the MFMA's M diff --git a/vllm/models/deepseek_v41/amd/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py similarity index 97% rename from vllm/models/deepseek_v41/amd/mono/stages/seam.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py index ea73e2ef5435..76cb6dba307d 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/seam.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py @@ -31,7 +31,7 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.mono.common.ops import ( +from vllm.models.common.specialized_decode.amd.mono.ops import ( batched_rounds, bf16_round, bf_hi, @@ -45,9 +45,13 @@ sum_partials, traced, ) -from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES, pair_layout -from vllm.models.deepseek_v41.amd.mono.common.sync import preg -from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( +from vllm.models.common.specialized_decode.amd.mono.plan import ( + THREADS, + WAVES, + pair_layout, +) +from vllm.models.common.specialized_decode.amd.mono.sync import preg +from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( ld_bf, ld_f32, plus, From 0d7f3e718befa6b2f8831a1f156c8e399e259e4c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Sun, 18 Oct 2026 17:39:34 -0500 Subject: [PATCH 57/79] revert: move back common code to dsv41 dir Signed-off-by: simondanielsson --- .../test_specialized_decode.py | 2 +- .../common/specialized_decode/amd/__init__.py | 7 ------- .../specialized_decode/amd/mono/__init__.py | 14 ------------- .../mono/common/__init__.py | 9 +++++++++ .../amd/specialized_decode/mono/common}/mx.py | 2 +- .../specialized_decode/mono/common}/ops.py | 5 ++++- .../mono/common}/peer_memory.py | 0 .../specialized_decode/mono/common}/plan.py | 20 +++++++++---------- .../specialized_decode/mono/common}/sync.py | 2 +- .../amd/specialized_decode/mono/layer.py | 8 ++++---- .../amd/specialized_decode/mono/runner.py | 6 ++---- .../amd/specialized_decode/mono/sources.py | 11 +++++----- .../specialized_decode/mono/stages/dims.py | 5 ++++- .../specialized_decode/mono/stages/gemv.py | 8 ++++---- .../amd/specialized_decode/mono/stages/moe.py | 12 +++++++---- .../mono/stages/moe_plan.py | 2 +- .../mono/stages/moe_shape.py | 5 ++++- .../specialized_decode/mono/stages/router.py | 7 +++++-- .../specialized_decode/mono/stages/seam.py | 6 +++--- 19 files changed, 66 insertions(+), 65 deletions(-) delete mode 100644 vllm/models/common/specialized_decode/amd/__init__.py delete mode 100644 vllm/models/common/specialized_decode/amd/mono/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/specialized_decode/mono/common/__init__.py rename vllm/models/{common/specialized_decode/amd/mono => deepseek_v41/amd/specialized_decode/mono/common}/mx.py (97%) rename vllm/models/{common/specialized_decode/amd/mono => deepseek_v41/amd/specialized_decode/mono/common}/ops.py (99%) rename vllm/models/{common/specialized_decode/amd/mono => deepseek_v41/amd/specialized_decode/mono/common}/peer_memory.py (100%) rename vllm/models/{common/specialized_decode/amd/mono => deepseek_v41/amd/specialized_decode/mono/common}/plan.py (87%) rename vllm/models/{common/specialized_decode/amd/mono => deepseek_v41/amd/specialized_decode/mono/common}/sync.py (98%) diff --git a/tests/models/common/specialized_decode/test_specialized_decode.py b/tests/models/common/specialized_decode/test_specialized_decode.py index e663825e9a91..fd8f1d4c008b 100644 --- a/tests/models/common/specialized_decode/test_specialized_decode.py +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -286,7 +286,7 @@ def test_run_takes_the_planned_step_to_the_kernel(config, model, decode): def test_the_contract_imports_without_a_vendor_runtime(monkeypatch): """The library loads where neither FlyDSL nor aiter is installed, so a - vendor's shared kernel code must stay out of its import graph.""" + model's kernel code must stay out of its import graph.""" for name in list(sys.modules): if name.split(".")[0] in ("flydsl", "aiter") or "specialized_decode" in name: monkeypatch.delitem(sys.modules, name, raising=False) diff --git a/vllm/models/common/specialized_decode/amd/__init__.py b/vllm/models/common/specialized_decode/amd/__init__.py deleted file mode 100644 index 93c43d34f99e..000000000000 --- a/vllm/models/common/specialized_decode/amd/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kernel code AMD specialized decodes share, one subpackage a kernel flavor. - -Nothing here is imported by the specialized decode contract itself, which stays -free of vendor runtimes so that it loads on any platform. -""" diff --git a/vllm/models/common/specialized_decode/amd/mono/__init__.py b/vllm/models/common/specialized_decode/amd/mono/__init__.py deleted file mode 100644 index 8eb2e8ce3069..000000000000 --- a/vllm/models/common/specialized_decode/amd/mono/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, -# Copyright (C) 2026, Advanced Micro Devices, Inc.: -# atom/mono/__init__.py -"""Mechanisms shared by the monokernels, the FlyDSL persistent-kernel flavor of -specialized decode: ``ops`` (wave / math primitives, the TP group's peer -buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX arithmetic), ``plan`` -(execution model, layouts, build keys) and ``peer_memory`` (symmetric peer -buffers). - -A model's own kernels live under its ``specialized_decode/mono/``, and name this -package's ``plan.SHARED`` in their build key so that an edit here rebuilds them. -""" diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/__init__.py new file mode 100644 index 000000000000..6c5a829f28de --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/__init__.py +"""This model's mono kernels' shared mechanisms: ``ops`` (wave / math +primitives, the TP group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` +(FP8 / MX arithmetic), ``plan`` (execution model, layouts, build keys) and +``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/common/specialized_decode/amd/mono/mx.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/mx.py similarity index 97% rename from vllm/models/common/specialized_decode/amd/mono/mx.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/common/mx.py index 21cc405cb046..69a88bd082f0 100644 --- a/vllm/models/common/specialized_decode/amd/mono/mx.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/mx.py @@ -14,7 +14,7 @@ from flydsl.expr import rocdl from flydsl.expr.typing import T -from vllm.models.common.specialized_decode.amd.mono.ops import CM_NT +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import CM_NT FP8_MAX = 448.0 # the f8f6f4 MFMA's operand formats (cbsz / blgp) diff --git a/vllm/models/common/specialized_decode/amd/mono/ops.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/ops.py similarity index 99% rename from vllm/models/common/specialized_decode/amd/mono/ops.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/common/ops.py index 16cafbfdb57e..9a094b3e0f5a 100644 --- a/vllm/models/common/specialized_decode/amd/mono/ops.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/ops.py @@ -35,7 +35,10 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T, as_ir_value -from vllm.models.common.specialized_decode.amd.mono.plan import THREADS, WAVES +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( + THREADS, + WAVES, +) # ---------------------------------------------------------------- primitives # gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device diff --git a/vllm/models/common/specialized_decode/amd/mono/peer_memory.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/peer_memory.py similarity index 100% rename from vllm/models/common/specialized_decode/amd/mono/peer_memory.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/common/peer_memory.py diff --git a/vllm/models/common/specialized_decode/amd/mono/plan.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/plan.py similarity index 87% rename from vllm/models/common/specialized_decode/amd/mono/plan.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/common/plan.py index d37f5b2f3c4d..15ca9cdcb8af 100644 --- a/vllm/models/common/specialized_decode/amd/mono/plan.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/plan.py @@ -99,33 +99,33 @@ def padded(n: int, align: int) -> int: # ---------------------------------------------------------------- build keys _SCALARS = (int, float, bool, str, enum.Enum) -SHARED = Path(__file__).resolve().parent -"""This package, which every model's digest covers alongside its own.""" +_ROOT = Path(__file__).resolve().parents[1] # the mono package @functools.cache -def source_digest(*roots: Path) -> str: - """sha256 of every ``.py`` file under ``roots`` (files or directories), each - file's path within its root and its contents: one model's kernel sources, so - an edit to another model's rebuilds none of these. +def source_digest(*paths: str) -> str: + """sha256 of every ``.py`` file under ``paths`` (files or directories, + relative to the mono package), each file's path and contents: one model's + kernel sources, so an edit to another model's rebuilds none of these. Args: - roots: Absolute paths. A model passes its own package and ``SHARED``. + paths: Paths relative to the mono package. Returns: The digest's first 16 hex characters. Raises: - FileNotFoundError: A root holds no ``.py`` file. + FileNotFoundError: A path holds no ``.py`` file. """ h = hashlib.sha256() - for root in roots: + for rel in paths: + root = _ROOT / rel files = [root] if root.is_file() else sorted(root.rglob("*.py")) if not files: raise FileNotFoundError(f"no kernel sources at {root}") for f in files: - h.update(str(f.relative_to(root.parent)).encode()) + h.update(str(f.relative_to(_ROOT)).encode()) h.update(f.read_bytes()) return h.hexdigest()[:16] diff --git a/vllm/models/common/specialized_decode/amd/mono/sync.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/sync.py similarity index 98% rename from vllm/models/common/specialized_decode/amd/mono/sync.py rename to vllm/models/deepseek_v41/amd/specialized_decode/mono/common/sync.py index 5fcc355c32fc..be2276b580dc 100644 --- a/vllm/models/common/specialized_decode/amd/mono/sync.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/sync.py @@ -14,7 +14,7 @@ from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T -from vllm.models.common.specialized_decode.amd.mono.ops import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( CM_DEV, CM_SYS, bf16_pair, diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py index a3f5b6c40eb2..6b321baf0d6b 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py @@ -22,19 +22,19 @@ from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import Int32, Int64, T -from vllm.models.common.specialized_decode.amd.mono.mx import clamp_fp8 -from vllm.models.common.specialized_decode.amd.mono.ops import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.mx import clamp_fp8 +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( butterfly, fp8_pack4, peer_bases, traced, ) -from vllm.models.common.specialized_decode.amd.mono.plan import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( WAVES, first_task, key_tuple, ) -from vllm.models.common.specialized_decode.amd.mono.sync import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import ( Mailbox, preg, publish, diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py index 91bfdf9f0803..80a1c97ae15f 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py @@ -22,7 +22,7 @@ import torch -from vllm.models.common.specialized_decode.amd.mono.plan import BLOCKS +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import BLOCKS from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, Dims from .layer import ( @@ -157,9 +157,7 @@ class DSV41MonoLayer: def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"): _ensure_writable_flydsl_cache() - from vllm.models.common.specialized_decode.amd.mono.peer_memory import ( - PeerBuffer, - ) + from .common.peer_memory import PeerBuffer self.tp, self.rank = tp, rank self.d = Dims(tp) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py index fbf50ff495ab..1714374a9068 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The mono layer's kernel sources, in every build's JIT cache key -(``plan.key_tuple``): an edit to this package or to the shared AMD primitives -rebuilds.""" +(``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" -from pathlib import Path +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( + source_digest, +) -from vllm.models.common.specialized_decode.amd.mono.plan import SHARED, source_digest - -SOURCES = source_digest(Path(__file__).resolve().parent, SHARED) +SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py index 9b743a7709cc..f65f51041f4a 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py @@ -8,7 +8,10 @@ from dataclasses import dataclass -from vllm.models.common.specialized_decode.amd.mono.plan import Shard, padded +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( + Shard, + padded, +) # the model's full widths MOE_INTER = 2304 # a routed expert's intermediate width diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py index ce9714f0c23b..fbb5c257b696 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py @@ -12,18 +12,18 @@ from flydsl.expr import const_expr, range_constexpr from flydsl.expr.typing import T -from vllm.models.common.specialized_decode.amd.mono.mx import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.mx import ( UNIT_SCALE, mfma_scaled, ) -from vllm.models.common.specialized_decode.amd.mono.ops import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( CM_NT, fresh, rsrc, traced, ) -from vllm.models.common.specialized_decode.amd.mono.plan import WAVES -from vllm.models.common.specialized_decode.amd.mono.sync import POLL_MAX +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import WAVES +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import POLL_MAX ROWS = 16 # GEMV rows a task TILE = 16 # an MFMA's N: the tokens a GEMV pass takes (``token_tiles``) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py index 958676744cbe..d4dd6c3c2473 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py @@ -49,7 +49,7 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T -from vllm.models.common.specialized_decode.amd.mono.mx import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.mx import ( FP4, FP8, FP8_MAX, @@ -61,7 +61,7 @@ rcp_pow2, scale_index, ) -from vllm.models.common.specialized_decode.amd.mono.ops import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( CM_DEV, CM_NT, _hw_f32, @@ -84,14 +84,18 @@ sum_partials, traced, ) -from vllm.models.common.specialized_decode.amd.mono.plan import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( BLOCKS, THREADS, WAVES, first_task, pair_layout, ) -from vllm.models.common.specialized_decode.amd.mono.sync import preg, publish, sreg +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import ( + preg, + publish, + sreg, +) from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( moe_plan as mplan, ) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py index a46eeefc4948..c5606b264aa7 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py @@ -15,7 +15,7 @@ by the kernel on traced Int32 and by the CPU property test on Python ints. """ -from vllm.models.common.specialized_decode.amd.mono.plan import cdiv, imin +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import cdiv, imin TILE = 16 # picks a pick tile: an MFMA's N diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py index b2d1eb20785c..9de13e87b5b8 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py @@ -12,7 +12,10 @@ import flydsl.expr as fx -from vllm.models.common.specialized_decode.amd.mono.plan import BLOCKS, WAVES +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( + BLOCKS, + WAVES, +) from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( moe_plan as mplan, ) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py index 967259733099..f2fa69dc9bb2 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py @@ -20,8 +20,11 @@ from flydsl.expr import const_expr, range_constexpr from flydsl.expr.typing import T -from vllm.models.common.specialized_decode.amd.mono.ops import mfma_bf16, rsrc -from vllm.models.common.specialized_decode.amd.mono.plan import WAVES +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( + mfma_bf16, + rsrc, +) +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import WAVES HIDDEN = 5120 ROWS = 16 # a tile's rows: the MFMA's M diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py index 76cb6dba307d..aae238376d37 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py @@ -31,7 +31,7 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T -from vllm.models.common.specialized_decode.amd.mono.ops import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( batched_rounds, bf16_round, bf_hi, @@ -45,12 +45,12 @@ sum_partials, traced, ) -from vllm.models.common.specialized_decode.amd.mono.plan import ( +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( THREADS, WAVES, pair_layout, ) -from vllm.models.common.specialized_decode.amd.mono.sync import preg +from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import preg from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( ld_bf, ld_f32, From a45d50d2b2d6903ad5d04934997f8649a3c66863 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Sun, 18 Oct 2026 17:56:55 -0500 Subject: [PATCH 58/79] feat: allow kernel padding and sseparate width from kernel_width Signed-off-by: simondanielsson --- .../common/specialized_decode/reference.py | 2 +- .../test_specialized_decode.py | 36 ++++++++++++++++++- tests/models/test_dsv41_specialized_decode.py | 2 +- .../common/specialized_decode/dispatch.py | 17 ++++++--- .../common/specialized_decode/guards.py | 9 +++-- vllm/models/common/specialized_decode/spec.py | 25 ++++++++++--- .../specialized_decode/runtime/decode.py | 2 +- 7 files changed, 78 insertions(+), 15 deletions(-) diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py index 5eea922602b8..051abaa2dd45 100644 --- a/tests/models/common/specialized_decode/reference.py +++ b/tests/models/common/specialized_decode/reference.py @@ -75,7 +75,7 @@ def __init__(self, vllm_config, model: ReferenceModel, widths: tuple[int, ...]): self.compiled: set[Any] = set() def run(self, step: Step, input_ids, inputs_embeds, **kwargs): - self.launched.append(step.width) + self.launched.append(step.kernel_width) hidden = inputs_embeds for layer in self.model.layers: weight = self.weights.get(layer.index) diff --git a/tests/models/common/specialized_decode/test_specialized_decode.py b/tests/models/common/specialized_decode/test_specialized_decode.py index fd8f1d4c008b..27b7f207aafc 100644 --- a/tests/models/common/specialized_decode/test_specialized_decode.py +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -181,6 +181,17 @@ def test_capture_sizes_become_the_widths(config): assert config.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 16, 32] +def test_a_spec_can_keep_the_engines_capture_sizes(config): + """A spec that pads serves whatever sizes the deployment captures, so it + opts out of having them rewritten.""" + spec = replace(reference.REFERENCE, set_capture_sizes=False) + config.compilation_config.cudagraph_capture_sizes = [1, 3, 4, 16, 32] + + assert select_decode_spec(config, spec) is spec + + assert config.compilation_config.cudagraph_capture_sizes == [1, 3, 4, 16, 32] + + # Build @@ -253,10 +264,33 @@ def test_a_served_decode_step_takes_the_kernel(config, decode): with forward(config, CUDAGraphMode.FULL, {"prefill": False}): step = decode.plan(rows) assert isinstance(step, Step) - assert step.width == min(decode.widths) + assert step.width == step.kernel_width == min(decode.widths) assert step.positions is rows +def test_a_spec_that_pads_runs_a_short_step_at_the_next_width(config, model): + """``pad_steps`` serves a width the spec was not built for by launching at + the next one up, which is the width the kernel is handed.""" + spec = replace(reference.REFERENCE, pad_steps=True) + decode = SpecializedDecode(model, config, spec) + decode.bind_caches() + assert decode.widths == (1, 2, 4, 8) + + with forward(config, CUDAGraphMode.FULL, {"prefill": False}): + step = decode.plan(positions(3)) + over = decode.plan(positions(9)) + + assert step is not None and step.width == 3 and step.kernel_width == 4 + assert over is None + + +def test_without_padding_an_unserved_width_takes_vllms_path(config, decode): + """The default is exact: a width the spec was not built for falls back, + rather than silently running on a wider launch.""" + with forward(config, CUDAGraphMode.FULL, {"prefill": False}): + assert decode.plan(positions(3)) is None + + @pytest.mark.parametrize( "case", ["pipeline stage", "width", "piecewise", "unbound caches", "skip_step"] ) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 769c720da36f..491848bd7553 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -553,7 +553,7 @@ def launch_whole(plan, *args, **kwargs): decode.kernel.launch_whole = launch_whole - step = Step(width=1, positions=torch.zeros(1, dtype=torch.long)) + step = Step(width=1, kernel_width=1, positions=torch.zeros(1, dtype=torch.long)) out = decode.kernel.run(step, None, torch.zeros(1, HIDDEN)) assert out.shape == (1, HIDDEN) diff --git a/vllm/models/common/specialized_decode/dispatch.py b/vllm/models/common/specialized_decode/dispatch.py index 82968a1bded8..898cc5e7d831 100644 --- a/vllm/models/common/specialized_decode/dispatch.py +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -83,7 +83,7 @@ def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): if not agreed: raise DecodeUnsupported(f"{refused}: another TP rank refused") self.vllm_config = vllm_config - self.widths = spec.widths(vllm_config) + self.widths = tuple(sorted(spec.widths(vllm_config))) self._bound = False self.kernel: KernelT = spec.build(vllm_config, model, self.widths) kernel = self.kernel @@ -106,6 +106,13 @@ def bind_caches(self) -> None: self.kernel.bind_caches(caches) self._bound = True + def _kernel_width(self, width: int) -> int | None: + """The width the kernel launches at for a step of ``width`` rows, or + None when the spec serves no width for it.""" + if self.spec.pad_steps: + return next((w for w in self.widths if w >= width), None) + return width if width in self.widths else None + def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | None: """This step's plan, or None to run the model's own forward instead. @@ -113,10 +120,11 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non serve it, as described by the spec. """ width = positions.size(-1) + kernel_width = self._kernel_width(width) why: str | None if intermediate_tensors is not None: why = "pipeline stage input" - elif width not in self.widths: + elif kernel_width is None: why = "width" elif not is_forward_context_available(): why = "no forward context" @@ -125,13 +133,14 @@ def plan(self, positions: torch.Tensor, intermediate_tensors=None) -> Step | Non elif not self._bound: why = "caches not bound" else: - why = self.spec.skip_step(self.vllm_config, self.kernel, width) + why = self.spec.skip_step(self.vllm_config, self.kernel, kernel_width) if why is not None: logger.info_once( "Specialized decode %s: step on vLLM's path (%s)", self.spec.name, why ) return None - return Step(width=width, positions=positions) + assert kernel_width is not None + return Step(width=width, kernel_width=kernel_width, positions=positions) def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: """Run the specialized kernel, returning what the model's own forward would diff --git a/vllm/models/common/specialized_decode/guards.py b/vllm/models/common/specialized_decode/guards.py index 822053f50b8a..4d30aed75df4 100644 --- a/vllm/models/common/specialized_decode/guards.py +++ b/vllm/models/common/specialized_decode/guards.py @@ -129,7 +129,8 @@ def _select_spec( def _set_capture_sizes(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> None: - """Amend the capture sizes with the supported widths to avoid padding.""" + """Amend the capture sizes with the supported widths, so that no step is + refused for its width alone.""" cc = vllm_config.compilation_config widths = spec.widths(vllm_config) sizes = cc.cudagraph_capture_sizes or [] @@ -145,7 +146,8 @@ def select_decode_spec( vllm_config: "VllmConfig", specs: DecodeSpec[KernelT] | Sequence[DecodeSpec[KernelT]], ) -> DecodeSpec[KernelT] | None: - """Select the first compatible spec and set the capture sizes, if enabled. + """Select the first compatible spec and, unless it declines them, set the + capture sizes to its widths. Args: vllm_config: The engine's config, before the weights are loaded. @@ -161,7 +163,8 @@ def select_decode_spec( "Specialized decode refused: " + "; ".join(f"{n}: {', '.join(rs)}" for n, rs in refusals.items()) ) - _set_capture_sizes(spec, vllm_config) + if spec.set_capture_sizes: + _set_capture_sizes(spec, vllm_config) if vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: logger.warning( "Specialized decode %s under eager mode: for debugging only, serve " diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py index 2deafe490705..4359cc9fa147 100644 --- a/vllm/models/common/specialized_decode/spec.py +++ b/vllm/models/common/specialized_decode/spec.py @@ -48,7 +48,15 @@ class Step: """One decode step the kernel takes, as ``SpecializedDecode.plan`` decided it.""" width: int - """The step's rows, one of the widths the spec serves.""" + """The step's rows, as the model's forward was handed them.""" + + kernel_width: int + """The width the kernel launches at, one of the widths the spec serves. + + Equal to ``width``, unless the spec pads a short step up to the next width + it was built for, in which case it is the larger of the two and the kernel + owns the padding. + """ positions: "torch.Tensor" """The step's positions, as the model's forward was handed them.""" @@ -104,8 +112,7 @@ class DecodeSpec(Generic[KernelT]): name: Log and metric name. supports: Features this spec serves. Default deny: a feature in use that is absent from it refuses the deployment. - widths: The rows a launch serves, from the engine's config. These - become the CUDA graph capture sizes. + widths: The rows a launch serves, from the engine's config. refuse_config: Conditions under which the specialized decode kernel cannot run, based on the config. refuse_model: Conditions under which the specialized decode kernel cannot run, @@ -113,8 +120,16 @@ class DecodeSpec(Generic[KernelT]): refuse_caches: Conditions under which the specialized decode kernel cannot run, based on the cache. skip_step: Whether to run the specialized decode kernel for this forward call, - or to fall back to the model's own forward. + or to fall back to the model's own forward. Takes the step's + ``kernel_width``. build: This rank's kernel. + pad_steps: Whether a step of fewer rows than a served width runs on the + next width up, the kernel padding it. Default deny: a step whose + rows are not themselves a served width falls back. + set_capture_sizes: Whether to make the served widths the CUDA graph + capture sizes, so that no step is refused for its width alone. Turn + it off for a spec that pads, which serves the engine's own capture + sizes as they are. """ @@ -126,3 +141,5 @@ class DecodeSpec(Generic[KernelT]): refuse_caches: Callable[["VllmConfig", KernelT, Mapping[str, Any]], list[str]] skip_step: Callable[["VllmConfig", KernelT, int], str | None] build: Callable[["VllmConfig", Any, tuple[int, ...]], KernelT] + pad_steps: bool = False + set_capture_sizes: bool = True diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py index d62b1505d0dd..38deeab488e9 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py @@ -223,7 +223,7 @@ def step_plan(self, width: int) -> StepPlan | None: def run(self, step: Step, input_ids, inputs_embeds, **kwargs): """Normalized hidden states (and the aux streams) of a planned step.""" - plan = self.step_plan(step.width) + plan = self.step_plan(step.kernel_width) assert plan is not None x = inputs_embeds if x is None: From cea9c9e45444dac7876f769cd97904966b1b9fd4 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Sun, 18 Oct 2026 18:34:01 -0500 Subject: [PATCH 59/79] chore: resutrcture files for dsv41 to follow kimi Signed-off-by: simondanielsson --- tests/models/test_dsv41_specialized_decode.py | 4 +- vllm/models/deepseek_v41/amd/model.py | 5 +- .../amd/specialized_decode/__init__.py | 551 ------------------ .../specialized_decode/runtime/__init__.py | 5 + .../amd/specialized_decode/runtime/decode.py | 433 ++++++++++++++ .../amd/specialized_decode/spec.py | 142 +++++ 6 files changed, 586 insertions(+), 554 deletions(-) create mode 100644 vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py create mode 100644 vllm/models/deepseek_v41/amd/specialized_decode/spec.py diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 491848bd7553..98959d327562 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -16,7 +16,7 @@ from vllm.models.common.specialized_decode.dispatch import SpecializedDecode from vllm.models.common.specialized_decode.guards import select_decode_spec from vllm.models.common.specialized_decode.spec import DecodeUnsupported, Step -from vllm.models.deepseek_v41.amd import specialized_decode as sd +from vllm.models.deepseek_v41.amd.specialized_decode import spec as sd SPECS = (sd.DECODE_SPEC,) LAYERS = 6 @@ -267,7 +267,7 @@ def _metadata(model, **overrides) -> dict: class TestDsv41Spec(DecodeConformance): specs = SPECS - specs_module = "vllm.models.deepseek_v41.amd.specialized_decode" + specs_module = "vllm.models.deepseek_v41.amd.specialized_decode.spec" @pytest.fixture(autouse=True) def _stubs(self, monkeypatch): diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 10202a255339..5809f4830050 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -60,7 +60,7 @@ DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention -from vllm.models.deepseek_v41.amd.specialized_decode import DECODE_SPEC, DSV41Kernel +from vllm.models.deepseek_v41.amd.specialized_decode.spec import DECODE_SPEC from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors @@ -72,6 +72,9 @@ from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask if typing.TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.specialized_decode.runtime.decode import ( + DSV41Kernel, + ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata logger = init_logger(__name__) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py index 4a999c81180b..208f01a7cb5e 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py @@ -1,553 +1,2 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's specialized decode.""" - -from collections.abc import Mapping -from dataclasses import dataclass -from functools import partial -from typing import TYPE_CHECKING, Any, cast - -import torch - -import vllm.envs as envs -from vllm.forward_context import get_forward_context -from vllm.logger import init_logger -from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature, Step -from vllm.platforms import current_platform - -if TYPE_CHECKING: - from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer - from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterSparseSWAMetadata - -logger = init_logger(__name__) - -RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes -SWA_WIDTH = 128 # a causal decode token's window slots -MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) -# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats -CDNA_VERSIONS = (4,) - - -def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: - """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's - data rows, then its scale words: the kernels index both from it).""" - kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) - return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) - - -def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: - return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) - - -def _select( - layer: "DeepseekV4DecoderLayer", - positions: torch.Tensor, - normed: torch.Tensor, - qr: torch.Tensor, -) -> None: - """An index source's compressor and top-k selection, from K1's normed input - and bf16 qr: ``_forward_csa2_full`` / ``_forward_csa2_reindex`` and the - indexer op, serial.""" - attn = layer.attn - comp, indexer = attn.compressor, attn.indexer - assert indexer is not None - if comp is not None: - kv_score = torch.mm( - normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32 - ) - latent = comp(kv_score, positions) - comp.insert_cache(latent, positions, attn.rotary_emb) - indexer._produce_k(latent, positions, attn.indexer_rotary_emb) - weights, _ = indexer.weights_proj(normed) - q, q_scale, weights = indexer.forward_q( - qr, None, weights, positions, attn.indexer_rotary_emb - ) - assert weights is not None - indexer.indexer_op( - normed, (q, q_scale) if q_scale is not None else q, None, weights - ) - - -@dataclass(frozen=True) -class LayerPlan: - """One decoder layer's launch, settled at load. - - An ``ffn_only`` layer keeps vLLM's attention and its launch takes wo_b's - unreduced output, so its weights are all it has; the rest is what a - whole-layer launch reads off the layer. - """ - - layer: "DeepseekV4DecoderLayer" - weights: Any - ffn_only: bool - ratio: int = 0 - index: bool = False - swa_key: str | None = None - comp_key: str | None = None - folds: tuple[bool, ...] = () - - -class DSV41Kernel: - def __init__(self, vllm_config, model, widths: tuple[int, ...]): - from vllm.models.common.specialized_decode.guards import tp_group - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( - MAX_TOKENS, - DSV41MonoLayer, - ) - - assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS - rank, world, cpu_group = tp_group() - device = next(model.parameters()).device - self.model = model - self.runner = DSV41MonoLayer(world, rank, cpu_group, device) - self.embed_layer = getattr(model, "start_layer", 0) - # A decode step is never trimmed, so ``run`` owns every layer, the ones - # past the replay cut included. - self.layer_end = getattr(model, "end_layer", len(model.layers)) - self.widths = widths - self.swa_keys: list[str] = [] - self.comp_keys: list[str] = [] - self.plans: dict[int, LayerPlan] = {} - - tp_size = vllm_config.parallel_config.tensor_parallel_size - for idx in range(self.embed_layer, self.layer_end): - layer = model.layers[idx] - plan = self._plan_layer(layer, tp_size) - if plan is None: - continue - self.plans[layer.attn.layer_id] = plan - if plan.ffn_only: - continue - for key, keys in ( - (plan.swa_key, self.swa_keys), - (plan.comp_key, self.comp_keys), - ): - if key is not None and key not in keys: - keys.append(key) - self.cache_layers = tuple(sorted(set(self.swa_keys) | set(self.comp_keys))) - - def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: - """Info needed to run the specialized decode kernel for this layer.""" - attn = layer.attn - why = _skip_layer(layer) - if why is not None: - logger.info_once( - "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", - attn.layer_id, - why, - ) - return None - why = _ffn_only(layer) - weights = _weights(layer, why is not None, tp_size) - if why is not None: - logger.info_once( - "DSv4.1 specialized decode: the FFN launch for layers with %s", why - ) - # A decode step hands wo_b's TP partial to the FFN launch; the layer - # reduces it on the steps that keep vLLM's path. - attn.wo_b.reduce_results = False - return LayerPlan(layer, weights, ffn_only=True) - comp_key = attn.compressed_cache_prefix - # K1's seam folds the owed post into the residual, except where the - # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == self.embed_layer or layer.engram is not None - return LayerPlan( - layer, - weights, - ffn_only=False, - ratio=attn.compress_ratio, - index=attn.indexer is not None, - swa_key=attn.swa_cache_layer.prefix, - comp_key=comp_key, - folds=(True, False) if settled else (True,), - ) - - def bind_caches(self, caches: Mapping[str, Any]) -> None: - """Nothing to take: a launch reads the cache off its attention layer.""" - - def warmup_keys(self) -> list[Any]: - keys: dict[Any, None] = {} - for plan in self.plans.values(): - for width in self.widths: - if plan.ffn_only: - keys[width, "ffn"] = None - else: - for fold in plan.folds: - keys[width, plan.ratio, fold, plan.index] = None - return list(keys) - - def compile(self, key: Any) -> None: - self.runner.scratch(key[0]) - if key[1] == "ffn": - self.runner.ffn_kernel(key[0]) - else: - width, ratio, fold, index = key - self.runner.kernels(width, ratio, fold=fold, index=index) - - def launch_whole( - self, - plan: LayerPlan, - x: torch.Tensor, - positions: torch.Tensor, - residual: torch.Tensor | None, - post_mix: torch.Tensor | None, - res_mix: torch.Tensor | None, - pre_mix: torch.Tensor | None, - engram_hashes: torch.Tensor | None = None, - engram_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, ...] | None: - """Run a decode layer as one logical kernel.""" - layer, runner = plan.layer, self.runner - attn = layer.attn - first = residual is None - if first and x.dim() != 2: - return None - tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) - if not all(t is not None and t.is_contiguous() for t in tensors): - return None - fold = True - if first: - residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() - pre_mix = runner.first_pre_mix(x.shape[0]) - post_mix = res_mix = pre_mix # unread without the fold - fold = False - elif layer.engram is not None and engram_hashes is not None: - residual = layer.mhc_post(x, residual, post_mix, res_mix) - hashes = engram_hashes[:, layer.engram.layer_hash_index] - residual = layer.engram(residual, hashes, engram_mask) - fold = False - md = cast("dict[str, Any]", get_forward_context().attn_metadata) - swa = md[cast("str", plan.swa_key)] - kw: dict[str, Any] = {} - if plan.index: - kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) - if plan.ratio: - comp = md[cast("str", plan.comp_key)] - kw.update( - topk_indices=attn.topk_indices_buffer, - comp_cache=_record_view( - attn._compressed_kv_cache(), comp.block_size // plan.ratio - ), - comp_block_table=comp.block_table, - ) - return runner.forward( - plan.weights, - x, - residual, - post_mix, - res_mix, - pre_mix, - positions, - swa.slot_mapping, - _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), - swa.decode_swa_indices, - swa.decode_swa_lens, - swa.token_to_req_indices, - fold=fold, - **kw, - ) - - def launch_ffn( - self, - plan: LayerPlan, - part: torch.Tensor, - residual: torch.Tensor, - post_mix: torch.Tensor, - res_mix: torch.Tensor, - pre_mix: torch.Tensor, - ) -> tuple[torch.Tensor, ...] | None: - """The FFN seam, MoE and both reductions as one launch, from wo_b's - unreduced output after vLLM's attention.""" - tensors = (part, residual, post_mix, res_mix, pre_mix) - if not all(t.is_contiguous() for t in tensors): - return None - return self.runner.ffn(plan.weights, *tensors) - - def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - model, positions = self.model, step.positions - ( - hidden, - input_ids, - pre_mix, - engram_hashes, - engram_mask, - full_num_tokens, - ) = model.prologue( - input_ids, positions, None, inputs_embeds, kwargs.get("lookback_token_ids") - ) - residual = post_mix = res_mix = None - aux_hidden_states: list[torch.Tensor] = [] - final_aux_recon = None - for idx in range(self.embed_layer, self.layer_end): - layer = model.layers[idx] - plan = self.plans.get(layer.attn.layer_id) - ffn = out = None - if plan is None: - # TODO: check if we should crash here - pass - elif plan.ffn_only: - ffn = partial(self.launch_ffn, plan) - else: - # Run entire layer - out = self.launch_whole( - plan, - hidden, - positions, - residual, - post_mix, - res_mix, - pre_mix, - engram_hashes, - engram_mask, - ) - if out is None: - # Run standard model implementation, with injected specialized ffn - out = layer( - hidden, - positions, - input_ids, - pre_mix, - post_mix, - res_mix, - residual, - engram_hashes, - engram_mask, - decode_ffn=ffn, - ) - hidden, residual, post_mix, res_mix, pre_mix = out - aux_recon = model.after_layer( - idx, - hidden, - residual, - post_mix, - res_mix, - full_num_tokens, - aux_hidden_states, - replay=False, - ) - if aux_recon is not None: - final_aux_recon = aux_recon - return model.epilogue( - hidden, - positions, - input_ids, - pre_mix, - post_mix, - res_mix, - residual, - aux_hidden_states, - final_aux_recon, - self.layer_end > self.embed_layer, - full_num_tokens, - replay=False, - ) - - -def _ffn_only(layer) -> str | None: - """Why this layer's attention seam and attention stay on vLLM's path, which - leaves it the FFN launch, or None when it runs whole.""" - from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer - - attn = layer.attn - if attn.compressor is not None and attn.indexer is None: - return "a compressor without an indexer" - if attn.indexer is not None and not isinstance( - attn.indexer, DeepseekV41RocmMxfp4Indexer - ): - return f"a {type(attn.indexer).__name__}" - if attn.compress_ratio not in (0, 1, 2): - return f"compress ratio {attn.compress_ratio}" - if attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": - return f"a {attn.kv_cache_dtype} KV cache" - return None - - -def _widths(vllm_config) -> tuple[int, ...]: - """The configured capture sizes the kernels can serve, so a deployment - that asked for fewer graphs keeps them; every size at all when it asked for - none.""" - sizes = vllm_config.compilation_config.cudagraph_capture_sizes or [] - served = tuple(s for s in sorted(sizes) if 1 <= s <= MAX_ROWS) - return served or tuple(range(1, MAX_ROWS + 1)) - - -def _refuse_config(vllm_config) -> list[str]: - """Every reason this config cannot run the kernels, before load.""" - reasons = [] - if vllm_config.parallel_config.tensor_parallel_size not in (2, 4): - reasons.append("needs tensor parallel size 2 or 4") - if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: - reasons.append( - "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" - ) - try: - import vllm.models.deepseek_v41.amd.specialized_decode.mono.runner # noqa: F401 - except ImportError as err: - reasons.append(f"needs FlyDSL and AITER's FlyDSL helpers ({err})") - return reasons - - -def _refuse_model(model) -> list[str]: - """Every reason the loaded model and this GPU cannot run the kernels. The - kernels read every one of the 384 experts on each rank, TP-sharded, in - AITER's A8W4 MXFP4 layout, from CTAs that all stay resident.""" - from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import BLOCKS - from vllm.platforms.rocm import get_cdna_version - - reasons = [] - if get_cdna_version() not in CDNA_VERSIONS: - reasons.append(f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}") - cus = current_platform.num_compute_units(torch.accelerator.current_device_index()) - if cus < BLOCKS: - reasons.append(f"needs {BLOCKS} compute units, the GPU has {cus}") - backends = set() - for layer in model.layers: - experts = getattr(getattr(layer, "ffn", None), "experts", None) - routed = getattr(experts, "routed_experts", None) - if routed is not None: - backends.add(getattr(routed.quant_method, "mxfp4_backend", None)) - wrong = backends - {Mxfp4MoeBackend.AITER_MXFP4_BF16} - if wrong: - reasons.append(f"needs the AITER_MXFP4_BF16 MoE backend, not {wrong}") - return reasons - - -def _skip_layer(layer) -> str | None: - """Why this layer keeps vLLM's path for every step, or None when it takes a - launch (whole or FFN only: ``_plan_layer`` picks which).""" - ffn = layer.ffn - if layer.use_sequence_parallel: - return "sequence-parallel layer" - if not layer.fuse_seam_norm: - return "needs aiter's fused mHC seam" - if ( - ffn.shared_experts is None - or ffn.gate.tid2eid is not None - or ffn.n_routed_experts != 384 - or ffn.n_activated_experts != 6 - or ffn.routed_scaling_factor != 1.5 - or ffn.swiglu_limit != 10.0 - or ffn.scoring_func != "sqrtsoftplus" - or not ffn.renormalize - ): - return "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" - return None - - -def _skip_step(vllm_config, decode: DSV41Kernel, width: int) -> str | None: - """Why this step keeps vLLM's path at every layer, or None when the kernels - take it. One verdict for every kernel layer, from the batch the whole TP - group sees alike.""" - md = get_forward_context().attn_metadata - if not decode.plans: - return "no layer takes a launch" - if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches - return "no attention metadata" - replay = decode.model.decoder_replay_layers - if replay is not None and replay.uses_replay_batch(): - return "replay batch step" - if any(md.get(k) is None for k in decode.comp_keys): - return "layer metadata missing" - for key in decode.swa_keys: - swa = cast("DeepseekV41ROCMAiterSparseSWAMetadata | None", md.get(key)) - if swa is None: - return "layer metadata missing" - if swa.num_prefills != 0 or swa.num_decodes == 0: - return "not a decode-only step" - if swa.num_decode_tokens > width: - return "more decode tokens than rows" - if swa.decode_swa_width != SWA_WIDTH: - return f"SWA width {swa.decode_swa_width}" - if ( - swa.decode_swa_indices is None - or swa.decode_swa_lens is None - or swa.token_to_req_indices is None - ): - return "SWA decode metadata missing" - return None - - -def _weights(layer, ffn_only: bool, world: int): - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( - MonoLayerWeights, - ) - - a, f = layer.attn, layer.ffn - e, sh = f.experts.routed_experts, f.shared_experts - weights = MonoLayerWeights( - attn=None if ffn_only else _attn_weights(a, world), - hc_attn_fn=layer.hc_attn_fn, - hc_attn_scale=layer.hc_attn_scale, - hc_attn_base=layer.hc_attn_base, - attn_norm=layer.attn_norm.weight, - hc_ffn_fn=layer.hc_ffn_fn, - hc_ffn_scale=layer.hc_ffn_scale, - hc_ffn_base=layer.hc_ffn_base, - ffn_norm=layer.ffn_norm.weight, - gate_w=f.gate.weight, - bias=f.gate.e_score_correction_bias, - w13=e.w13_weight, - w13_s=e.w13_weight_scale, - w2=e.w2_weight, - w2_s=e.w2_weight_scale, - sgu=sh.gate_up_proj.weight, - sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), - sw2=sh.down_proj.weight, - sw2_s=_scale_bytes(sh.down_proj.weight_scale), - ) - weights.check(world) - return weights - - -def _attn_weights(a, world: int): - from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import Dims - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import AttnWeights - - attn = AttnWeights( - layer_id=a.layer_id, - wqkv=a.fused_wqa_wkv.weight, - wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), - q_norm=a.q_norm.weight, - kv_norm=a.kv_norm.weight, - wq_b=a.wq_b.weight, - wq_b_scale=_scale_bytes(a.wq_b.weight_scale), - wo_a=a.wo_a.weight, - wo_a_scale=_scale_bytes(a.wo_a.weight_scale), - wo_b=a.wo_b.weight, - wo_b_scale=_scale_bytes(a.wo_b.weight_scale), - attn_sink=a.attn_sink, - cos_sin=a.rotary_emb.cos_sin_cache, - ratio=a.compress_ratio, - ) - attn.check(Dims(world)) - return attn - - -def _refuse_caches(vllm_config, decode: DSV41Kernel, caches) -> list[str]: - """Every reason the kernels cannot read these caches: they index each as - [blocks, block, RECORD] bytes.""" - reasons = [] - for name, cache in caches.items(): - if not isinstance(cache, torch.Tensor): - reasons.append(f"{name}: no cache tensor") - continue - if cache.element_size() != 1: - reasons.append(f"{name}: {cache.dtype} records, needs bytes") - if cache.dim() < 2: - reasons.append(f"{name}: {cache.dim()} dims, needs [blocks, block, bytes]") - elif cache.stride(0) < RECORD: - reasons.append(f"{name}: block stride {cache.stride(0)} under {RECORD}") - return reasons - - -DECODE_SPEC = DecodeSpec[DSV41Kernel]( - name="dsv41-cdna4", - supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), - widths=_widths, - refuse_config=_refuse_config, - refuse_model=_refuse_model, - refuse_caches=_refuse_caches, - skip_step=_skip_step, - build=DSV41Kernel, -) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py new file mode 100644 index 000000000000..2fc7c032f0b6 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1 specialized decode: a FlyDSL decode launch a layer over the +stock backbone, enabled by ``--specialized-decode``. What the deployment has to +satisfy to run it is declared in ``..spec``.""" diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py new file mode 100644 index 000000000000..1f8db64687f6 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py @@ -0,0 +1,433 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1's decode kernel: one launch a layer over the stock backbone.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING, Any, cast + +import torch + +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.models.common.specialized_decode.spec import Step + +from ..spec import MAX_ROWS, RECORD + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + +logger = init_logger(__name__) + + +def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: + """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's + data rows, then its scale words: the kernels index both from it).""" + kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) + + +def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _select( + layer: "DeepseekV4DecoderLayer", + positions: torch.Tensor, + normed: torch.Tensor, + qr: torch.Tensor, +) -> None: + """An index source's compressor and top-k selection, from K1's normed input + and bf16 qr: ``_forward_csa2_full`` / ``_forward_csa2_reindex`` and the + indexer op, serial.""" + attn = layer.attn + comp, indexer = attn.compressor, attn.indexer + assert indexer is not None + if comp is not None: + kv_score = torch.mm( + normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32 + ) + latent = comp(kv_score, positions) + comp.insert_cache(latent, positions, attn.rotary_emb) + indexer._produce_k(latent, positions, attn.indexer_rotary_emb) + weights, _ = indexer.weights_proj(normed) + q, q_scale, weights = indexer.forward_q( + qr, None, weights, positions, attn.indexer_rotary_emb + ) + assert weights is not None + indexer.indexer_op( + normed, (q, q_scale) if q_scale is not None else q, None, weights + ) + + +@dataclass(frozen=True) +class LayerPlan: + """One decoder layer's launch, settled at load. + + An ``ffn_only`` layer keeps vLLM's attention and its launch takes wo_b's + unreduced output, so its weights are all it has; the rest is what a + whole-layer launch reads off the layer. + """ + + layer: "DeepseekV4DecoderLayer" + weights: Any + ffn_only: bool + ratio: int = 0 + index: bool = False + swa_key: str | None = None + comp_key: str | None = None + folds: tuple[bool, ...] = () + + +class DSV41Kernel: + def __init__(self, vllm_config, model, widths: tuple[int, ...]): + from vllm.models.common.specialized_decode.guards import tp_group + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + MAX_TOKENS, + DSV41MonoLayer, + ) + + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + rank, world, cpu_group = tp_group() + device = next(model.parameters()).device + self.model = model + self.runner = DSV41MonoLayer(world, rank, cpu_group, device) + self.embed_layer = getattr(model, "start_layer", 0) + # A decode step is never trimmed, so ``run`` owns every layer, the ones + # past the replay cut included. + self.layer_end = getattr(model, "end_layer", len(model.layers)) + self.widths = widths + self.swa_keys: list[str] = [] + self.comp_keys: list[str] = [] + self.plans: dict[int, LayerPlan] = {} + + tp_size = vllm_config.parallel_config.tensor_parallel_size + for idx in range(self.embed_layer, self.layer_end): + layer = model.layers[idx] + plan = self._plan_layer(layer, tp_size) + if plan is None: + continue + self.plans[layer.attn.layer_id] = plan + if plan.ffn_only: + continue + for key, keys in ( + (plan.swa_key, self.swa_keys), + (plan.comp_key, self.comp_keys), + ): + if key is not None and key not in keys: + keys.append(key) + self.cache_layers = tuple(sorted(set(self.swa_keys) | set(self.comp_keys))) + + def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: + """Info needed to run the specialized decode kernel for this layer.""" + attn = layer.attn + why = _skip_layer(layer) + if why is not None: + logger.info_once( + "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", + attn.layer_id, + why, + ) + return None + why = _ffn_only(layer) + weights = _weights(layer, why is not None, tp_size) + if why is not None: + logger.info_once( + "DSv4.1 specialized decode: the FFN launch for layers with %s", why + ) + # A decode step hands wo_b's TP partial to the FFN launch; the layer + # reduces it on the steps that keep vLLM's path. + attn.wo_b.reduce_results = False + return LayerPlan(layer, weights, ffn_only=True) + comp_key = attn.compressed_cache_prefix + # K1's seam folds the owed post into the residual, except where the + # residual arrives settled: the embedding, or after the Engram write. + settled = attn.layer_id == self.embed_layer or layer.engram is not None + return LayerPlan( + layer, + weights, + ffn_only=False, + ratio=attn.compress_ratio, + index=attn.indexer is not None, + swa_key=attn.swa_cache_layer.prefix, + comp_key=comp_key, + folds=(True, False) if settled else (True,), + ) + + def bind_caches(self, caches: Mapping[str, Any]) -> None: + """Nothing to take: a launch reads the cache off its attention layer.""" + + def warmup_keys(self) -> list[Any]: + keys: dict[Any, None] = {} + for plan in self.plans.values(): + for width in self.widths: + if plan.ffn_only: + keys[width, "ffn"] = None + else: + for fold in plan.folds: + keys[width, plan.ratio, fold, plan.index] = None + return list(keys) + + def compile(self, key: Any) -> None: + self.runner.scratch(key[0]) + if key[1] == "ffn": + self.runner.ffn_kernel(key[0]) + else: + width, ratio, fold, index = key + self.runner.kernels(width, ratio, fold=fold, index=index) + + def launch_whole( + self, + plan: LayerPlan, + x: torch.Tensor, + positions: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + pre_mix: torch.Tensor | None, + engram_hashes: torch.Tensor | None = None, + engram_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, ...] | None: + """Run a decode layer as one logical kernel.""" + layer, runner = plan.layer, self.runner + attn = layer.attn + first = residual is None + if first and x.dim() != 2: + return None + tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) + if not all(t is not None and t.is_contiguous() for t in tensors): + return None + fold = True + if first: + residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + pre_mix = runner.first_pre_mix(x.shape[0]) + post_mix = res_mix = pre_mix # unread without the fold + fold = False + elif layer.engram is not None and engram_hashes is not None: + residual = layer.mhc_post(x, residual, post_mix, res_mix) + hashes = engram_hashes[:, layer.engram.layer_hash_index] + residual = layer.engram(residual, hashes, engram_mask) + fold = False + md = cast("dict[str, Any]", get_forward_context().attn_metadata) + swa = md[cast("str", plan.swa_key)] + kw: dict[str, Any] = {} + if plan.index: + kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) + if plan.ratio: + comp = md[cast("str", plan.comp_key)] + kw.update( + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // plan.ratio + ), + comp_block_table=comp.block_table, + ) + return runner.forward( + plan.weights, + x, + residual, + post_mix, + res_mix, + pre_mix, + positions, + swa.slot_mapping, + _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), + swa.decode_swa_indices, + swa.decode_swa_lens, + swa.token_to_req_indices, + fold=fold, + **kw, + ) + + def launch_ffn( + self, + plan: LayerPlan, + part: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + pre_mix: torch.Tensor, + ) -> tuple[torch.Tensor, ...] | None: + """The FFN seam, MoE and both reductions as one launch, from wo_b's + unreduced output after vLLM's attention.""" + tensors = (part, residual, post_mix, res_mix, pre_mix) + if not all(t.is_contiguous() for t in tensors): + return None + return self.runner.ffn(plan.weights, *tensors) + + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + model, positions = self.model, step.positions + ( + hidden, + input_ids, + pre_mix, + engram_hashes, + engram_mask, + full_num_tokens, + ) = model.prologue( + input_ids, positions, None, inputs_embeds, kwargs.get("lookback_token_ids") + ) + residual = post_mix = res_mix = None + aux_hidden_states: list[torch.Tensor] = [] + final_aux_recon = None + for idx in range(self.embed_layer, self.layer_end): + layer = model.layers[idx] + plan = self.plans.get(layer.attn.layer_id) + ffn = out = None + if plan is None: + # TODO: check if we should crash here + pass + elif plan.ffn_only: + ffn = partial(self.launch_ffn, plan) + else: + # Run entire layer + out = self.launch_whole( + plan, + hidden, + positions, + residual, + post_mix, + res_mix, + pre_mix, + engram_hashes, + engram_mask, + ) + if out is None: + # Run standard model implementation, with injected specialized ffn + out = layer( + hidden, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + decode_ffn=ffn, + ) + hidden, residual, post_mix, res_mix, pre_mix = out + aux_recon = model.after_layer( + idx, + hidden, + residual, + post_mix, + res_mix, + full_num_tokens, + aux_hidden_states, + replay=False, + ) + if aux_recon is not None: + final_aux_recon = aux_recon + return model.epilogue( + hidden, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_states, + final_aux_recon, + self.layer_end > self.embed_layer, + full_num_tokens, + replay=False, + ) + + +def _ffn_only(layer) -> str | None: + """Why this layer's attention seam and attention stay on vLLM's path, which + leaves it the FFN launch, or None when it runs whole.""" + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + attn = layer.attn + if attn.compressor is not None and attn.indexer is None: + return "a compressor without an indexer" + if attn.indexer is not None and not isinstance( + attn.indexer, DeepseekV41RocmMxfp4Indexer + ): + return f"a {type(attn.indexer).__name__}" + if attn.compress_ratio not in (0, 1, 2): + return f"compress ratio {attn.compress_ratio}" + if attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": + return f"a {attn.kv_cache_dtype} KV cache" + return None + + +def _skip_layer(layer) -> str | None: + """Why this layer keeps vLLM's path for every step, or None when it takes a + launch (whole or FFN only: ``_plan_layer`` picks which).""" + ffn = layer.ffn + if layer.use_sequence_parallel: + return "sequence-parallel layer" + if not layer.fuse_seam_norm: + return "needs aiter's fused mHC seam" + if ( + ffn.shared_experts is None + or ffn.gate.tid2eid is not None + or ffn.n_routed_experts != 384 + or ffn.n_activated_experts != 6 + or ffn.routed_scaling_factor != 1.5 + or ffn.swiglu_limit != 10.0 + or ffn.scoring_func != "sqrtsoftplus" + or not ffn.renormalize + ): + return "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + return None + + +def _weights(layer, ffn_only: bool, world: int): + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + MonoLayerWeights, + ) + + a, f = layer.attn, layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + weights = MonoLayerWeights( + attn=None if ffn_only else _attn_weights(a, world), + hc_attn_fn=layer.hc_attn_fn, + hc_attn_scale=layer.hc_attn_scale, + hc_attn_base=layer.hc_attn_base, + attn_norm=layer.attn_norm.weight, + hc_ffn_fn=layer.hc_ffn_fn, + hc_ffn_scale=layer.hc_ffn_scale, + hc_ffn_base=layer.hc_ffn_base, + ffn_norm=layer.ffn_norm.weight, + gate_w=f.gate.weight, + bias=f.gate.e_score_correction_bias, + w13=e.w13_weight, + w13_s=e.w13_weight_scale, + w2=e.w2_weight, + w2_s=e.w2_weight_scale, + sgu=sh.gate_up_proj.weight, + sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), + sw2=sh.down_proj.weight, + sw2_s=_scale_bytes(sh.down_proj.weight_scale), + ) + weights.check(world) + return weights + + +def _attn_weights(a, world: int): + from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import AttnWeights + + attn = AttnWeights( + layer_id=a.layer_id, + wqkv=a.fused_wqa_wkv.weight, + wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=a.wq_b.weight, + wq_b_scale=_scale_bytes(a.wq_b.weight_scale), + wo_a=a.wo_a.weight, + wo_a_scale=_scale_bytes(a.wo_a.weight_scale), + wo_b=a.wo_b.weight, + wo_b_scale=_scale_bytes(a.wo_b.weight_scale), + attn_sink=a.attn_sink, + cos_sin=a.rotary_emb.cos_sin_cache, + ratio=a.compress_ratio, + ) + attn.check(Dims(world)) + return attn diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py new file mode 100644 index 000000000000..31d123f00521 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1's specialized decode.""" + +from typing import TYPE_CHECKING, cast + +import torch + +import vllm.envs as envs +from vllm.forward_context import get_forward_context +from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterSparseSWAMetadata + + from .runtime.decode import DSV41Kernel + +RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes +SWA_WIDTH = 128 # a causal decode token's window slots +MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats +CDNA_VERSIONS = (4,) + + +def _widths(vllm_config) -> tuple[int, ...]: + """The configured capture sizes the kernels can serve, so a deployment + that asked for fewer graphs keeps them; every size at all when it asked for + none.""" + sizes = vllm_config.compilation_config.cudagraph_capture_sizes or [] + served = tuple(s for s in sorted(sizes) if 1 <= s <= MAX_ROWS) + return served or tuple(range(1, MAX_ROWS + 1)) + + +def _refuse_config(vllm_config) -> list[str]: + """Every reason this config cannot run the kernels, before load.""" + reasons = [] + if vllm_config.parallel_config.tensor_parallel_size not in (2, 4): + reasons.append("needs tensor parallel size 2 or 4") + if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + reasons.append( + "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" + ) + try: + import vllm.models.deepseek_v41.amd.specialized_decode.mono.runner # noqa: F401 + except ImportError as err: + reasons.append(f"needs FlyDSL and AITER's FlyDSL helpers ({err})") + return reasons + + +def _refuse_model(model) -> list[str]: + """Every reason the loaded model and this GPU cannot run the kernels. The + kernels read every one of the 384 experts on each rank, TP-sharded, in + AITER's A8W4 MXFP4 layout, from CTAs that all stay resident.""" + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import BLOCKS + from vllm.platforms.rocm import get_cdna_version + + reasons = [] + if get_cdna_version() not in CDNA_VERSIONS: + reasons.append(f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}") + cus = current_platform.num_compute_units(torch.accelerator.current_device_index()) + if cus < BLOCKS: + reasons.append(f"needs {BLOCKS} compute units, the GPU has {cus}") + backends = set() + for layer in model.layers: + experts = getattr(getattr(layer, "ffn", None), "experts", None) + routed = getattr(experts, "routed_experts", None) + if routed is not None: + backends.add(getattr(routed.quant_method, "mxfp4_backend", None)) + wrong = backends - {Mxfp4MoeBackend.AITER_MXFP4_BF16} + if wrong: + reasons.append(f"needs the AITER_MXFP4_BF16 MoE backend, not {wrong}") + return reasons + + +def _refuse_caches(vllm_config, decode: "DSV41Kernel", caches) -> list[str]: + """Every reason the kernels cannot read these caches: they index each as + [blocks, block, RECORD] bytes.""" + reasons = [] + for name, cache in caches.items(): + if not isinstance(cache, torch.Tensor): + reasons.append(f"{name}: no cache tensor") + continue + if cache.element_size() != 1: + reasons.append(f"{name}: {cache.dtype} records, needs bytes") + if cache.dim() < 2: + reasons.append(f"{name}: {cache.dim()} dims, needs [blocks, block, bytes]") + elif cache.stride(0) < RECORD: + reasons.append(f"{name}: block stride {cache.stride(0)} under {RECORD}") + return reasons + + +def _skip_step(vllm_config, decode: "DSV41Kernel", width: int) -> str | None: + """Why this step keeps vLLM's path at every layer, or None when the kernels + take it. One verdict for every kernel layer, from the batch the whole TP + group sees alike.""" + md = get_forward_context().attn_metadata + if not decode.plans: + return "no layer takes a launch" + if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + return "no attention metadata" + replay = decode.model.decoder_replay_layers + if replay is not None and replay.uses_replay_batch(): + return "replay batch step" + if any(md.get(k) is None for k in decode.comp_keys): + return "layer metadata missing" + for key in decode.swa_keys: + swa = cast("DeepseekV41ROCMAiterSparseSWAMetadata | None", md.get(key)) + if swa is None: + return "layer metadata missing" + if swa.num_prefills != 0 or swa.num_decodes == 0: + return "not a decode-only step" + if swa.num_decode_tokens > width: + return "more decode tokens than rows" + if swa.decode_swa_width != SWA_WIDTH: + return f"SWA width {swa.decode_swa_width}" + if ( + swa.decode_swa_indices is None + or swa.decode_swa_lens is None + or swa.token_to_req_indices is None + ): + return "SWA decode metadata missing" + return None + + +def _build(vllm_config, model, widths: tuple[int, ...]) -> "DSV41Kernel": + from .runtime.decode import DSV41Kernel + + return DSV41Kernel(vllm_config, model, widths) + + +DECODE_SPEC = DecodeSpec["DSV41Kernel"]( + name="dsv41-cdna4", + supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), + widths=_widths, + refuse_config=_refuse_config, + refuse_model=_refuse_model, + refuse_caches=_refuse_caches, + skip_step=_skip_step, + build=_build, +) From 307e2e7e2e3ed06b98d0188e67379887dd02cea7 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 19 Oct 2026 10:17:13 -0500 Subject: [PATCH 60/79] feat: port dspark monokernel Signed-off-by: simondanielsson --- tests/models/test_dsv41_specialized_decode.py | 169 +++++++++++++++++- vllm/models/deepseek_v41/amd/dspark.py | 92 ++++++++-- vllm/models/deepseek_v41/amd/model_state.py | 8 +- .../amd/specialized_decode/runtime/decode.py | 148 +++++++++++++-- .../amd/specialized_decode/spec.py | 100 +++++++++-- 5 files changed, 463 insertions(+), 54 deletions(-) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 98959d327562..5832d5ed8280 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -25,8 +25,9 @@ class _StubRunner: - def __init__(self, world, rank, group, device): + def __init__(self, world, rank, group, device, **routing): self.world = world + self.routing = routing self.built: list = [] self.epoch = torch.zeros(1, dtype=torch.int32) @@ -567,6 +568,172 @@ def launch_whole(plan, *args, **kwargs): ] +DRAFT_LAYERS = 2 +DRAFT_EXPERTS = 128 +DRAFT_TOPK = 3 +DRAFT_TOKENS = 5 + + +class _DraftModel(torch.nn.Module): + """The DSpark draft head's backbone: its own layers, routing and stream.""" + + def __init__(self, config): + super().__init__() + self.config = config.model_config.hf_config + self.vllm_config = config + self.layers = torch.nn.ModuleList( + _Layer( + LAYERS + i, + attn=dict(compress_ratio=0), + ffn=dict( + n_routed_experts=DRAFT_EXPERTS, n_activated_experts=DRAFT_TOPK + ), + ) + for i in range(DRAFT_LAYERS) + ) + self.hc_mult = 4 + + def prologue(self, input_ids, positions, inputs_embeds): + return inputs_embeds, input_ids, positions.shape[0] + + def epilogue(self, hidden, *rest): + return hidden + + +def make_draft_config(**overrides) -> Any: + config = make_config(**overrides) + config.speculative_config = SimpleNamespace( + num_speculative_tokens=DRAFT_TOKENS, method="dspark" + ) + hf = config.model_config.hf_config + hf.sliding_window = sd.SWA_WIDTH + hf.dspark_n_routed_experts = DRAFT_EXPERTS + hf.dspark_num_experts_per_tok = DRAFT_TOPK + return config + + +def _build_draft(monkeypatch, config=None, model=None): + _gpu(monkeypatch) + config = config or make_draft_config() + model = model if model is not None else _DraftModel(config) + sfc = config.compilation_config.static_forward_context + for layer in model.layers: + layer.attn.swa_cache_layer.kv_cache = make_cache() + sfc[layer.attn.swa_cache_layer.prefix] = layer.attn.swa_cache_layer + spec = select_decode_spec(config, sd.DRAFT_DECODE_SPEC) + return SpecializedDecode(model, config, spec), model + + +def _draft_metadata(model, **overrides) -> dict: + width = sd.draft_swa_width(model.vllm_config) + fields = dict( + decode_swa_width=width, + decode_swa_indices=torch.zeros(8, 1, width, dtype=torch.int32), + ) + return { + layer.attn.swa_cache_layer.prefix: _swa(**(fields | overrides)) + for layer in model.layers + } + + +def test_draft_rows_span_the_window_and_the_drafts(): + """The draft's SWA rows are a block's non-causal window, the causal window + plus the drafts, padded; not a decode token's.""" + config = make_draft_config() + width = sd.draft_swa_width(config) + assert width >= sd.SWA_WIDTH + DRAFT_TOKENS + assert width != sd.SWA_WIDTH + + +def test_draft_layers_take_the_drafts_own_routing_and_rows(monkeypatch): + """Every draft layer takes a launch, from a runner built for the draft's + 128 experts, top-3 and non-causal rows, not the backbone's.""" + decode, model = _build_draft(monkeypatch) + kernel = decode.kernel + assert set(kernel.plans) == {layer.attn.layer_id for layer in model.layers} + assert kernel.runner.routing == dict( + experts=DRAFT_EXPERTS, + topk=DRAFT_TOPK, + ug_groups=sd.DRAFT_UG_GROUPS, + swa_width=sd.draft_swa_width(model.vllm_config), + ) + # The draft reads no compressed keys, so it binds only its SWA caches. + assert kernel.cache_layers == tuple( + sorted(layer.attn.swa_cache_layer.prefix for layer in model.layers) + ) + decode.bind_caches() + + +def test_draft_layer_with_compressed_keys_keeps_vllm_path(monkeypatch): + """The draft has no FFN-only launch: a layer the kernels cannot take whole + runs as vLLM's.""" + config = make_draft_config() + model = _DraftModel(config) + for layer in model.layers: + layer.attn.compressor = object() + decode, _ = _build_draft(monkeypatch, config=config, model=model) + assert not decode.kernel.plans + + +def test_draft_first_layer_takes_the_broadcast_embedding(monkeypatch): + """The draft's first layer is handed the embedding already broadcast to the + hc copies, so the launch takes it as the residual and the first copy as the + stream.""" + decode, model = _build_draft(monkeypatch) + positions = torch.zeros(1, dtype=torch.long) + embeds = torch.zeros(1, 4, HIDDEN) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _draft_metadata(model)): + out = decode.kernel.launch_whole( + _plan(decode, LAYERS), embeds, positions, None, None, None, None + ) + assert out[0] == "whole" and not out[2]["fold"] + assert out[1][1].shape == (1, HIDDEN) + assert out[1][2].shape == (1, 4, HIDDEN) + + +def test_draft_and_target_take_only_their_own_rows(monkeypatch): + """Each refuses the other's SWA width, so a draft step cannot run on the + backbone's kernels or the other way round.""" + draft, draft_model = _build_draft(monkeypatch) + target, target_model = _build(monkeypatch) + + def skip(decode, metadata): + with forward(decode.vllm_config, CUDAGraphMode.FULL, metadata): + return decode.spec.skip_step(decode.vllm_config, decode.kernel, 1) + + assert skip(draft, _draft_metadata(draft_model)) is None + assert skip(target, _metadata(target_model)) is None + assert "SWA width" in skip( + draft, _draft_metadata(draft_model, decode_swa_width=sd.SWA_WIDTH) + ) + assert "SWA width" in skip( + target, + _metadata( + target_model, decode_swa_width=sd.draft_swa_width(draft_model.vllm_config) + ), + ) + + +def test_draft_run_takes_every_layer_the_way_its_launch_allows(monkeypatch): + decode, model = _build_draft(monkeypatch) + + def launch_whole(plan, *args, **kwargs): + layer_id = plan.layer.attn.layer_id + model.layers[layer_id - LAYERS].taken.append(("launch", layer_id)) + return None if layer_id == LAYERS else _stream() + + decode.kernel.launch_whole = launch_whole + + step = Step(width=1, kernel_width=1, positions=torch.zeros(1, dtype=torch.long)) + out = decode.kernel.run(step, None, torch.zeros(1, 4, HIDDEN)) + + assert out.shape == (1, HIDDEN) + assert [layer.taken for layer in model.layers] == [ + [("launch", LAYERS), ("layer", LAYERS, False)], + [("launch", LAYERS + 1)], + ] + + def test_launch_declines_a_layout_it_cannot_read(monkeypatch): """A later pipeline stage's first layer gets a stream without the mixes: ``skip_step`` cannot see that, so the launch declines it.""" diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 14c6c97257b1..4fc9f13dbf59 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -11,6 +11,7 @@ """ from collections.abc import Iterable +from typing import TYPE_CHECKING import regex as re import torch @@ -42,6 +43,8 @@ sp_padding_mask, sp_shard, ) +from vllm.models.common.specialized_decode import SpecializedDecode, select_decode_spec +from vllm.models.deepseek_v41.amd.specialized_decode.spec import DRAFT_DECODE_SPEC from .model import ( DeepseekV4DecoderLayer, @@ -50,8 +53,24 @@ _use_sequence_parallel, ) +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.specialized_decode.runtime.decode import ( + DSV41DraftKernel, + ) + logger = init_logger(__name__) +# The draft heads built in this process, for the model runner's cache checks: +# the draft is the speculator's, out of the target's ModelState's reach. +_DRAFT_DECODES: list[SpecializedDecode] = [] + + +def bind_draft_decode_caches() -> None: + """Check every DSpark draft's caches, as the runner binds them.""" + for decode in _DRAFT_DECODES: + decode.bind_caches() + + # MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders): # fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``. _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") @@ -177,12 +196,20 @@ def precompute_and_store_context_kv( continue _insert_context_kv(attn, kv, context_positions, slot_mapping) - def forward( + def prologue( self, input_ids: torch.Tensor, positions: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: + inputs_embeds: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + """Everything the draft's decoder layers need, before the first of them + runs. + + Returns: + The stream the first layer takes, the (possibly sharded) input ids, + and the batch's token count before any sequence-parallel shard. + + """ if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) full_num_tokens = positions.shape[0] @@ -196,7 +223,38 @@ def forward( input_ids = sp_shard(input_ids) # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) + return hidden_states, input_ids, full_num_tokens + def epilogue( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + pre_mix: torch.Tensor | None, + full_num_tokens: int, + ) -> torch.Tensor: + """Everything after the draft's last decoder layer: the final post and + collapse.""" + hidden_states = mhc_post_torch(hidden_states, residual, post_mix, res_mix) + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] + pre_mix = sp_all_gather(pre_mix)[:full_num_tokens] + # Collapse the hc copies with the pre-mix from the last layer's FFN + # mixes, the mix the reference forward_head applies (hc_pre with the + # last block's ffn pre-mix). Return the PRE-norm head hidden; + # compute_logits applies self.norm. + return DeepseekV4DecoderLayer._hc_collapse(hidden_states, pre_mix) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden_states, input_ids, full_num_tokens = self.prologue( + input_ids, positions, inputs_embeds + ) residual = post_mix = res_mix = pre_mix = None for layer in self.layers: hidden_states, residual, post_mix, res_mix, pre_mix = layer( @@ -208,16 +266,9 @@ def forward( res_mix, residual, ) - hidden_states = mhc_post_torch(hidden_states, residual, post_mix, res_mix) - if self.use_sequence_parallel: - hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] - pre_mix = sp_all_gather(pre_mix)[:full_num_tokens] - # Collapse the hc copies with the pre-mix from the last layer's FFN - # mixes — the mix the reference forward_head applies (hc_pre with the - # last block's ffn pre-mix). Return the PRE-norm head hidden; - # compute_logits applies self.norm. - hidden_states = DeepseekV4DecoderLayer._hc_collapse(hidden_states, pre_mix) - return hidden_states + return self.epilogue( + hidden_states, residual, post_mix, res_mix, pre_mix, full_num_tokens + ) def _insert_context_kv( @@ -300,6 +351,11 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() assert vllm_config.speculative_config is not None + self.vllm_config = vllm_config + self.specialized_decode_spec = select_decode_spec( + vllm_config, DRAFT_DECODE_SPEC + ) + self.specialized_decode: SpecializedDecode[DSV41DraftKernel] | None = None self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config self.quant_config = vllm_config.quant_config @@ -350,6 +406,10 @@ def forward( inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor: # Returns the pre-norm collapsed head hidden ([T, hidden_size]). + if self.specialized_decode is not None: + step = self.specialized_decode.plan(positions) + if step is not None: + return self.specialized_decode.run(step, input_ids, inputs_embeds) return self.model(input_ids, positions, inputs_embeds) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -495,7 +555,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: # ROCm linears and fused-MoE runners finalize their own quantized # parameters. NVIDIA-only MegaMoE/WO-A requantization is not needed. - return + if self.specialized_decode_spec is not None and self.specialized_decode is None: + self.specialized_decode = SpecializedDecode( + self.model, self.vllm_config, self.specialized_decode_spec + ) + _DRAFT_DECODES.append(self.specialized_decode) def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v41/amd/model_state.py b/vllm/models/deepseek_v41/amd/model_state.py index a3e1ef027058..b7e412c0ef1e 100644 --- a/vllm/models/deepseek_v41/amd/model_state.py +++ b/vllm/models/deepseek_v41/amd/model_state.py @@ -9,12 +9,16 @@ class DeepseekV41ROCMModelState(DeepseekV41ModelState): - """Checks the specialized decode's caches for every cache the runner binds: the - throwaway one memory profiling allocates, then the real one.""" + """Checks the specialized decode's caches, the backbone's and the DSpark + draft's, for every cache the runner binds: the throwaway one memory + profiling allocates, then the real one.""" def initialize_kv_cache( self, kv_cache_config: KVCacheConfig, block_tables: BlockTables ) -> None: + from vllm.models.deepseek_v41.amd.dspark import bind_draft_decode_caches + super().initialize_kv_cache(kv_cache_config, block_tables) if self.model.specialized_decode is not None: self.model.specialized_decode.bind_caches() + bind_draft_decode_caches() diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py index 1f8db64687f6..aa236e9186bc 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1's decode kernel: one launch a layer over the stock backbone.""" +"""DeepSeek-V4.1's decode kernels: one launch a layer over the stock backbone, +and the same over the DSpark draft head's layers.""" from collections.abc import Mapping from dataclasses import dataclass @@ -13,7 +14,16 @@ from vllm.logger import init_logger from vllm.models.common.specialized_decode.spec import Step -from ..spec import MAX_ROWS, RECORD +from ..spec import ( + DRAFT_UG_GROUPS, + EXPERTS, + MAX_ROWS, + RECORD, + SWA_WIDTH, + TOPK, + draft_routing, + draft_swa_width, +) if TYPE_CHECKING: from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer @@ -80,8 +90,26 @@ class LayerPlan: folds: tuple[bool, ...] = () -class DSV41Kernel: - def __init__(self, vllm_config, model, widths: tuple[int, ...]): +class LayerLaunches: + """The launches of one stack of DSv4.1 decoder layers: the backbone's, or + the DSpark draft's, which has its own routing, SWA rows and runner. + + A subclass owns the layer loop in ``run``; everything the launches need is + settled here, at load. + """ + + draft = False + + def __init__( + self, + vllm_config, + model, + widths: tuple[int, ...], + experts: int, + topk: int, + ug_groups: int, + swa_width: int, + ): from vllm.models.common.specialized_decode.guards import tp_group from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( MAX_TOKENS, @@ -92,7 +120,17 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): rank, world, cpu_group = tp_group() device = next(model.parameters()).device self.model = model - self.runner = DSV41MonoLayer(world, rank, cpu_group, device) + self.experts, self.topk, self.swa_width = experts, topk, swa_width + self.runner = DSV41MonoLayer( + world, + rank, + cpu_group, + device, + experts=experts, + topk=topk, + ug_groups=ug_groups, + swa_width=swa_width, + ) self.embed_layer = getattr(model, "start_layer", 0) # A decode step is never trimmed, so ``run`` owns every layer, the ones # past the replay cut included. @@ -105,7 +143,7 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): tp_size = vllm_config.parallel_config.tensor_parallel_size for idx in range(self.embed_layer, self.layer_end): layer = model.layers[idx] - plan = self._plan_layer(layer, tp_size) + plan = self._plan_layer(layer, tp_size, first=idx == self.embed_layer) if plan is None: continue self.plans[layer.attn.layer_id] = plan @@ -119,10 +157,10 @@ def __init__(self, vllm_config, model, widths: tuple[int, ...]): keys.append(key) self.cache_layers = tuple(sorted(set(self.swa_keys) | set(self.comp_keys))) - def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: + def _plan_layer(self, layer, tp_size: int, first: bool) -> LayerPlan | None: """Info needed to run the specialized decode kernel for this layer.""" attn = layer.attn - why = _skip_layer(layer) + why = _skip_layer(layer, self.experts, self.topk) if why is not None: logger.info_once( "DSv4.1 specialized decode: layer %d on vLLM's path (%s)", @@ -131,7 +169,15 @@ def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: ) return None why = _ffn_only(layer) - weights = _weights(layer, why is not None, tp_size) + if why is None and self.draft and (attn.compress_ratio or attn.indexer): + why = "compressed keys" + if why is not None and self.draft: + # The draft has no FFN launch: its layers run whole or not at all. + logger.info_once( + "DSv4.1 specialized decode: the draft layers on vLLM's path (%s)", why + ) + return None + weights = _weights(layer, why is not None, tp_size, self.experts) if why is not None: logger.info_once( "DSv4.1 specialized decode: the FFN launch for layers with %s", why @@ -140,10 +186,10 @@ def _plan_layer(self, layer, tp_size: int) -> LayerPlan | None: # reduces it on the steps that keep vLLM's path. attn.wo_b.reduce_results = False return LayerPlan(layer, weights, ffn_only=True) - comp_key = attn.compressed_cache_prefix + comp_key = attn.compressed_cache_prefix if attn.compress_ratio else None # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write. - settled = attn.layer_id == self.embed_layer or layer.engram is not None + settled = first or layer.engram is not None return LayerPlan( layer, weights, @@ -193,14 +239,19 @@ def launch_whole( layer, runner = plan.layer, self.runner attn = layer.attn first = residual is None - if first and x.dim() != 2: + # The draft's first layer takes the embedding already broadcast to the + # hc copies ([rows, hc, dim]), the backbone's the 2-D embedding. + if first and x.dim() != (3 if self.draft else 2): return None tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) if not all(t is not None and t.is_contiguous() for t in tensors): return None fold = True if first: - residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + if self.draft: + residual, x = x, x[:, 0].contiguous() + else: + residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() pre_mix = runner.first_pre_mix(x.shape[0]) post_mix = res_mix = pre_mix # unread without the fold fold = False @@ -256,6 +307,21 @@ def launch_ffn( return None return self.runner.ffn(plan.weights, *tensors) + +class DSV41Kernel(LayerLaunches): + """The backbone's decode: 384 experts, top-6, causal SWA windows.""" + + def __init__(self, vllm_config, model, widths: tuple[int, ...]): + super().__init__( + vllm_config, + model, + widths, + experts=EXPERTS, + topk=TOPK, + ug_groups=2, + swa_width=SWA_WIDTH, + ) + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: model, positions = self.model, step.positions ( @@ -336,6 +402,48 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: ) +class DSV41DraftKernel(LayerLaunches): + """The DSpark draft's decode: the draft's own routing and non-causal SWA + rows, over the draft head's layers.""" + + draft = True + + def __init__(self, vllm_config, model, widths: tuple[int, ...]): + experts, topk = draft_routing(vllm_config) + super().__init__( + vllm_config, + model, + widths, + experts=experts, + topk=topk, + ug_groups=DRAFT_UG_GROUPS, + swa_width=draft_swa_width(vllm_config), + ) + + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + model, positions = self.model, step.positions + hidden, input_ids, full_num_tokens = model.prologue( + input_ids, positions, inputs_embeds + ) + residual = post_mix = res_mix = pre_mix = None + for idx in range(self.embed_layer, self.layer_end): + layer = model.layers[idx] + plan = self.plans.get(layer.attn.layer_id) + out = None + if plan is not None: + out = self.launch_whole( + plan, hidden, positions, residual, post_mix, res_mix, pre_mix + ) + if out is None: + out = layer( + hidden, positions, input_ids, pre_mix, post_mix, res_mix, residual + ) + hidden, residual, post_mix, res_mix, pre_mix = out + return model.epilogue( + hidden, residual, post_mix, res_mix, pre_mix, full_num_tokens + ) + + def _ffn_only(layer) -> str | None: """Why this layer's attention seam and attention stay on vLLM's path, which leaves it the FFN launch, or None when it runs whole.""" @@ -355,7 +463,7 @@ def _ffn_only(layer) -> str | None: return None -def _skip_layer(layer) -> str | None: +def _skip_layer(layer, experts: int, topk: int) -> str | None: """Why this layer keeps vLLM's path for every step, or None when it takes a launch (whole or FFN only: ``_plan_layer`` picks which).""" ffn = layer.ffn @@ -366,18 +474,20 @@ def _skip_layer(layer) -> str | None: if ( ffn.shared_experts is None or ffn.gate.tid2eid is not None - or ffn.n_routed_experts != 384 - or ffn.n_activated_experts != 6 + or ffn.n_routed_experts != experts + or ffn.n_activated_experts != topk or ffn.routed_scaling_factor != 1.5 or ffn.swiglu_limit != 10.0 or ffn.scoring_func != "sqrtsoftplus" or not ffn.renormalize ): - return "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + return ( + f"MoE routing outside the kernels' ({experts} / {topk}, sqrtsoftplus, x1.5)" + ) return None -def _weights(layer, ffn_only: bool, world: int): +def _weights(layer, ffn_only: bool, world: int, experts: int): from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( MonoLayerWeights, ) @@ -405,7 +515,7 @@ def _weights(layer, ffn_only: bool, world: int): sw2=sh.down_proj.weight, sw2_s=_scale_bytes(sh.down_proj.weight_scale), ) - weights.check(world) + weights.check(world, experts) return weights diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py index 31d123f00521..a85cfdeb15f5 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py @@ -14,15 +14,43 @@ if TYPE_CHECKING: from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterSparseSWAMetadata - from .runtime.decode import DSV41Kernel + from .runtime.decode import DSV41DraftKernel, DSV41Kernel, LayerLaunches RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes SWA_WIDTH = 128 # a causal decode token's window slots MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +EXPERTS = 384 # the backbone's routed experts, top-TOPK of them a token +TOPK = 6 +# FP4 groups of an ug task: the draft's 128-expert unions are small +DRAFT_UG_GROUPS = 1 # the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats CDNA_VERSIONS = (4,) +def draft_routing(vllm_config) -> tuple[int, int]: + """The DSpark draft's routed experts and top-k (128 / 3 on V4.1-Flash).""" + config = vllm_config.model_config.hf_config + return ( + getattr(config, "dspark_n_routed_experts", 0) or EXPERTS, + getattr(config, "dspark_num_experts_per_tok", 0) or TOPK, + ) + + +def draft_swa_width(vllm_config) -> int: + """The DSpark draft's SWA rows: a block's non-causal window, not a causal + decode token's.""" + from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, + ) + + config = vllm_config.model_config.hf_config + speculative = vllm_config.speculative_config + assert speculative is not None + return get_dspark_swa_index_width( + config.sliding_window, speculative.num_speculative_tokens + ) + + def _widths(vllm_config) -> tuple[int, ...]: """The configured capture sizes the kernels can serve, so a deployment that asked for fewer graphs keeps them; every size at all when it asked for @@ -74,7 +102,7 @@ def _refuse_model(model) -> list[str]: return reasons -def _refuse_caches(vllm_config, decode: "DSV41Kernel", caches) -> list[str]: +def _refuse_caches(vllm_config, decode: "LayerLaunches", caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as [blocks, block, RECORD] bytes.""" reasons = [] @@ -91,20 +119,9 @@ def _refuse_caches(vllm_config, decode: "DSV41Kernel", caches) -> list[str]: return reasons -def _skip_step(vllm_config, decode: "DSV41Kernel", width: int) -> str | None: - """Why this step keeps vLLM's path at every layer, or None when the kernels - take it. One verdict for every kernel layer, from the batch the whole TP - group sees alike.""" - md = get_forward_context().attn_metadata - if not decode.plans: - return "no layer takes a launch" - if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches - return "no attention metadata" - replay = decode.model.decoder_replay_layers - if replay is not None and replay.uses_replay_batch(): - return "replay batch step" - if any(md.get(k) is None for k in decode.comp_keys): - return "layer metadata missing" +def _skip_swa(decode, md: dict, width: int) -> str | None: + """Why the SWA rows of this step are outside what the kernels read, or None + when they take it.""" for key in decode.swa_keys: swa = cast("DeepseekV41ROCMAiterSparseSWAMetadata | None", md.get(key)) if swa is None: @@ -113,7 +130,7 @@ def _skip_step(vllm_config, decode: "DSV41Kernel", width: int) -> str | None: return "not a decode-only step" if swa.num_decode_tokens > width: return "more decode tokens than rows" - if swa.decode_swa_width != SWA_WIDTH: + if swa.decode_swa_width != decode.swa_width: return f"SWA width {swa.decode_swa_width}" if ( swa.decode_swa_indices is None @@ -124,15 +141,51 @@ def _skip_step(vllm_config, decode: "DSV41Kernel", width: int) -> str | None: return None +def _skip_step(vllm_config, decode: "DSV41Kernel", width: int) -> str | None: + """Why this step keeps vLLM's path at every layer, or None when the kernels + take it. One verdict for every kernel layer, from the batch the whole TP + group sees alike.""" + md = get_forward_context().attn_metadata + if not decode.plans: + return "no layer takes a launch" + if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + return "no attention metadata" + replay = decode.model.decoder_replay_layers + if replay is not None and replay.uses_replay_batch(): + return "replay batch step" + if any(md.get(k) is None for k in decode.comp_keys): + return "layer metadata missing" + return _skip_swa(decode, md, width) + + +def _skip_draft_step(vllm_config, decode: "DSV41DraftKernel", width: int) -> str | None: + """Why this draft step keeps vLLM's path at every layer. The draft reads no + compressed keys and never runs a replay batch, so its SWA rows decide it.""" + md = get_forward_context().attn_metadata + if not decode.plans: + return "no draft layer takes a launch" + if not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + return "no attention metadata" + return _skip_swa(decode, md, width) + + def _build(vllm_config, model, widths: tuple[int, ...]) -> "DSV41Kernel": from .runtime.decode import DSV41Kernel return DSV41Kernel(vllm_config, model, widths) +def _build_draft(vllm_config, model, widths: tuple[int, ...]) -> "DSV41DraftKernel": + from .runtime.decode import DSV41DraftKernel + + return DSV41DraftKernel(vllm_config, model, widths) + + +_SUPPORTS = frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}) + DECODE_SPEC = DecodeSpec["DSV41Kernel"]( name="dsv41-cdna4", - supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), + supports=_SUPPORTS, widths=_widths, refuse_config=_refuse_config, refuse_model=_refuse_model, @@ -140,3 +193,14 @@ def _build(vllm_config, model, widths: tuple[int, ...]) -> "DSV41Kernel": skip_step=_skip_step, build=_build, ) + +DRAFT_DECODE_SPEC = DecodeSpec["DSV41DraftKernel"]( + name="dsv41-dspark-cdna4", + supports=_SUPPORTS, + widths=_widths, + refuse_config=_refuse_config, + refuse_model=_refuse_model, + refuse_caches=_refuse_caches, + skip_step=_skip_draft_step, + build=_build_draft, +) From d6678803a148eaf85ff03ff3296447b5c5cbe8dd Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 08:48:24 -0500 Subject: [PATCH 61/79] revert: rename files Signed-off-by: simondanielsson --- .../models/amd/test_dsv41_mono_gfx942.py | 20 ++++----- tests/models/test_dsv41_mono_numerics.py | 16 ++++---- tests/models/test_dsv41_specialized_decode.py | 2 +- .../{specialized_decode => }/mono/README.md | 0 .../{specialized_decode => }/mono/__init__.py | 0 .../mono/attention/__init__.py | 0 .../mono/attention/back.py | 0 .../mono/attention/device.py | 0 .../mono/attention/front.py | 0 .../mono/attention/gemv.py | 0 .../mono/attention/plan.py | 0 .../mono/common/__init__.py | 0 .../mono/common/arch.py | 0 .../mono/common/gfx942.py | 0 .../mono/common/mx.py | 4 +- .../mono/common/ops.py | 11 ++--- .../mono/common/peer_memory.py | 0 .../mono/common/plan.py | 0 .../mono/common/sync.py | 2 +- .../{specialized_decode => }/mono/keylist.py | 0 .../{specialized_decode => }/mono/layer.py | 23 ++--------- .../{specialized_decode => }/mono/runner.py | 0 .../{specialized_decode => }/mono/sources.py | 4 +- .../mono/stages/__init__.py | 0 .../mono/stages/dims.py | 5 +-- .../mono/stages/gemv.py | 12 +++--- .../mono/stages/moe.py | 41 +++++++------------ .../mono/stages/moe_plan.py | 4 +- .../mono/stages/moe_shape.py | 24 +++-------- .../mono/stages/router.py | 7 +--- .../mono/stages/seam.py | 12 ++---- .../mono/weights942.py | 0 .../amd/specialized_decode/runtime/decode.py | 8 ++-- .../specialized_decode/runtime/weights942.py | 10 ++--- .../amd/specialized_decode/spec.py | 6 +-- 35 files changed, 76 insertions(+), 135 deletions(-) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/README.md (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/attention/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/attention/back.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/attention/device.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/attention/front.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/attention/gemv.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/attention/plan.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/arch.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/gfx942.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/mx.py (94%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/ops.py (97%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/peer_memory.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/plan.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/common/sync.py (98%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/keylist.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/layer.py (98%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/runner.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/sources.py (73%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/__init__.py (100%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/dims.py (94%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/gemv.py (96%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/moe.py (98%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/moe_plan.py (93%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/moe_shape.py (93%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/router.py (94%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/stages/seam.py (97%) rename vllm/models/deepseek_v41/amd/{specialized_decode => }/mono/weights942.py (100%) diff --git a/tests/kernels/models/amd/test_dsv41_mono_gfx942.py b/tests/kernels/models/amd/test_dsv41_mono_gfx942.py index 5cb86c8fc1f4..4fb29c29b1cb 100644 --- a/tests/kernels/models/amd/test_dsv41_mono_gfx942.py +++ b/tests/kernels/models/amd/test_dsv41_mono_gfx942.py @@ -47,27 +47,27 @@ def _on_gfx942() -> bool: from flydsl.expr import const_expr, gpu, range_constexpr # noqa: E402 from flydsl.expr.typing import Int64, T # noqa: E402 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention import ( # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.attention import ( # noqa: E402 gemv as G, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.device import ( # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.attention.device import ( # noqa: E402 THREADS, WAVES, fp8x8_bf16, rsrc, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.gfx942 import ( # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.common.gfx942 import ( # noqa: E402 fp4x8_fnuz, mfma_bf16_k32, ocp_to_fnuz, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.stages import ( # noqa: E402 gemv as SG, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.stages import ( # noqa: E402 moe as M, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.weights942 import ( # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.weights942 import ( # noqa: E402 EVEN_ODD, codes_tile_rounds, dense_copy, @@ -100,7 +100,7 @@ def _rel_err(got: torch.Tensor, ref: torch.Tensor) -> float: def _keylist_reference(slot, pos, swa_idx, swa_lens, t2r, ratio, topk, bt, cb): - from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import ( + from vllm.models.deepseek_v41.amd.mono.attention.plan import ( KEYS, TOPK, ) @@ -126,8 +126,8 @@ def _keylist_reference(slot, pos, swa_idx, swa_lens, t2r, ratio, topk, bt, cb): @pytest.mark.parametrize("ratio", [0, 1, 2, 4]) def test_keylist_matches_torch_reference(ratio): - from vllm.models.deepseek_v41.amd.specialized_decode.mono import keylist - from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import ( + from vllm.models.deepseek_v41.amd.mono import keylist + from vllm.models.deepseek_v41.amd.mono.attention.plan import ( KEYS, TOPK, WINDOW, @@ -619,7 +619,7 @@ def test_contrib_942_is_bf16_of_product(routing): every pick of a 42-row step, with w the pick's route weight from the route table (RouteShape.pick_w). down_combine rounds this same f32 product on gfx950, so the gfx942 combine adds the same values.""" - from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.moe_shape import ( + from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( RouteShape, ) diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py index e8d8425ddbfd..c36064386acf 100644 --- a/tests/models/test_dsv41_mono_numerics.py +++ b/tests/models/test_dsv41_mono_numerics.py @@ -310,17 +310,17 @@ def _check_layers(rank: int, tp: int, device) -> None: from vllm.model_executor.model_loader.utils import process_weights_after_loading from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE - from vllm.models.deepseek_v41.amd.rocm import ( - compute_global_topk_ragged_indices_and_indptr, - ) - from vllm.models.deepseek_v41.amd.specialized_decode.mono.layer import ( + from vllm.models.deepseek_v41.amd.mono.layer import ( scratch_layout, ) - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + from vllm.models.deepseek_v41.amd.mono.runner import ( AttnWeights, DSV41MonoLayer, MonoLayerWeights, ) + from vllm.models.deepseek_v41.amd.rocm import ( + compute_global_topk_ragged_indices_and_indptr, + ) from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( @@ -606,13 +606,13 @@ def _replay_942(rank: int, tp: int, device) -> None: from vllm.model_executor.model_loader.utils import process_weights_after_loading from vllm.models.deepseek_v41.amd import mono_decode from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE - from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import Dims - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.mono.runner import ( AttnWeights, DSV41MonoLayer, MonoLayerWeights, ) - from vllm.models.deepseek_v41.amd.specialized_decode.mono.weights942 import ( + from vllm.models.deepseek_v41.amd.mono.weights942 import ( linear_copy, ) from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 5832d5ed8280..9bdad89a0ddd 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -50,7 +50,7 @@ def ffn(self, *args): return ("ffn", args) -_MONO = "vllm.models.deepseek_v41.amd.specialized_decode.mono" +_MONO = "vllm.models.deepseek_v41.amd.mono" def _stub_modules(monkeypatch) -> None: diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/README.md rename to vllm/models/deepseek_v41/amd/mono/README.md diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/__init__.py rename to vllm/models/deepseek_v41/amd/mono/__init__.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/__init__.py b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/__init__.py rename to vllm/models/deepseek_v41/amd/mono/attention/__init__.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/back.py rename to vllm/models/deepseek_v41/amd/mono/attention/back.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/device.py b/vllm/models/deepseek_v41/amd/mono/attention/device.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/device.py rename to vllm/models/deepseek_v41/amd/mono/attention/device.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/front.py rename to vllm/models/deepseek_v41/amd/mono/attention/front.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/gemv.py b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/gemv.py rename to vllm/models/deepseek_v41/amd/mono/attention/gemv.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/plan.py b/vllm/models/deepseek_v41/amd/mono/attention/plan.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/attention/plan.py rename to vllm/models/deepseek_v41/amd/mono/attention/plan.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/mono/common/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/__init__.py rename to vllm/models/deepseek_v41/amd/mono/common/__init__.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/arch.py b/vllm/models/deepseek_v41/amd/mono/common/arch.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/arch.py rename to vllm/models/deepseek_v41/amd/mono/common/arch.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/gfx942.py b/vllm/models/deepseek_v41/amd/mono/common/gfx942.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/gfx942.py rename to vllm/models/deepseek_v41/amd/mono/common/gfx942.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/mx.py b/vllm/models/deepseek_v41/amd/mono/common/mx.py similarity index 94% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/mx.py rename to vllm/models/deepseek_v41/amd/mono/common/mx.py index 5a10887ca6c2..e0e914825c00 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/mx.py +++ b/vllm/models/deepseek_v41/amd/mono/common/mx.py @@ -14,8 +14,8 @@ from flydsl.expr import rocdl from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.arch import GFX942 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import CM_NT +from vllm.models.deepseek_v41.amd.mono.common.arch import GFX942 +from vllm.models.deepseek_v41.amd.mono.common.ops import CM_NT # gfx942 quantizes to FNUZ with 224: the same bytes as OCP with 448 and a scale # code one higher (common/gfx942.py). diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/ops.py b/vllm/models/deepseek_v41/amd/mono/common/ops.py similarity index 97% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/ops.py rename to vllm/models/deepseek_v41/amd/mono/common/ops.py index e54efed638ba..5b35d2a570d6 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/ops.py +++ b/vllm/models/deepseek_v41/amd/mono/common/ops.py @@ -35,14 +35,9 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T, as_ir_value -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.arch import GFX942 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.gfx942 import ( - mfma_bf16_k32, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( - THREADS, - WAVES, -) +from vllm.models.deepseek_v41.amd.mono.common.arch import GFX942 +from vllm.models.deepseek_v41.amd.mono.common.gfx942 import mfma_bf16_k32 +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES # ---------------------------------------------------------------- primitives # gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/peer_memory.py b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/peer_memory.py rename to vllm/models/deepseek_v41/amd/mono/common/peer_memory.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/plan.py b/vllm/models/deepseek_v41/amd/mono/common/plan.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/plan.py rename to vllm/models/deepseek_v41/amd/mono/common/plan.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/sync.py b/vllm/models/deepseek_v41/amd/mono/common/sync.py similarity index 98% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/common/sync.py rename to vllm/models/deepseek_v41/amd/mono/common/sync.py index be2276b580dc..0220ce12c32a 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/common/sync.py +++ b/vllm/models/deepseek_v41/amd/mono/common/sync.py @@ -14,7 +14,7 @@ from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( +from vllm.models.deepseek_v41.amd.mono.common.ops import ( CM_DEV, CM_SYS, bf16_pair, diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/keylist.py b/vllm/models/deepseek_v41/amd/mono/keylist.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/keylist.py rename to vllm/models/deepseek_v41/amd/mono/keylist.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py similarity index 98% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py rename to vllm/models/deepseek_v41/amd/mono/layer.py index 2d4272abaf26..f4f976b9a33e 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/layer.py +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -23,25 +23,6 @@ from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import Int32, Int64, T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.mx import clamp_fp8 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( - butterfly, - fp8_pack4, - peer_bases, - traced, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( - WAVES, - first_task, - key_tuple, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import ( - Mailbox, - preg, - publish, - sreg, -) - from .attention import back as mb_back from .attention import front as mb_front from .attention.device import ( @@ -54,6 +35,10 @@ rsrc, ) from .attention.plan import HIDDEN, WINDOW, Dims, back_scratch, front_scratch +from .common.mx import clamp_fp8 +from .common.ops import butterfly, fp8_pack4, peer_bases, traced +from .common.plan import WAVES, first_task, key_tuple +from .common.sync import Mailbox, preg, publish, sreg from .sources import SOURCES from .stages import moe, seam from .stages.dims import Dims as StageDims diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/runner.py rename to vllm/models/deepseek_v41/amd/mono/runner.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py b/vllm/models/deepseek_v41/amd/mono/sources.py similarity index 73% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py rename to vllm/models/deepseek_v41/amd/mono/sources.py index 1714374a9068..49f34d8fed67 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/sources.py +++ b/vllm/models/deepseek_v41/amd/mono/sources.py @@ -3,8 +3,6 @@ """The mono layer's kernel sources, in every build's JIT cache key (``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( - source_digest, -) +from vllm.models.deepseek_v41.amd.mono.common.plan import source_digest SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/__init__.py b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/__init__.py rename to vllm/models/deepseek_v41/amd/mono/stages/__init__.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py b/vllm/models/deepseek_v41/amd/mono/stages/dims.py similarity index 94% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py rename to vllm/models/deepseek_v41/amd/mono/stages/dims.py index f65f51041f4a..6fd119c77206 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/dims.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/dims.py @@ -8,10 +8,7 @@ from dataclasses import dataclass -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( - Shard, - padded, -) +from vllm.models.deepseek_v41.amd.mono.common.plan import Shard, padded # the model's full widths MOE_INTER = 2304 # a routed expert's intermediate width diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py similarity index 96% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py rename to vllm/models/deepseek_v41/amd/mono/stages/gemv.py index 309eff748a3a..0627071ea073 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/gemv.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py @@ -12,24 +12,24 @@ from flydsl.expr import const_expr, range_constexpr from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.arch import GFX942 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.gfx942 import ( +from vllm.models.deepseek_v41.amd.mono.common.arch import GFX942 +from vllm.models.deepseek_v41.amd.mono.common.gfx942 import ( mfma_fp8_scaled, pair_i64, pin_acc, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.mx import ( +from vllm.models.deepseek_v41.amd.mono.common.mx import ( UNIT_SCALE, mfma_scaled, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( +from vllm.models.deepseek_v41.amd.mono.common.ops import ( CM_NT, fresh, rsrc, traced, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import WAVES -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import POLL_MAX +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES +from vllm.models.deepseek_v41.amd.mono.common.sync import POLL_MAX ROWS = 16 # GEMV rows a task # An MFMA's N: the tokens or picks that one 16x16 MFMA takes. diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/mono/stages/moe.py similarity index 98% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py rename to vllm/models/deepseek_v41/amd/mono/stages/moe.py index 6a717ccf19cb..e7f773b39693 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe.py @@ -49,14 +49,14 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.arch import GFX942 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.gfx942 import ( +from vllm.models.deepseek_v41.amd.mono.common.arch import GFX942 +from vllm.models.deepseek_v41.amd.mono.common.gfx942 import ( fp4x8_fnuz, mfma_fp8, pair_i64, pin_acc, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.mx import ( +from vllm.models.deepseek_v41.amd.mono.common.mx import ( FP4, FP8, FP8_MAX, @@ -68,7 +68,7 @@ rcp_pow2, scale_index, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( +from vllm.models.deepseek_v41.amd.mono.common.ops import ( CM_DEV, CM_NT, _hw_f32, @@ -91,29 +91,21 @@ sum_partials, traced, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( - stamp as stamp_lds, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( +from vllm.models.deepseek_v41.amd.mono.common.ops import stamp as stamp_lds +from vllm.models.deepseek_v41.amd.mono.common.plan import ( BLOCKS, THREADS, WAVES, first_task, pair_layout, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import ( - preg, - publish, - sreg, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( - moe_plan as mplan, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.dims import UG_PART +from vllm.models.deepseek_v41.amd.mono.common.sync import preg, publish, sreg +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART # helpers the kernel calls are imported by name: FlyDSL keys a build by the # source of the same-directory functions it names, never through a module -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( MFMA_N, ROWS, TILE, @@ -127,23 +119,18 @@ tile_rows, token_tiles, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.moe_shape import ( +from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( MASK_WORDS, MoeBuild, route_shape, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.router import ( - PARTS as ROUTER_PARTS, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.router import ( +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS +from vllm.models.deepseek_v41.amd.mono.stages.router import ( router_tile_mfma, router_weight_loads, router_x_loads, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.seam import ( - HIDDEN, - attn_peer_bytes, -) +from vllm.models.deepseek_v41.amd.mono.stages.seam import HIDDEN, attn_peer_bytes # a pick tile is a token tile's MFMA N: ``gather_x8r`` fills one from the other assert mplan.TILE == TILE diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py similarity index 93% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py rename to vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py index a5e633a9ab4f..be25080796f0 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_plan.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py @@ -15,8 +15,8 @@ by the kernel on traced Int32 and by the CPU property test on Python ints. """ -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.arch import GFX942 -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import cdiv, imin +from vllm.models.deepseek_v41.amd.mono.common.arch import GFX942 +from vllm.models.deepseek_v41.amd.mono.common.plan import cdiv, imin # The picks a pick tile takes: a token tile's (``stages.gemv.TILE``), whose x # rows the ug task of the tile gathers into LDS. 6 on gfx942, 16 elsewhere. diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py similarity index 93% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py rename to vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py index b7dea897c0cf..23fd36ebbf68 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/moe_shape.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py @@ -12,25 +12,11 @@ import flydsl.expr as fx -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( - BLOCKS, - WAVES, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages import ( - moe_plan as mplan, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.dims import ( - UG_PART, - Dims, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( - ROWS, - TILE, - token_tiles, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.router import ( - PARTS as ROUTER_PARTS, -) +from vllm.models.deepseek_v41.amd.mono.common.plan import BLOCKS, WAVES +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART, Dims +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ROWS, TILE, token_tiles +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS EXPERTS, TOPK = 384, 6 # the target's routing (a build's: ``MoeBuild``) ROUTER0 = 8 # the router's first task: past the xq tasks (a token each) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py b/vllm/models/deepseek_v41/amd/mono/stages/router.py similarity index 94% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py rename to vllm/models/deepseek_v41/amd/mono/stages/router.py index f2fa69dc9bb2..a7d6f411a816 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/router.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/router.py @@ -20,11 +20,8 @@ from flydsl.expr import const_expr, range_constexpr from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( - mfma_bf16, - rsrc, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import WAVES +from vllm.models.deepseek_v41.amd.mono.common.ops import mfma_bf16, rsrc +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES HIDDEN = 5120 ROWS = 16 # a tile's rows: the MFMA's M diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/mono/stages/seam.py similarity index 97% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py rename to vllm/models/deepseek_v41/amd/mono/stages/seam.py index aae238376d37..ea73e2ef5435 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/mono/stages/seam.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/seam.py @@ -31,7 +31,7 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.ops import ( +from vllm.models.deepseek_v41.amd.mono.common.ops import ( batched_rounds, bf16_round, bf_hi, @@ -45,13 +45,9 @@ sum_partials, traced, ) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.plan import ( - THREADS, - WAVES, - pair_layout, -) -from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.sync import preg -from vllm.models.deepseek_v41.amd.specialized_decode.mono.stages.gemv import ( +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES, pair_layout +from vllm.models.deepseek_v41.amd.mono.common.sync import preg +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( ld_bf, ld_f32, plus, diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/mono/weights942.py b/vllm/models/deepseek_v41/amd/mono/weights942.py similarity index 100% rename from vllm/models/deepseek_v41/amd/specialized_decode/mono/weights942.py rename to vllm/models/deepseek_v41/amd/mono/weights942.py diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py index b28f0204c28c..11d48601e6a7 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py @@ -110,7 +110,7 @@ def __init__( swa_width: int, ): from vllm.models.common.specialized_decode.guards import tp_group - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + from vllm.models.deepseek_v41.amd.mono.runner import ( MAX_TOKENS, DSV41MonoLayer, ) @@ -563,7 +563,7 @@ def _skip_layer(layer, experts: int, topk: int, tp_size: int) -> str | None: def _weights(layer, ffn_only: bool, world: int, experts: int): - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import ( + from vllm.models.deepseek_v41.amd.mono.runner import ( MonoLayerWeights, ) @@ -599,8 +599,8 @@ def _weights(layer, ffn_only: bool, world: int, experts: int): def _attn_weights(a, world: int): - from vllm.models.deepseek_v41.amd.specialized_decode.mono.attention.plan import Dims - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import AttnWeights + from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights attn = AttnWeights( layer_id=a.layer_id, diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py index ddbf1c2f2f1c..05b81c2b5ae4 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py @@ -48,7 +48,7 @@ def _capture_plain_experts(experts) -> None: def process_weights_after_loading(layer): if layer is experts: - from ..mono.weights942 import moe_copies + from ...mono.weights942 import moe_copies _keep( layer, @@ -75,7 +75,7 @@ def _copy_linear(linear) -> None: def process_weights_after_loading(layer): out = original(layer) if layer is linear: - from ..mono.weights942 import linear_copy + from ...mono.weights942 import linear_copy _keep(layer, linear_copy(layer.weight, layer.weight_scale)) return out @@ -100,7 +100,7 @@ def copy_at_load(layer: "DeepseekV4DecoderLayer", ffn_only: bool) -> None: def layer_weights(layer: "DeepseekV4DecoderLayer", ffn_only: bool, world: int): """The layer's tensors as the gfx942 kernels read them: the dense MXFP8 linears and the routed experts as the copies made at load.""" - from ..mono.runner import MonoLayerWeights + from ...mono.runner import MonoLayerWeights f = layer.ffn e, sh = f.experts.routed_experts, f.shared_experts @@ -134,8 +134,8 @@ def layer_weights(layer: "DeepseekV4DecoderLayer", ffn_only: bool, world: int): def _attn_weights(a, world: int): - from ..mono.attention.plan import Dims - from ..mono.runner import AttnWeights + from ...mono.attention.plan import Dims + from ...mono.runner import AttnWeights wqkv, wq_b, wo_a, wo_b = (getattr(a, name).mono942 for name in ATTN_LINEARS) attn = AttnWeights( diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py index 4bcceb495bb9..7e3c167ef4c9 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py @@ -101,7 +101,7 @@ def _refuse_config(vllm_config) -> list[str]: "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" ) try: - import vllm.models.deepseek_v41.amd.specialized_decode.mono.runner # noqa: F401 + import vllm.models.deepseek_v41.amd.mono.runner # noqa: F401 except ImportError as err: reasons.append(f"needs FlyDSL and AITER's FlyDSL helpers ({err})") return reasons @@ -112,10 +112,10 @@ def _refuse_model(model) -> list[str]: kernels read every one of the 384 experts on each rank, TP-sharded, in AITER's A8W4 MXFP4 layout, from CTAs that all stay resident.""" from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend - from vllm.models.deepseek_v41.amd.specialized_decode.mono.common.arch import ( + from vllm.models.deepseek_v41.amd.mono.common.arch import ( target_arch, ) - from vllm.models.deepseek_v41.amd.specialized_decode.mono.runner import BLOCKS + from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS from vllm.platforms.rocm import get_cdna_version reasons = [] From f7fd7fe7415f1aebfcd44a94a555fae00c804298 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 08:51:27 -0500 Subject: [PATCH 62/79] chore: revert pure foramtting changes Signed-off-by: simondanielsson --- .../amd/test_dsv41_gfx942_draft_lm_head.py | 22 +++++++-------- .../amd/test_dsv41_gfx942_fused_markov.py | 4 +-- .../amd/test_dsv41_gfx942_sparse_indexer.py | 4 +-- .../models/amd/test_dsv41_gfx942_topk.py | 4 +-- .../models/amd/test_dsv41_mono_gfx942.py | 27 +++++-------------- tests/models/common/specialized_decode/kit.py | 8 +----- tests/models/test_dsv41_mono_numerics.py | 8 ++---- 7 files changed, 26 insertions(+), 51 deletions(-) diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py b/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py index 551498fcb3bc..2ccd60dc5c78 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py @@ -2,13 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """draft_lm_head.compute_logits, the DSpark draft's base logits on gfx942. -With --specialized-decode on gfx942, the draft's 5 rows of one request get a sixth -row of zeros, so that vLLM's dispatch (rocm_unquantized_gemm) takes the torch -GEMM instead of wvSplitK, and the logits of the first 5 rows come back. Every -other row count, and every call without --specialized-decode, keeps vLLM's dispatch, -which sends up to 5 rows to wvSplitK. The tests count the wvSplitK calls of -vLLM's dispatch on the draft's lm_head shard at TP 4 and compare the logits -with an fp32 GEMM. +With --specialized-decode on gfx942, the draft's 5 rows of one request get a +sixth row of zeros, so that vLLM's dispatch (rocm_unquantized_gemm) takes the +torch GEMM instead of wvSplitK, and the logits of the first 5 rows come back. +Every other row count, and every call without --specialized-decode, keeps +vLLM's dispatch, which sends up to 5 rows to wvSplitK. The tests count the +wvSplitK calls of vLLM's dispatch on the draft's lm_head shard at TP 4 and +compare the logits with an fp32 GEMM. """ import pytest @@ -34,8 +34,8 @@ def _skip_unless_gfx942(): def _set_gate(monkeypatch, on: bool): from vllm.models.deepseek_v41.amd import ops - # The model records --specialized-decode when it is created. These tests have no - # model, so they set the recorded flag themselves. + # The model records --specialized-decode when it is created. These tests + # have no model, so they set the recorded flag themselves. monkeypatch.setattr(ops, "_specialized_decode", on) ops.enabled.cache_clear() @@ -95,8 +95,8 @@ def test_draft_lm_head_takes_5_rows_off_wvsplitk( finally: _set_gate(monkeypatch, on=False) - # Only the 5 rows with --specialized-decode leave wvSplitK. Every other call - # keeps vLLM's dispatch, which takes wvSplitK for up to 5 rows. + # Only the 5 rows with --specialized-decode leave wvSplitK. Every other + # call keeps vLLM's dispatch, which takes wvSplitK for up to 5 rows. padded = specialized_decode and rows == draft_lm_head.ROWS assert wvsplitk_rows == ([] if padded or rows > 5 else [rows]) assert got.shape == (rows, VOCAB_SHARD) diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py index 56955b6ca994..6bbd553dffc4 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py @@ -35,8 +35,8 @@ def _set_gate(monkeypatch, on: bool): from vllm.models.deepseek_v41.amd import ops - # The model records --specialized-decode when it is created. These tests have no - # model, so they set the recorded flag themselves. + # The model records --specialized-decode when it is created. These tests + # have no model, so they set the recorded flag themselves. monkeypatch.setattr(ops, "_specialized_decode", on) ops.enabled.cache_clear() diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py b/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py index f22ffa1c1e42..32d7ea714f66 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py @@ -51,8 +51,8 @@ def _on_gfx942() -> bool: def gfx942_step(monkeypatch): """Turns on AITER and the gfx942 decode launches for one test, and makes the workspace that vLLM's logits call takes its output from. The model - records --specialized-decode when it is created, and this test has no model, so it - sets the recorded flag itself.""" + records --specialized-decode when it is created, and this test has no + model, so it sets the recorded flag itself.""" from vllm._aiter_ops import rocm_aiter_ops from vllm.models.deepseek_v41.amd import ops from vllm.v1.attention.ops.rocm_aiter_mla_sparse import paged_mqa_logits_module diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_topk.py b/tests/kernels/models/amd/test_dsv41_gfx942_topk.py index 99da3154077c..241f2835dd6a 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_topk.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_topk.py @@ -41,8 +41,8 @@ def _on_gfx942() -> bool: @pytest.fixture def gfx942_decode(monkeypatch): """Turns on the gfx942 decode helpers for one test. The model records - --specialized-decode when it is created, and this test has no model, so it sets - the recorded flag itself.""" + --specialized-decode when it is created, and this test has no model, so it + sets the recorded flag itself.""" from vllm.models.deepseek_v41.amd import ops monkeypatch.setattr(ops, "_specialized_decode", True) diff --git a/tests/kernels/models/amd/test_dsv41_mono_gfx942.py b/tests/kernels/models/amd/test_dsv41_mono_gfx942.py index 4fb29c29b1cb..2609e56495ad 100644 --- a/tests/kernels/models/amd/test_dsv41_mono_gfx942.py +++ b/tests/kernels/models/amd/test_dsv41_mono_gfx942.py @@ -47,9 +47,7 @@ def _on_gfx942() -> bool: from flydsl.expr import const_expr, gpu, range_constexpr # noqa: E402 from flydsl.expr.typing import Int64, T # noqa: E402 -from vllm.models.deepseek_v41.amd.mono.attention import ( # noqa: E402 - gemv as G, -) +from vllm.models.deepseek_v41.amd.mono.attention import gemv as G # noqa: E402 from vllm.models.deepseek_v41.amd.mono.attention.device import ( # noqa: E402 THREADS, WAVES, @@ -61,12 +59,8 @@ def _on_gfx942() -> bool: mfma_bf16_k32, ocp_to_fnuz, ) -from vllm.models.deepseek_v41.amd.mono.stages import ( # noqa: E402 - gemv as SG, -) -from vllm.models.deepseek_v41.amd.mono.stages import ( # noqa: E402 - moe as M, -) +from vllm.models.deepseek_v41.amd.mono.stages import gemv as SG # noqa: E402 +from vllm.models.deepseek_v41.amd.mono.stages import moe as M # noqa: E402 from vllm.models.deepseek_v41.amd.mono.weights942 import ( # noqa: E402 EVEN_ODD, codes_tile_rounds, @@ -100,10 +94,7 @@ def _rel_err(got: torch.Tensor, ref: torch.Tensor) -> float: def _keylist_reference(slot, pos, swa_idx, swa_lens, t2r, ratio, topk, bt, cb): - from vllm.models.deepseek_v41.amd.mono.attention.plan import ( - KEYS, - TOPK, - ) + from vllm.models.deepseek_v41.amd.mono.attention.plan import KEYS, TOPK S = slot.shape[0] kt = torch.full((S, KEYS), -1, dtype=torch.int64) @@ -127,11 +118,7 @@ def _keylist_reference(slot, pos, swa_idx, swa_lens, t2r, ratio, topk, bt, cb): @pytest.mark.parametrize("ratio", [0, 1, 2, 4]) def test_keylist_matches_torch_reference(ratio): from vllm.models.deepseek_v41.amd.mono import keylist - from vllm.models.deepseek_v41.amd.mono.attention.plan import ( - KEYS, - TOPK, - WINDOW, - ) + from vllm.models.deepseek_v41.amd.mono.attention.plan import KEYS, TOPK, WINDOW g = torch.Generator().manual_seed(ratio) S, reqs, cb, blocks = 6, 2, 64, 3000 @@ -619,9 +606,7 @@ def test_contrib_942_is_bf16_of_product(routing): every pick of a 42-row step, with w the pick's route weight from the route table (RouteShape.pick_w). down_combine rounds this same f32 product on gfx950, so the gfx942 combine adds the same values.""" - from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( - RouteShape, - ) + from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import RouteShape experts, topk = routing rs = RouteShape(experts, topk, 42, 4) diff --git a/tests/models/common/specialized_decode/kit.py b/tests/models/common/specialized_decode/kit.py index 274657471bae..d712451fad0a 100644 --- a/tests/models/common/specialized_decode/kit.py +++ b/tests/models/common/specialized_decode/kit.py @@ -1,12 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The specialized decode conformance kit (RFC #60904): what a model's own -``DecodeSpec`` must satisfy. - -The library's behaviour is tested once, in -``tests/models/common/specialized_decode/test_specialized_decode.py``. Only what -a model answers for itself belongs here, so a model pays for it in CI. -""" +"""What a model's ``DecodeSpec`` must satisfy.""" import subprocess import sys diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py index c36064386acf..04488957d828 100644 --- a/tests/models/test_dsv41_mono_numerics.py +++ b/tests/models/test_dsv41_mono_numerics.py @@ -310,9 +310,7 @@ def _check_layers(rank: int, tp: int, device) -> None: from vllm.model_executor.model_loader.utils import process_weights_after_loading from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE - from vllm.models.deepseek_v41.amd.mono.layer import ( - scratch_layout, - ) + from vllm.models.deepseek_v41.amd.mono.layer import scratch_layout from vllm.models.deepseek_v41.amd.mono.runner import ( AttnWeights, DSV41MonoLayer, @@ -612,9 +610,7 @@ def _replay_942(rank: int, tp: int, device) -> None: DSV41MonoLayer, MonoLayerWeights, ) - from vllm.models.deepseek_v41.amd.mono.weights942 import ( - linear_copy, - ) + from vllm.models.deepseek_v41.amd.mono.weights942 import linear_copy from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.utils.torch_utils import set_default_torch_dtype From 84f458aa8eaeb83896fdf31074b68213b334518d Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 09:16:10 -0500 Subject: [PATCH 63/79] fix: failing tests after gfx942 merge Signed-off-by: simondanielsson --- tests/models/test_dsv41_specialized_decode.py | 42 +++++++++++++++---- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 9bdad89a0ddd..03fa1d7f964f 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -54,7 +54,7 @@ def ffn(self, *args): def _stub_modules(monkeypatch) -> None: - """The FlyDSL runner package, as a spec-shaped stub.""" + """The FlyDSL runner and arch modules, as spec-shaped stubs.""" runner: Any = ModuleType(f"{_MONO}.runner") runner.BLOCKS = BLOCKS runner.MAX_TOKENS = 64 @@ -63,9 +63,12 @@ def _stub_modules(monkeypatch) -> None: runner.AttnWeights = _weights_stub plan: Any = ModuleType(f"{_MONO}.attention.plan") plan.Dims = lambda world: world + arch: Any = ModuleType(f"{_MONO}.common.arch") + arch.target_arch = lambda: "gfx950" for name, module in ( (f"{_MONO}.runner", runner), (f"{_MONO}.attention.plan", plan), + (f"{_MONO}.common.arch", arch), ): monkeypatch.setitem(sys.modules, name, module) @@ -306,13 +309,27 @@ def _rocm_platform() -> ModuleType: return sys.modules["vllm.platforms.rocm"] -def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS) -> None: +@pytest.fixture(autouse=True) +def _unpin_row_cap(): + """The row cap is cached per process, and these tests move the GPU under + it.""" + sd.max_rows.cache_clear() + yield + sd.max_rows.cache_clear() + + +def _gpu(monkeypatch, cdna: int = 4, cus: int = BLOCKS, target: str = "") -> None: + arch = "gfx942" if cdna == 3 else "gfx950" _stub_modules(monkeypatch) # The compute-unit probe reads the current accelerator, which a CPU run has - # none of. + # none of, and the arch probes read the GPU and FlyDSL's ARCH setting. monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) monkeypatch.setattr(sd.current_platform, "num_compute_units", lambda _: cus) monkeypatch.setattr(_rocm_platform(), "get_cdna_version", lambda: cdna) + monkeypatch.setattr(sd, "_gpu_arch", lambda: arch) + monkeypatch.setattr( + sys.modules[f"{_MONO}.common.arch"], "target_arch", lambda: target or arch + ) def _build(monkeypatch, config=None, model=None, **layer_overrides): @@ -430,15 +447,18 @@ def test_profile_run_takes_vllm_path(monkeypatch): def test_widths_follow_the_capture_sizes(monkeypatch): + _gpu(monkeypatch) config = make_config() config.compilation_config.cudagraph_capture_sizes = [1, 2, 8, 64, 512] assert sd._widths(config) == (1, 2, 8) -def test_widths_default_to_every_row(monkeypatch): +@pytest.mark.parametrize("cdna,rows", [(4, sd.MAX_ROWS), (3, sd.MAX_ROWS_GFX942)]) +def test_widths_default_to_every_row_the_gpu_takes(monkeypatch, cdna, rows): + _gpu(monkeypatch, cdna=cdna) config = make_config() config.compilation_config.cudagraph_capture_sizes = [512] - assert sd._widths(config) == tuple(range(1, sd.MAX_ROWS + 1)) + assert sd._widths(config) == tuple(range(1, rows + 1)) @pytest.mark.parametrize("tp", [1, 3, 8]) @@ -453,12 +473,20 @@ def test_a4w4_moe_refuses(monkeypatch): assert any("A8W4" in r for r in sd._refuse_config(make_config())) -@pytest.mark.parametrize("cdna", [3, 5]) +@pytest.mark.parametrize("cdna", [2, 5]) def test_other_isa_refuses(monkeypatch, cdna): config = make_config() model = TestDsv41Spec().make_model(config) _gpu(monkeypatch, cdna=cdna) - with pytest.raises(DecodeUnsupported, match="CDNA4"): + with pytest.raises(DecodeUnsupported, match="CDNA3/4"): + SpecializedDecode(model, config, select_decode_spec(config, SPECS)) + + +def test_arch_env_against_the_gpu_refuses(monkeypatch): + config = make_config() + model = TestDsv41Spec().make_model(config) + _gpu(monkeypatch, target="gfx942") + with pytest.raises(DecodeUnsupported, match="ARCH"): SpecializedDecode(model, config, select_decode_spec(config, SPECS)) From 3284a1f932684bb92c4c4e2d359a772cb5908a03 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 09:26:12 -0500 Subject: [PATCH 64/79] chore: revert redundant docstring changes Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/mono/__init__.py | 4 +++- .../models/deepseek_v41/amd/mono/common/__init__.py | 6 +++--- vllm/models/deepseek_v41/amd/mono/common/plan.py | 13 +------------ 3 files changed, 7 insertions(+), 16 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py index d8de17ee19bc..3c96927c8c02 100644 --- a/vllm/models/deepseek_v41/amd/mono/__init__.py +++ b/vllm/models/deepseek_v41/amd/mono/__init__.py @@ -1,3 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1-Flash mono decode layer, in two FlyDSL launches (README.md).""" +"""DeepSeek-V4.1-Flash mono decode layer: one decoder layer of a decode step in +two persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The +host side is ``runner``; vLLM docks it in ``..specialized_decode.runtime``.""" diff --git a/vllm/models/deepseek_v41/amd/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/mono/common/__init__.py index 6c5a829f28de..397642c6b441 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/__init__.py +++ b/vllm/models/deepseek_v41/amd/mono/common/__init__.py @@ -3,7 +3,7 @@ # Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, # Copyright (C) 2026, Advanced Micro Devices, Inc.: # atom/mono/__init__.py -"""This model's mono kernels' shared mechanisms: ``ops`` (wave / math -primitives, the TP group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` -(FP8 / MX arithmetic), ``plan`` (execution model, layouts, build keys) and +"""The mono kernels' shared mechanisms: ``ops`` (wave / math primitives, the TP +group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX +arithmetic), ``plan`` (execution model, layouts, build keys) and ``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/deepseek_v41/amd/mono/common/plan.py b/vllm/models/deepseek_v41/amd/mono/common/plan.py index 15ca9cdcb8af..2b05d7a073e9 100644 --- a/vllm/models/deepseek_v41/amd/mono/common/plan.py +++ b/vllm/models/deepseek_v41/amd/mono/common/plan.py @@ -106,18 +106,7 @@ def padded(n: int, align: int) -> int: def source_digest(*paths: str) -> str: """sha256 of every ``.py`` file under ``paths`` (files or directories, relative to the mono package), each file's path and contents: one model's - kernel sources, so an edit to another model's rebuilds none of these. - - Args: - paths: Paths relative to the mono package. - - Returns: - The digest's first 16 hex characters. - - Raises: - FileNotFoundError: A path holds no ``.py`` file. - - """ + kernel sources, so an edit to another model's rebuilds none of these.""" h = hashlib.sha256() for rel in paths: root = _ROOT / rel From 4adf124b96caf32796241a243670447fad2eb81b Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 14:02:26 -0500 Subject: [PATCH 65/79] fix: old import in test Signed-off-by: simondanielsson --- tests/models/test_dsv41_mono_numerics.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py index 04488957d828..cc7378c231a2 100644 --- a/tests/models/test_dsv41_mono_numerics.py +++ b/tests/models/test_dsv41_mono_numerics.py @@ -602,7 +602,6 @@ def _replay_942(rank: int, tp: int, device) -> None: from vllm.distributed import get_tp_group from vllm.engine.arg_utils import EngineArgs from vllm.model_executor.model_loader.utils import process_weights_after_loading - from vllm.models.deepseek_v41.amd import mono_decode from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims from vllm.models.deepseek_v41.amd.mono.runner import ( @@ -611,6 +610,8 @@ def _replay_942(rank: int, tp: int, device) -> None: MonoLayerWeights, ) from vllm.models.deepseek_v41.amd.mono.weights942 import linear_copy + from vllm.models.deepseek_v41.amd.specialized_decode.runtime import weights942 + from vllm.models.deepseek_v41.amd.specialized_decode.spec import INTER_TP4 from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.utils.torch_utils import set_default_torch_dtype @@ -633,18 +634,18 @@ def _replay_942(rank: int, tp: int, device) -> None: with set_current_vllm_config(vc), set_default_torch_dtype(bf), device: moe = DeepseekV4MoE(vc, prefix="model.layers.21.ffn") e, sh = moe.experts.routed_experts, moe.shared_experts - # The hooks that the model sets at load (mono_decode._copy_942_at_load). - # The routed experts are copied from the plain tensors before vLLM's MoE + # The hooks that the model sets at load (weights942.copy_at_load). The + # routed experts are copied from the plain tensors before vLLM's MoE # backend converts them, and the shared expert from the bf16 weights that # vLLM dequantizes at load. - mono_decode._capture_plain_experts(e) + weights942._capture_plain_experts(e) for lin in (sh.gate_up_proj, sh.down_proj): - mono_decode._copy_linear_at_load(lin) + weights942._copy_linear(lin) _random_moe( moe, torch.Generator(device=device).manual_seed(2001 + rank), torch.Generator(device=device).manual_seed(2000), - inter=mono_decode.INTER_TP4, + inter=INTER_TP4, ) process_weights_after_loading(moe, vc.model_config, device) From 85797038fa593e3c490b7eac1ada77912744d353 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 14:02:50 -0500 Subject: [PATCH 66/79] fix: update the recipe Signed-off-by: simondanielsson --- .../deepseek-ai/DeepSeek-V4.1-Flash.yaml | 45 ++++++++++--------- vllm/config/kernel.py | 2 +- 2 files changed, 26 insertions(+), 21 deletions(-) diff --git a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml index a7e734a35d7c..34fe5669f769 100644 --- a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml +++ b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml @@ -41,10 +41,11 @@ model: min_vllm_version: "0.31.0" default_frontend: rust docker_image: - # ROCm 10.0 nightly channel. MI355X's --monokernel (vllm-project/vllm#60397, - # #61082), paged MXFP4 indexer (#58671) and Gluon sparse MLA (#53492) are on - # main but not in 0.31.0, so AMD stays on the nightly image the AMD rows were - # verified on; the MI355X pip tab is hidden in the variants below. + # ROCm 10.0 nightly channel. MI355X's --specialized-decode + # (vllm-project/vllm#60397, #61082, #61311), paged MXFP4 indexer (#58671) + # and Gluon sparse MLA (#53492) are on main but not in 0.31.0, so AMD stays + # on the nightly image the AMD rows were verified on; the MI355X pip tab is + # hidden in the variants below. amd: "vllm/vllm-openai-rocm:nightly-rocm100" architecture: moe parameter_count: "552B" @@ -147,8 +148,9 @@ variants: extra_env: VLLM_USE_V2_MODEL_RUNNER: "1" PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" - # --monokernel and the ROCm MXFP4 indexer path are on main, not in the - # 0.31.0 wheel, which rejects the flag. Drop this once a release ships them. + # --specialized-decode and the ROCm MXFP4 indexer path are on main, not in + # the 0.31.0 wheel, which rejects the flag. Drop this once a release ships + # them. mi355x: install: pip: false @@ -291,12 +293,13 @@ strategy_overrides: # otherwise falls back to 64. - "--block-size" - "128" - # Mono decode kernels (vllm-project/vllm#60397, #61082): two persistent - # launches a layer on all 40 backbone layers for decode steps of up to - # 48 rows, i.e. up to concurrency 8 with five DSpark drafts. Larger - # steps take the regular path, so the gain fades as concurrency rises. - # Replaces the VLLM_ROCM_MONO_DECODE environment variable. - - "--monokernel" + # Specialized decode kernels (vllm-project/vllm#60397, #61082, + # #61311): two persistent launches a layer on all 40 backbone layers + # for decode steps of up to 48 rows, i.e. up to concurrency 8 with + # five DSpark drafts. Larger steps take the regular path, so the gain + # fades as concurrency rises. Replaces the VLLM_ROCM_MONO_DECODE + # environment variable and the --monokernel flag. + - "--specialized-decode" extra_env: # Aiter's Gluon sparse-MLA kernel, gfx950 only # (vllm-project/vllm#53492). @@ -675,9 +678,9 @@ guide: | The launch-day `deepseekv41-flash-0909` image is no longer needed. - On AMD the Docker tab uses `vllm/vllm-openai-rocm:nightly-rocm100`, the ROCm 10.0 nightly the AMD rows were verified on. **MI355X needs it**: the TP2 command uses - `--monokernel` and the ROCm MXFP4 indexer, which are on `main` but not in 0.31.0, - so the Install block offers only Docker on MI355X. MI300X, MI325X and MI350X also - run from the ROCm pip wheel. + `--specialized-decode` and the ROCm MXFP4 indexer, which are on `main` but not + in 0.31.0, so the Install block offers only Docker on MI355X. MI300X, MI325X and + MI350X also run from the ROCm pip wheel. - The Rust OpenAI frontend is selected by default in the command builder. Switch to Python if you encounter unsupported features or compatibility issues. - Expect a long first load: `VLLM_ENGINE_READY_TIMEOUT_S=3600` is set for that reason. @@ -788,14 +791,16 @@ guide: | default `cpu_offload` on through `VLLM_PLE_CPU_OFFLOAD`, so set it explicitly rather than inheriting the default. - `--monokernel` selects the mono decode layer from + `--specialized-decode` selects the mono decode layer from [vllm-project/vllm#60397](https://github.com/vllm-project/vllm/pull/60397), which [vllm-project/vllm#61082](https://github.com/vllm-project/vllm/pull/61082) extended to all 40 backbone layers and moved from the - `VLLM_ROCM_MONO_DECODE` environment variable to this flag. Two persistent - launches cover a layer for decode steps of up to 48 rows, which with five - DSpark drafts is concurrency 8; larger steps fall back to the regular path, so - the gain is largest at low concurrency and fades as concurrency rises. + `VLLM_ROCM_MONO_DECODE` environment variable to a flag, named `--monokernel` + until [vllm-project/vllm#61311](https://github.com/vllm-project/vllm/pull/61311) + renamed it. Two persistent launches cover a layer for decode steps of up to 48 + rows, which with five DSpark drafts is concurrency 8; larger steps fall back + to the regular path, so the gain is largest at low concurrency and fades as + concurrency rises. Decoder SWA bounded replay ([vllm-project/vllm#59039](https://github.com/vllm-project/vllm/pull/59039)) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 7090fb8e24c7..f94dec9f5cfe 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -294,7 +294,7 @@ class KernelConfig: specialized_decode: bool = False """If True, run decode steps on the model's specialized decode: a kernel - that runs most or all of a decoder layer in a single launch. A deployment + that runs most or all of a decoder layer in a single launch. A deployment that cannot run it is refused at startup.""" moe_backend: MoEBackend = "auto" From c927890927c14a636ec8c9682569ae12fc536715 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 18:27:32 -0500 Subject: [PATCH 67/79] fix: bug in launching gfx942 mono Signed-off-by: simondanielsson --- vllm/models/deepseek_v41/amd/dspark.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 2f9815fcaabe..d89b13758bb1 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -576,7 +576,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if self.model.confidence_head is not None and not loaded_confidence_head: self.model.confidence_head = None - self.process_weights_after_loading() logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params From fb577dc5bdf10a475453cc698695d1eae119e27a Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 21 Oct 2026 15:22:08 -0500 Subject: [PATCH 68/79] fix: things broken for minimax after this refactor Signed-off-by: simondanielsson --- .../minimax_m3/amd_specialized_decode/test_model.py | 8 ++++---- .../minimax_m3/amd_specialized_decode/test_runtime.py | 4 ++-- tests/v1/worker/test_spec_decode_embed_sharing_pp.py | 2 +- vllm/models/minimax_m3/amd/specialized_decode/README.md | 2 +- .../minimax_m3/amd/specialized_decode/runtime/spec.py | 2 +- .../minimax_m3/amd/specialized_decode/runtime/weights.py | 2 +- 6 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_model.py b/tests/models/minimax_m3/amd_specialized_decode/test_model.py index 1df6d9071075..b15ef5c5e245 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_model.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Layer- and model-level precision of --monokernel on a mini MiniMax M3 at +"""Layer- and model-level precision of --specialized-decode on a mini MiniMax M3 at the real per-layer shapes, TP4, with random weights (no checkpoint): layers 0-2 dense and stock, 3-4 sparse and fused, the last one ending in the step-end handshake. Engines load dummy weights, so vLLM's real MXFP4 / AITER post-load @@ -12,7 +12,7 @@ residual, post-attention norm, MoE) on captured decode-step inputs, native vs the fused launch; - model: teacher-forced decode (trace replay) with full-vocabulary logprobs, - --monokernel eager and under FULL graphs (with frequent epoch resets) vs + --specialized-decode eager and under FULL graphs (with frequent epoch resets) vs native eager; - fail-stop: a late rank fault under FULL graphs returns no token.""" @@ -337,7 +337,7 @@ def _engine(vllm_runner, model_dir, mono, graphs=False, **extra): kernel_config={ "enable_jit_warmup": False, "moe_backend": "aiter", - "monokernel": mono, + "specialized_decode": mono, }, enforce_eager=not graphs, compilation_config={"cudagraph_mode": "FULL_DECODE_ONLY" if graphs else "NONE"}, @@ -395,7 +395,7 @@ def _reset_state(worker): @pytest.mark.parametrize("graphs", [False, True], ids=["eager", "full_graphs"]) def test_model_matches_native(vllm_runner, model_dir, native, graphs): - """Teacher-forced decode with --monokernel against native eager. Under + """Teacher-forced decode with --specialized-decode against native eager. Under FULL graphs (async scheduling) the epoch resets every 3 sampled steps, on every rank at the same steps, and no wait times out.""" with _engine( diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py index c7cfba8ba068..c36e59cd7f5f 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Host side of --monokernel: config and weight refusals, per-step routing, +"""Host side of --specialized-decode: config and weight refusals, per-step routing, the per-forward plan, cache binding and the fail-stop gate.""" from collections import Counter @@ -105,7 +105,7 @@ def _vllm_config(layers=60, **parallel): resolve_indexer_kv_dtype=lambda _default: "fp8" ), compilation_config=SimpleNamespace(mode=CompilationMode.NONE), - kernel_config=SimpleNamespace(moe_backend="aiter", monokernel=True), + kernel_config=SimpleNamespace(moe_backend="aiter", specialized_decode=True), lora_config=None, kv_transfer_config=SimpleNamespace(kv_connector=None), aux_output_config=SimpleNamespace(enable_return_routed_experts=False), diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py index 9e26d33df6fb..a6c51c123db9 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -160,7 +160,7 @@ def m3_config(): ), quant_config=None, cache_config=None, - kernel_config=SimpleNamespace(monokernel=False), + kernel_config=SimpleNamespace(specialized_decode=False), ) diff --git a/vllm/models/minimax_m3/amd/specialized_decode/README.md b/vllm/models/minimax_m3/amd/specialized_decode/README.md index 0bbdfdb12ff3..d6a8e9da2edf 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/README.md +++ b/vllm/models/minimax_m3/amd/specialized_decode/README.md @@ -1,6 +1,6 @@ # MiniMax M3 specialized decode -`--monokernel` on MI355X/gfx950 runs each sparse layer (3-59) as vLLM's input +`--specialized-decode` on MI355X/gfx950 runs each sparse layer (3-59) as vLLM's input norm and attention followed by one persistent FlyDSL launch for the rest of the layer (`mono/ffn.py`): the attention TP all-reduce and residual add, the Gemma post-attention norm, the FP32 router and top-4, the shared and routed experts, diff --git a/vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py b/vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py index 9172b30fadd8..78758652f14c 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py +++ b/vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py @@ -99,7 +99,7 @@ def refusal_reasons(self, config: VllmConfig) -> list[str]: def specialized_decode_config(config: VllmConfig) -> DecodeConfig | None: - if not config.kernel_config.monokernel: + if not config.kernel_config.specialized_decode: return None layers = config.model_config.hf_text_config.num_hidden_layers spec = replace(DECODE_CONFIG, last_layer=layers - 1) diff --git a/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py b/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py index 80a83157ef70..60185712e3df 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py +++ b/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py @@ -42,7 +42,7 @@ def effective_activation_mode(experts) -> EffectiveActivationMode: def _require(condition: bool, message: str) -> None: if not condition: - raise ValueError(f"--monokernel {message}.") + raise ValueError(f"--specialized-decode {message}.") def _tensor( From adbcbab7f20cf01d25e74308cf62076b7e85583d Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 21 Oct 2026 16:35:55 -0500 Subject: [PATCH 69/79] fix: ddrafter width bug Signed-off-by: simondanielsson --- .../deepseek_v41/amd/specialized_decode/spec.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py index 7e3c167ef4c9..d252ff71effc 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py @@ -85,6 +85,19 @@ def _widths(vllm_config) -> tuple[int, ...]: return served or tuple(range(1, rows + 1)) +def _draft_widths(vllm_config) -> tuple[int, ...]: + """A draft step's rows: one block of draft tokens a request, so multiples of + the block rather than the target's capture sizes. + + Intuitively [block, 2*block, ..., seqs*block or max_rows] + """ + speculative = vllm_config.speculative_config + assert speculative is not None + block = speculative.num_speculative_tokens + seqs = vllm_config.scheduler_config.max_num_seqs + return tuple(range(block, min(block * seqs, max_rows()) + 1, block)) + + def _gpu_arch() -> str: """The GPU's architecture without its feature flags, for example gfx942.""" props = torch.cuda.get_device_properties(torch.accelerator.current_device_index()) @@ -244,7 +257,8 @@ def _build_draft(vllm_config, model, widths: tuple[int, ...]) -> "DSV41DraftKern DRAFT_DECODE_SPEC = DecodeSpec["DSV41DraftKernel"]( name="dsv41-dspark-cdna4", supports=_SUPPORTS, - widths=_widths, + widths=_draft_widths, + set_capture_sizes=False, refuse_config=_refuse_config, refuse_model=_refuse_model, refuse_caches=_refuse_caches, From 52cd25ab681faadc8f9cfe6461b3c01427502b3e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 22 Oct 2026 08:12:28 -0500 Subject: [PATCH 70/79] feat: add --no-monokernel Signed-off-by: simondanielsson --- tests/engine/test_arg_utils.py | 29 +++++++++++++++++++++++++++++ vllm/engine/arg_utils.py | 13 +++++++------ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 27d4493475fb..d7fe98754d88 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -805,6 +805,35 @@ def test_ir_op_priority(): ).create_engine_config() +@pytest.mark.parametrize( + ("argv", "expected"), + [ + ([], None), + (["--specialized-decode"], True), + (["--no-specialized-decode"], False), + ], +) +def test_specialized_decode_cli(argv, expected): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = EngineArgs.from_cli_args(parser.parse_args(argv)) + assert args.specialized_decode is expected + + +@pytest.mark.parametrize( + ("flag", "expected"), [(None, True), (True, True), (False, False)] +) +def test_specialized_decode_overrides_the_kernel_config(flag, expected): + """``--no-specialized-decode`` turns off a specialized decode that a config + file or the Python API asked for, and an absent flag leaves it alone.""" + from vllm.config.kernel import KernelConfig + + config = EngineArgs( + specialized_decode=flag, + kernel_config=KernelConfig(specialized_decode=True), + ).create_engine_config() + assert config.kernel_config.specialized_decode is expected + + @pytest.mark.parametrize( ("input_json", "expected_json"), [ diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 2e80a418a8c9..bf858c409596 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -770,7 +770,7 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) - specialized_decode: bool = get_field(KernelConfig, "specialized_decode") + specialized_decode: bool | None = None worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1829,9 +1829,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) - kernel_group.add_argument( - "--specialized-decode", **kernel_kwargs["specialized_decode"] - ) + specialized_decode_kwargs = kernel_kwargs["specialized_decode"] + # None, not the config's False, so an absent flag overrides nothing. + specialized_decode_kwargs["default"] = None + kernel_group.add_argument("--specialized-decode", **specialized_decode_kwargs) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -2786,8 +2787,8 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune - if self.specialized_decode: - kernel_config.specialized_decode = True + if self.specialized_decode is not None: + kernel_config.specialized_decode = self.specialized_decode if self.moe_backend != "auto": kernel_config.moe_backend = self.moe_backend if self.linear_backend != "auto": From b62a69dde8e1f30647b8cd0ce516beff2c363025 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 22 Oct 2026 08:26:06 -0500 Subject: [PATCH 71/79] revert: all nvidia changes Signed-off-by: simondanielsson --- vllm/models/kimi_k3/nvidia/model.py | 28 ++- vllm/models/kimi_k3/nvidia/model_state.py | 6 +- .../specialized_decode/runtime/__init__.py | 136 ++++++++++++- .../specialized_decode/runtime/decode.py | 112 +++++------ .../specialized_decode/runtime/providers.py | 2 +- .../kimi_k3/nvidia/specialized_decode/spec.py | 185 ------------------ 6 files changed, 206 insertions(+), 263 deletions(-) delete mode 100644 vllm/models/kimi_k3/nvidia/specialized_decode/spec.py diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index ea9c1f5b923d..6368cbaa0f77 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -98,10 +98,6 @@ sp_reduce_scatter, sp_shard, ) -from vllm.models.common.specialized_decode import ( - SpecializedDecode, - select_decode_spec, -) from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.nvidia.latent_moe_runner import ( LatentMoERunner, @@ -126,12 +122,12 @@ KimiK3MultiModalProcessor, KimiK3ProcessingInfo, ) -from .specialized_decode.spec import DECODE_SPEC +from .specialized_decode.runtime import specialized_decode_config if TYPE_CHECKING: from vllm.v1.worker.gpu.model_states.interface import ModelState - from .specialized_decode.runtime.decode import KimiK3Kernel + from .specialized_decode.runtime.decode import SpecializedDecode logger = init_logger(__name__) @@ -450,7 +446,6 @@ def __init__( torch.empty(num_experts, dtype=torch.float32) ) - self.shared_experts: KimiMLP | None if self.num_shared_experts is not None: shared_intermediate_size = moe_intermediate_size * self.num_shared_experts self.shared_experts = KimiMLP( @@ -1437,18 +1432,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.vocab_size, scale=logit_scale ) - # --specialized-decode: a specialized decode executor over self.model, - # built by the model state once the weights are loaded. - self.decode_spec = select_decode_spec(vllm_config, DECODE_SPEC) - self.specialized_decode: SpecializedDecode[KimiK3Kernel] | None = None + # decode_backend=latency_specialized: a specialized decode executor over + # self.model, built by the model state once the weights are loaded. + self.specialized_decode_config = specialized_decode_config(vllm_config) + self.specialized_decode: SpecializedDecode | None = None # The final norm is applied in compute_logits unless the specialized decode # applies it in forward. self.needs_norm_in_logits = True def build_specialized_decode(self) -> None: - assert self.decode_spec is not None + from .specialized_decode.runtime.decode import SpecializedDecode + self.specialized_decode = SpecializedDecode( - self.model, self.vllm_config, self.decode_spec + self.model, self.specialized_decode_config ) # From here on forward() returns normalized hidden states on every step: the # specialized decode's last AttnRes applies the norm, and forward() applies it @@ -1463,7 +1459,7 @@ def get_model_state_cls(self) -> "type[ModelState]": def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: if self.specialized_decode is not None: - return self.specialized_decode.kernel.embed_input_ids(input_ids) + return self.specialized_decode.embed_input_ids(input_ids) return self.model.embed_input_ids(input_ids) def make_empty_intermediate_tensors( @@ -1486,13 +1482,13 @@ def forward( # type: ignore[override] if decode is not None: step = decode.plan(positions, intermediate_tensors) if step is not None: - return decode.run(step, input_ids, inputs_embeds, **kwargs) + return decode.run(step, input_ids, inputs_embeds) hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ) if not self.needs_norm_in_logits: assert decode is not None - hidden_states = decode.kernel.normalize(hidden_states) + hidden_states = decode.normalize(hidden_states) return hidden_states @classmethod diff --git a/vllm/models/kimi_k3/nvidia/model_state.py b/vllm/models/kimi_k3/nvidia/model_state.py index cbb3b2bba29d..d1d0e489b529 100644 --- a/vllm/models/kimi_k3/nvidia/model_state.py +++ b/vllm/models/kimi_k3/nvidia/model_state.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kimi K3's model runner state: the hybrid (KDA + MLA) state, plus the lifecycle of -the specialized decode when ``--specialized-decode`` selects it.""" +the specialized decode when ``decode_backend=latency_specialized`` selects it.""" from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState @@ -16,11 +16,11 @@ def __init__(self, vllm_config, model, encoder_cache, device): if hasattr(model, "get_language_model"): model = model.get_language_model() self.language_model = model - if model.decode_spec is not None: + if model.specialized_decode_config is not None: model.build_specialized_decode() def initialize_kv_cache(self, kv_cache_config, block_tables): super().initialize_kv_cache(kv_cache_config, block_tables) decode = self.language_model.specialized_decode if decode is not None: - decode.bind_caches() + decode.prepare_kv() diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py index a5e3b8dc95cc..dccfb9a5cbe0 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py @@ -1,5 +1,135 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, -enabled by ``--specialized-decode``. What the deployment has to satisfy to run it -is declared in ``..spec``.""" +"""Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, enabled +by ``additional_config={"decode_backend": "latency_specialized"}``.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class DecodeConfig: + max_context: int + mla_core: str + dcp: int + # Tokens per request in a verify step (1 + draft tokens; 1 without speculation). + q: int + capture_sizes: tuple[int, ...] + + +def specialized_decode_config(vllm_config) -> DecodeConfig | None: + """The deployment's specialized decode settings, None for the stock decode. + Raises when the specialized decode is requested for an unsupported deployment.""" + extra = vllm_config.additional_config + backend = extra.get("decode_backend", "auto") if isinstance(extra, dict) else "auto" + if backend == "auto": + return None + if backend != "latency_specialized": + raise ValueError("Supported decode_backend values are auto/latency_specialized") + reason = unsupported(vllm_config) + if reason is not None: + raise ValueError(f"Kimi K3 specialized decode does not support {reason}") + spec = vllm_config.speculative_config + return DecodeConfig( + max_context=vllm_config.model_config.max_model_len, + mla_core=extra.get("mla_core", "auto"), + dcp=vllm_config.parallel_config.decode_context_parallel_size, + q=1 + spec.num_speculative_tokens if spec is not None else 1, + capture_sizes=tuple( + vllm_config.compilation_config.cudagraph_capture_sizes or () + ), + ) + + +def kda_dtypes(vllm_config): + """(conv, recurrent) KDA state dtypes the deployment's cache config selects.""" + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator + + cache = vllm_config.cache_config + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, + cache.mamba_cache_dtype, + cache.mamba_ssm_cache_dtype, + ) + + +def unsupported(vllm_config) -> str | None: + import torch + + from vllm import envs + from vllm.config import CompilationMode, CUDAGraphMode + from vllm.platforms import current_platform + + parallel, spec = vllm_config.parallel_config, vllm_config.speculative_config + text = vllm_config.model_config.hf_text_config + moe = tuple( + getattr(text, name, None) + for name in ( + "hidden_act", + "activation_situ_beta", + "activation_situ_linear_beta", + "num_experts", + "num_experts_per_token", + "num_expert_group", + "topk_group", + "moe_router_activation_func", + ) + ) + compilation = vllm_config.compilation_config + graphs = compilation.cudagraph_mode + checks = ( + (not current_platform.is_device_capability_family(100), "this GPU (SM100/103)"), + ( + moe != ("situ", 4.0, 25.0, 896, 16, 1, 1, "sigmoid"), + "MoE configs but the one-group sigmoid top-16 SiTU MoE of Kimi K3", + ), + (parallel.tensor_parallel_size not in (8, 16), "TP other than 8/16"), + (parallel.pipeline_parallel_size != 1, "pipeline parallelism"), + (parallel.enable_expert_parallel, "expert parallelism"), + (parallel.use_ubatching, "DBO"), + ( + parallel.decode_context_parallel_size + not in (1, parallel.tensor_parallel_size), + "decode context parallel other than 1 or TP", + ), + (vllm_config.model_config.enable_sleep_mode, "sleep mode"), + (vllm_config.lora_config is not None, "LoRA"), + # The runner's ModelState hooks build the executor and prepare the KV layouts. + (not vllm_config.use_v2_model_runner, "the V1 model runner"), + # The decode loop runs in forward; a compiled forward would trace it. + (compilation.mode == CompilationMode.STOCK_TORCH_COMPILE, "torch.compile"), + # The loop runs only in FULL decode graphs; eager and PIECEWISE steps run stock. + ( + graphs is None or graphs.decode_mode() != CUDAGraphMode.FULL, + "deployments without FULL decode CUDA graphs (e.g. enforce-eager)", + ), + ( + spec is not None and spec.method != "dspark", + "speculative methods but DSpark", + ), + ( + spec is not None and spec.enable_adaptive_verification, + "adaptive verification", + ), + # KdaSplit keeps a BF16 conv state and an FP32 recurrent state in place; it + # maintains neither other cache dtypes nor ReplaySSM / RecoverSSM records. + ( + kda_dtypes(vllm_config) != (torch.bfloat16, torch.float32), + "KDA cache dtypes other than a BF16 conv state and an FP32 recurrent " + "state (mamba cache dtype)", + ), + ( + vllm_config.cache_config.use_replayssm + or vllm_config.cache_config.use_kda_recoverssm, + "ReplaySSM / RecoverSSM (--use-replayssm)", + ), + # The verify kernels take at most 8 tokens per request (KdaSplit's barriers). + ( + spec is not None and spec.num_speculative_tokens > 7, + "more than 7 speculative tokens", + ), + ( + spec is not None and envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM, + "the AttnRes aux stream", + ), + ) + return next((reason for failed, reason in checks if failed), None) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py index 38deeab488e9..e1563d38cbbe 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py @@ -3,17 +3,17 @@ """The specialized decode executor over the stock Kimi K3 backbone: one routing decision per step, then every layer's prepared providers in one loop.""" -from collections.abc import Mapping from functools import partial from typing import Any, NamedTuple import torch +from vllm.config import CUDAGraphMode from vllm.distributed import get_tp_group from vllm.distributed import tensor_model_parallel_all_reduce as all_reduce -from vllm.forward_context import get_forward_context -from vllm.models.common.specialized_decode import Step +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up from ..cutedsl.attn_res.prepared import PreparedAttnRes from ..cutedsl.common.embedding import PreparedEmbedding @@ -25,12 +25,11 @@ from ..cutedsl.moe.small_batch import chain as small_batch_chain from ..cutedsl.moe.tail_runtime import MoeTailBox, PreparedTail from ..cutedsl.oproj.prepared import OprojMailbox -from ..spec import PREPARED_MAX, tokens_per_request from .prefetch import PrefetchController -from .providers import Attention, pack_front +from .providers import BATCHES, PREPARED_MAX, Attention, pack_front -class StepPlan(NamedTuple): +class Step(NamedTuple): q: int # tokens per request metadata: list # per layer slot_mapping: Any @@ -66,26 +65,28 @@ def is_uniform_verify(md, q, b): return full and not md.num_decodes and tokens == md.num_spec_decodes * q -class KimiK3Kernel: +def caches_of(attention): + cache = getattr(attention, "kv_cache", ()) + return [cache] if isinstance(cache, torch.Tensor) else list(cache) + + +class SpecializedDecode: """Built once the weights are loaded; owns every prepared provider, the mailboxes and the decode buffers. Reads the backbone, never modifies it.""" - def __init__(self, vllm_config, backbone, widths: tuple[int, ...]): - self.backbone = backbone - self.buffers: dict[str, torch.Tensor] = {} - extra = vllm_config.additional_config - # Verify steps decode q tokens per request at M = requests * q tokens; the - # widths already round each capture size up to whole requests. - self.batches, self.q = tuple(widths), tokens_per_request(vllm_config) - self.queries = (1, self.q) if self.q > 1 else (1,) + def __init__(self, backbone, config): + self.backbone, self.buffers = backbone, {} + q = config.q + # Verify steps decode q tokens per request at M = requests * q tokens; FULL + # verify graphs round each capture size up to whole requests. + rounded = {round_up(s, q) for s in config.capture_sizes} + verify = {m for m in rounded if m <= PREPARED_MAX} + self.batches, self.q = tuple(sorted({*BATCHES, *verify})), q + self.queries = (1, q) if q > 1 else (1,) # The longest sequence the deployment serves bounds the MLA split count. - self.max_context = vllm_config.model_config.max_model_len - self.mla_core = ( - extra.get("mla_core", "auto") if isinstance(extra, dict) else ("auto") - ) + self.max_context, self.mla_core = config.max_context, config.mla_core # Decode context parallel over the TP ranks: MLA KV is sequence-sharded. - self.dcp = vllm_config.parallel_config.decode_context_parallel_size - self.dcp_box, self.mla_layers = None, 0 + self.dcp, self.dcp_box, self.mla_layers = config.dcp, None, 0 group = get_tp_group() self.tp, self.rank = group.world_size, group.rank_in_group self.group = group.device_group @@ -94,12 +95,6 @@ def __init__(self, vllm_config, backbone, widths: tuple[int, ...]): layers = backbone.layers[backbone.start_layer : backbone.end_layer] self.layers = layers = tuple(layers) self.build_providers() - self.cache_layers = tuple(layer_name(d.self_attn) for d in layers) - self.cache_kinds = { - name: self.attention[d.layer_idx].kind - for name, d in zip(self.cache_layers, layers) - } - self.caches: dict[int, list] = {} self.embedding = PreparedEmbedding(backbone.embed_tokens, self.tp, self.rank) self.before = { d.layer_idx: PreparedAttnRes.for_target(d, has_delta=d is not layers[0]) @@ -181,14 +176,12 @@ def build_providers(self): ) self.feed_forward[index] = moe - def bind_caches(self, caches: Mapping[str, Any]) -> None: - """Compile each layer's cores for its bound KV cache layout and keep the - layer's tensors for every launch (the profiling cache, then the real one).""" - for name, layer in zip(self.cache_layers, self.layers): - cache = caches[name] - tensors = [cache] if isinstance(cache, torch.Tensor) else list(cache) - self.caches[layer.layer_idx] = tensors - self.attention[layer.layer_idx].prepare_layout(tensors) + def prepare_kv(self): + """Compile each layer's cores for its bound KV cache layout (the profiling + cache, then the real one).""" + for layer in self.layers: + caches = caches_of(layer.self_attn) + self.attention[layer.layer_idx].prepare_layout(caches) def embed_input_ids(self, ids): rows = ids.shape[0] if ids.ndim == 1 and ids.is_contiguous() else 0 @@ -197,34 +190,44 @@ def embed_input_ids(self, ids): table = self.backbone.embed_tokens.weight return self.embedding.run(ids, table.new_empty((ids.shape[0], table.shape[1]))) - def step_plan(self, width: int) -> StepPlan | None: - """Every layer's metadata for a step of ``width`` rows, or None when the - prepared providers do not serve it. + def plan(self, positions, intermediate_tensors): + """This step's decode plan, or None to run the stock forward instead. - None for steps with prompt rows (prefills, chunked prefills, fresh - prompts), and for speculative steps whose requests do not all verify q - tokens (partial drafts, verify rows mixed with plain decode rows). + Steps that run the stock forward: + - pipeline-parallel stage inputs (``intermediate_tensors``); + - token counts without prepared providers (above 64 or not in ``batches``); + - steps outside a CUDA graph capture (profiling, warmup and the warmup run + before each capture), so every stock kernel runs before its capture; + - PIECEWISE steps, keyed by token count alone and so captured as stock; + - steps with prompt rows (prefills, chunked prefills, fresh prompts); + - speculative steps whose requests do not all verify q tokens (partial + drafts, verify rows mixed with plain decode rows). - Every predicate is the same on all TP ranks: the mailboxes of the - specialized decode need every rank to take the same path. + Every predicate is the same on all TP ranks: the mailboxes of the specialized + decode need every rank to take the same path. """ + b = positions.shape[0] + if intermediate_tensors is not None or b not in self.batches: + return None + if not torch.cuda.is_current_stream_capturing(): + return None + if not is_forward_context_available(): + return None context = get_forward_context() + if context.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + return None by_layer = context.attn_metadata rows: list[Any] = list(by_layer.values()) if isinstance(by_layer, dict) else [] q = self.q if any(getattr(md, "num_spec_decodes", 0) for md in rows) else 1 - if width % q: + if b % q: return None - metadata = [ - decode_metadata(d.self_attn, context, q, width) for d in self.layers - ] + metadata = [decode_metadata(d.self_attn, context, q, b) for d in self.layers] if any(md is None for md in metadata): return None - return StepPlan(q, metadata, context.slot_mapping) + return Step(q, metadata, context.slot_mapping) - def run(self, step: Step, input_ids, inputs_embeds, **kwargs): + def run(self, step, input_ids, inputs_embeds): """Normalized hidden states (and the aux streams) of a planned step.""" - plan = self.step_plan(step.kernel_width) - assert plan is not None x = inputs_embeds if x is None: x = self.embed_input_ids(input_ids) @@ -232,13 +235,12 @@ def run(self, step: Step, input_ids, inputs_embeds, **kwargs): # The stock aux streams: the AttnRes prefix entering each aux layer. aux_out = [torch.empty_like(x) for _ in self.aux_layers] aux = dict(zip(self.aux_layers, aux_out)) - b, q, slots = x.shape[0], plan.q, plan.slot_mapping + b, q, slots = x.shape[0], step.q, step.slot_mapping prefix, blocks = self.prefix[:b], self.blocks[:b] first, second = self.hidden[0][:b], self.hidden[1][:b] hidden, pre_ready = x, False - for layer, md in zip(self.layers, plan.metadata): - index = layer.layer_idx - caches = self.caches[index] + for layer, md in zip(self.layers, step.metadata): + index, caches = layer.layer_idx, caches_of(layer.self_attn) if not pre_ready: self.before[index].run(hidden, prefix, blocks, first, aux.get(index)) self.attention[index].run( diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py index 8bbfc1ada0c4..63e96bfb79b4 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Per-layer attention providers (input projection, KDA/MLA core, o-projection) and -the helpers ``KimiK3Kernel.build_providers`` uses.""" +the helpers ``SpecializedDecode.build_providers`` uses.""" from typing import Any diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py b/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py deleted file mode 100644 index 009d8236e4b0..000000000000 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/spec.py +++ /dev/null @@ -1,185 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kimi K3's specialized decode.""" - -from typing import TYPE_CHECKING - -import torch - -from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature -from vllm.utils.math_utils import round_up - -if TYPE_CHECKING: - from .runtime.decode import KimiK3Kernel - -BATCHES = (*range(1, 9), 16, 32, 64) -"""Decode rows the prepared providers serve whatever the deployment captures.""" - -PREPARED_MAX = BATCHES[-1] -"""The most rows a launch serves; larger steps run vLLM's path.""" - -KDA_STATE_DTYPES = (torch.bfloat16, torch.float32) -"""The (conv, recurrent) KDA states the prepared KDA cores keep in place.""" - -MOE = ("situ", 4.0, 25.0, 896, 16, 1, 1, "sigmoid") -"""The one-group sigmoid top-16 SiTU MoE the prepared experts are packed for.""" - -HIDDEN_SIZE = 7168 -"""The hidden size the decode buffers are shaped for.""" - -MAX_SPECULATIVE_TOKENS = 7 -"""The verify kernels' tokens per request, bounded by KdaSplit's barriers.""" - - -def tokens_per_request(vllm_config) -> int: - """A verify step's tokens per request: 1 + draft tokens, 1 without - speculation.""" - spec = vllm_config.speculative_config - return 1 + spec.num_speculative_tokens if spec is not None else 1 - - -def kda_dtypes(vllm_config): - """(conv, recurrent) KDA state dtypes the deployment's cache config selects.""" - from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator - - cache = vllm_config.cache_config - return MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, - cache.mamba_cache_dtype, - cache.mamba_ssm_cache_dtype, - ) - - -def _widths(vllm_config) -> tuple[int, ...]: - """Every width a launch serves: the prepared batches, plus each captured - size rounded up to whole requests, up to ``PREPARED_MAX``.""" - q = tokens_per_request(vllm_config) - sizes = vllm_config.compilation_config.cudagraph_capture_sizes or () - verify = {round_up(s, q) for s in sizes} - return tuple(sorted({*BATCHES, *(m for m in verify if m <= PREPARED_MAX)})) - - -def _refuse_config(vllm_config) -> list[str]: - """Every reason this config cannot run the kernels, before load.""" - from vllm import envs - from vllm.config import CompilationMode, CUDAGraphMode - from vllm.platforms import current_platform - - parallel, spec = vllm_config.parallel_config, vllm_config.speculative_config - text = vllm_config.model_config.hf_text_config - moe = tuple( - getattr(text, name, None) - for name in ( - "hidden_act", - "activation_situ_beta", - "activation_situ_linear_beta", - "num_experts", - "num_experts_per_token", - "num_expert_group", - "topk_group", - "moe_router_activation_func", - ) - ) - graphs = vllm_config.compilation_config.cudagraph_mode - reasons = [] - if not current_platform.is_device_capability_family(100): - reasons.append("needs an SM100/103 GPU") - if moe != MOE: - reasons.append(f"needs Kimi K3's SiTU MoE {MOE}, not {moe}") - if parallel.tensor_parallel_size not in (8, 16): - reasons.append("needs tensor parallel size 8 or 16") - if parallel.decode_context_parallel_size not in ( - 1, - parallel.tensor_parallel_size, - ): - reasons.append("needs decode context parallel size 1 or the TP size") - # The decode loop runs in forward; a compiled forward would trace it. - if vllm_config.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: - reasons.append("needs a forward that is not torch.compile'd") - # A launch only ever runs while a FULL decode graph is captured. - if graphs is None or graphs.decode_mode() != CUDAGraphMode.FULL: - reasons.append("needs FULL decode CUDA graphs (not enforce-eager)") - if spec is not None: - if spec.method != "dspark": - reasons.append(f"needs the DSpark speculative method, not {spec.method}") - if spec.enable_adaptive_verification: - reasons.append("needs adaptive verification off") - if spec.num_speculative_tokens > MAX_SPECULATIVE_TOKENS: - reasons.append(f"needs at most {MAX_SPECULATIVE_TOKENS} draft tokens") - if envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: - reasons.append("needs the AttnRes aux stream off while speculating") - # KdaSplit keeps a BF16 conv state and an FP32 recurrent state in place; it - # maintains neither other cache dtypes nor ReplaySSM / RecoverSSM records. - if kda_dtypes(vllm_config) != KDA_STATE_DTYPES: - reasons.append( - "needs a BF16 KDA conv state and an FP32 recurrent state " - "(--mamba-cache-dtype)" - ) - if vllm_config.cache_config.use_replayssm or ( - vllm_config.cache_config.use_kda_recoverssm - ): - reasons.append("needs ReplaySSM / RecoverSSM off (--use-replayssm)") - return reasons - - -def _refuse_model(backbone) -> list[str]: - """Every reason the loaded backbone cannot run the kernels.""" - reasons = [] - hidden = getattr(backbone.config, "hidden_size", None) - if hidden != HIDDEN_SIZE: - reasons.append(f"needs hidden size {HIDDEN_SIZE}, the model has {hidden}") - if not getattr(backbone, "num_attn_res_blocks", 0): - reasons.append("needs the backbone's AttnRes blocks") - return reasons - - -def _refuse_caches(vllm_config, kernel: "KimiK3Kernel", caches) -> list[str]: - """Every reason the prepared cores cannot compile against these caches: a - KDA layer's states are kept in place, and an MLA layer's KV is indexed as - one contiguous tensor.""" - reasons = [] - for name, kind in kernel.cache_kinds.items(): - cache = caches.get(name) - tensors = [cache] if isinstance(cache, torch.Tensor) else list(cache or ()) - if kind == "kda": - if len(tensors) < 2: - reasons.append(f"{name}: {len(tensors)} KDA states, needs 2") - continue - dtypes = tuple(t.dtype for t in tensors[:2]) - if dtypes != KDA_STATE_DTYPES: - reasons.append(f"{name}: KDA states {dtypes}, needs {KDA_STATE_DTYPES}") - elif not tensors: - reasons.append(f"{name}: no MLA KV cache") - elif not tensors[0].is_contiguous(): - reasons.append(f"{name}: MLA KV cache is not contiguous") - return reasons - - -def _skip_step(vllm_config, kernel: "KimiK3Kernel", width: int) -> str | None: - """Why this step keeps vLLM's path, or None when the kernels take it.""" - # Every stock kernel has to run before its graph is captured, so a launch - # only ever replaces one inside a capture: profiling, warm-up and the - # warm-up run before each capture all run vLLM's path. - if not torch.cuda.is_current_stream_capturing(): - return "not capturing a CUDA graph" - if kernel.step_plan(width) is None: - return "not a decode or uniform verify step the providers serve" - return None - - -def _build(vllm_config, model, widths: tuple[int, ...]) -> "KimiK3Kernel": - from .runtime.decode import KimiK3Kernel - - return KimiK3Kernel(vllm_config, model, widths) - - -DECODE_SPEC = DecodeSpec["KimiK3Kernel"]( - name="kimi-k3-sm100", - supports=frozenset({Feature.DECODE_CONTEXT_PARALLEL, Feature.SPECULATIVE_DECODING}), - widths=_widths, - refuse_config=_refuse_config, - refuse_model=_refuse_model, - refuse_caches=_refuse_caches, - skip_step=_skip_step, - build=_build, -) From 73aaead19d025509dc522ce20068caf6802f4ccb Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 22 Oct 2026 08:35:59 -0500 Subject: [PATCH 72/79] fix: mentions of monokernel in minimax Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 4 ++-- tests/models/minimax_m3/amd_specialized_decode/test_model.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 11d9a58fa590..b5ba15a36079 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -113,8 +113,8 @@ steps: - 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 diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_model.py b/tests/models/minimax_m3/amd_specialized_decode/test_model.py index 425490453309..9b037aedde30 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_model.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_model.py @@ -506,7 +506,7 @@ def _spec_requests(): def test_eagle3_verify_steps_match_native(vllm_runner, model_dir, draft_dir, capfd): - """EAGLE3 (random draft) with --monokernel under FULL graphs: the verify + """EAGLE3 (random draft) with --specialized-decode under FULL graphs: the verify steps of 1-4 requests (4, 8, 12 padded to 16, and 16 rows) run the fused launches (the device epoch moves on); the target's full-vocabulary logprobs along the generated tokens match native eager teacher-forced on From f9e8d8fa76cf04425d13119befc2a44a6d799ed7 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 22 Oct 2026 08:39:08 -0500 Subject: [PATCH 73/79] fix: tp_agree imports Signed-off-by: simondanielsson --- tests/models/minimax_m3/amd_specialized_decode/test_ops.py | 3 ++- vllm/models/common/specialized_decode/__init__.py | 3 ++- .../minimax_m3/amd/specialized_decode/runtime/decode.py | 5 ++--- .../minimax_m3/amd/specialized_decode/runtime/executor.py | 5 ++--- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py index c72a1fde334b..1cd25b08f720 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py @@ -611,6 +611,7 @@ def test_plain_payload_accesses_are_caught(stores, loads): def _stress_rank(rank, port, results): + from vllm.models.common.specialized_decode import tp_agree from vllm.models.deepseek_v41.amd.mono.common import peer_memory from vllm.models.minimax_m3.amd.specialized_decode.runtime import executor as ex @@ -650,7 +651,7 @@ def inputs(s, which): dist.barrier() # the executor's step hook, its epoch reset agreed over this group - ex.tp_agree = lambda ok: peer_memory.tp_agree(ok, dist.group.WORLD) + ex.tp_agree = lambda ok: tp_agree(ok, dist.group.WORLD) executor = object.__new__(ex.FfnExecutor) executor.rank, executor.runner, executor.peer = rank, runner, peer executor._layers = {0: None} diff --git a/vllm/models/common/specialized_decode/__init__.py b/vllm/models/common/specialized_decode/__init__.py index 76a7b2d5855d..0601c6298be1 100644 --- a/vllm/models/common/specialized_decode/__init__.py +++ b/vllm/models/common/specialized_decode/__init__.py @@ -3,7 +3,7 @@ """Specialized decode: a model's decode step in one launch.""" from vllm.models.common.specialized_decode.dispatch import SpecializedDecode -from vllm.models.common.specialized_decode.guards import select_decode_spec +from vllm.models.common.specialized_decode.guards import select_decode_spec, tp_agree from vllm.models.common.specialized_decode.spec import DecodeSpec, Step __all__ = [ @@ -11,4 +11,5 @@ "SpecializedDecode", "Step", "select_decode_spec", + "tp_agree", ] diff --git a/vllm/models/minimax_m3/amd/specialized_decode/runtime/decode.py b/vllm/models/minimax_m3/amd/specialized_decode/runtime/decode.py index 1de3b65612fb..bbfaa5989acf 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/runtime/decode.py +++ b/vllm/models/minimax_m3/amd/specialized_decode/runtime/decode.py @@ -12,6 +12,7 @@ from vllm.distributed import get_tp_group from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.models.common.specialized_decode import tp_agree as _tp_agree from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseMetadata from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -53,9 +54,7 @@ def close(self) -> None: ... def tp_agree(ok: bool) -> bool: - from vllm.models.deepseek_v41.amd.mono.common.peer_memory import tp_agree - - return tp_agree(ok, get_tp_group().cpu_group) + return _tp_agree(ok, get_tp_group().cpu_group) @dataclass(frozen=True) diff --git a/vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py b/vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py index 4bbe05daf37c..126076f40708 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py +++ b/vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py @@ -11,6 +11,7 @@ from vllm.distributed import get_tp_group from vllm.logger import init_logger +from vllm.models.common.specialized_decode import tp_agree as _tp_agree from vllm.platforms import current_platform from ..mono.ffn import ERROR_BYTES, ERROR_WORDS, peer_bytes @@ -36,9 +37,7 @@ class MonoHealth: def tp_agree(ok: bool) -> bool: - from vllm.models.deepseek_v41.amd.mono.common.peer_memory import tp_agree - - return tp_agree(ok, get_tp_group().cpu_group) + return _tp_agree(ok, get_tp_group().cpu_group) class FfnExecutor: From b6ae90bf1b163984364dfee52bfdbd59e81cd67a Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 22 Oct 2026 16:47:42 -0500 Subject: [PATCH 74/79] fix: revert monokernel rename and add '--specialized-decoded monokernel' alternative Signed-off-by: simondanielsson --- .../deepseek-ai/DeepSeek-V4.1-Flash.yaml | 45 +++++++++---------- tests/engine/test_arg_utils.py | 40 ++++++++++++----- .../amd/test_dsv41_gfx942_draft_lm_head.py | 30 ++++++------- .../amd/test_dsv41_gfx942_fused_markov.py | 4 +- .../amd/test_dsv41_gfx942_sparse_indexer.py | 4 +- .../models/amd/test_dsv41_gfx942_topk.py | 4 +- .../common/specialized_decode/reference.py | 2 +- .../test_specialized_decode.py | 2 +- .../amd_specialized_decode/test_model.py | 10 ++--- .../amd_specialized_decode/test_runtime.py | 6 ++- tests/models/test_dsv41_mono_layers.py | 2 +- tests/models/test_dsv41_specialized_decode.py | 4 +- .../test_spec_decode_embed_sharing_pp.py | 2 +- vllm/config/kernel.py | 12 +++-- vllm/engine/arg_utils.py | 17 +++++-- vllm/models/deepseek_v41/amd/dspark.py | 4 +- vllm/models/deepseek_v41/amd/ops/__init__.py | 16 +++---- .../deepseek_v41/amd/ops/draft_lm_head.py | 4 +- .../deepseek_v41/amd/ops/fused_markov.py | 2 +- .../deepseek_v41/amd/ops/sparse_indexer.py | 6 +-- .../deepseek_v41/amd/ops/swa_decode_meta.py | 2 +- vllm/models/deepseek_v41/amd/ops/topk_pack.py | 2 +- vllm/models/deepseek_v41/amd/rocm.py | 8 ++-- .../specialized_decode/runtime/__init__.py | 2 +- .../amd/specialized_decode/README.md | 2 +- .../amd/specialized_decode/runtime/weights.py | 2 +- 26 files changed, 132 insertions(+), 102 deletions(-) diff --git a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml index 58629c506aff..3cef9fd91a37 100644 --- a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml +++ b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml @@ -41,11 +41,10 @@ model: min_vllm_version: "0.31.0" default_frontend: rust docker_image: - # ROCm 10.0 nightly channel. MI355X's --specialized-decode - # (vllm-project/vllm#60397, #61082, #61311), paged MXFP4 indexer (#58671) - # and Gluon sparse MLA (#53492) are on main but not in 0.31.0, so AMD stays - # on the nightly image the AMD rows were verified on; the MI355X pip tab is - # hidden in the variants below. + # ROCm 10.0 nightly channel. MI355X's --monokernel (vllm-project/vllm#60397, + # #61082), paged MXFP4 indexer (#58671) and Gluon sparse MLA (#53492) are on + # main but not in 0.31.0, so AMD stays on the nightly image the AMD rows were + # verified on; the MI355X pip tab is hidden in the variants below. amd: "vllm/vllm-openai-rocm:nightly-rocm100" architecture: moe parameter_count: "552B" @@ -148,9 +147,8 @@ variants: extra_env: VLLM_USE_V2_MODEL_RUNNER: "1" PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" - # --specialized-decode and the ROCm MXFP4 indexer path are on main, not in - # the 0.31.0 wheel, which rejects the flag. Drop this once a release ships - # them. + # --monokernel and the ROCm MXFP4 indexer path are on main, not in the + # 0.31.0 wheel, which rejects the flag. Drop this once a release ships them. mi355x: install: pip: false @@ -293,13 +291,12 @@ strategy_overrides: # otherwise falls back to 64. - "--block-size" - "128" - # Specialized decode kernels (vllm-project/vllm#60397, #61082, - # #61311): two persistent launches a layer on all 40 backbone layers - # for decode steps of up to 48 rows, i.e. up to concurrency 8 with - # five DSpark drafts. Larger steps take the regular path, so the gain - # fades as concurrency rises. Replaces the VLLM_ROCM_MONO_DECODE - # environment variable and the --monokernel flag. - - "--specialized-decode" + # Mono decode kernels (vllm-project/vllm#60397, #61082): two persistent + # launches a layer on all 40 backbone layers for decode steps of up to + # 48 rows, i.e. up to concurrency 8 with five DSpark drafts. Larger + # steps take the regular path, so the gain fades as concurrency rises. + # Replaces the VLLM_ROCM_MONO_DECODE environment variable. + - "--monokernel" extra_env: # Aiter's Gluon sparse-MLA kernel, gfx950 only # (vllm-project/vllm#53492). @@ -699,9 +696,9 @@ guide: | The launch-day `deepseekv41-flash-0909` image is no longer needed. - On AMD the Docker tab uses `vllm/vllm-openai-rocm:nightly-rocm100`, the ROCm 10.0 nightly the AMD rows were verified on. **MI355X needs it**: the TP2 command uses - `--specialized-decode` and the ROCm MXFP4 indexer, which are on `main` but not - in 0.31.0, so the Install block offers only Docker on MI355X. MI300X, MI325X and - MI350X also run from the ROCm pip wheel. + `--monokernel` and the ROCm MXFP4 indexer, which are on `main` but not in 0.31.0, + so the Install block offers only Docker on MI355X. MI300X, MI325X and MI350X also + run from the ROCm pip wheel. - The Rust OpenAI frontend is selected by default in the command builder. Switch to Python if you encounter unsupported features or compatibility issues. - Expect a long first load: `VLLM_ENGINE_READY_TIMEOUT_S=3600` is set for that reason. @@ -812,16 +809,14 @@ guide: | default `cpu_offload` on through `VLLM_PLE_CPU_OFFLOAD`, so set it explicitly rather than inheriting the default. - `--specialized-decode` selects the mono decode layer from + `--monokernel` selects the mono decode layer from [vllm-project/vllm#60397](https://github.com/vllm-project/vllm/pull/60397), which [vllm-project/vllm#61082](https://github.com/vllm-project/vllm/pull/61082) extended to all 40 backbone layers and moved from the - `VLLM_ROCM_MONO_DECODE` environment variable to a flag, named `--monokernel` - until [vllm-project/vllm#61311](https://github.com/vllm-project/vllm/pull/61311) - renamed it. Two persistent launches cover a layer for decode steps of up to 48 - rows, which with five DSpark drafts is concurrency 8; larger steps fall back - to the regular path, so the gain is largest at low concurrency and fades as - concurrency rises. + `VLLM_ROCM_MONO_DECODE` environment variable to this flag. Two persistent + launches cover a layer for decode steps of up to 48 rows, which with five + DSpark drafts is concurrency 8; larger steps fall back to the regular path, so + the gain is largest at low concurrency and fades as concurrency rises. Decoder SWA bounded replay ([vllm-project/vllm#59039](https://github.com/vllm-project/vllm/pull/59039)) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index d7fe98754d88..348e04c91bb6 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -809,29 +809,49 @@ def test_ir_op_priority(): ("argv", "expected"), [ ([], None), - (["--specialized-decode"], True), - (["--no-specialized-decode"], False), + (["--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)) - assert args.specialized_decode is expected + 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( - ("flag", "expected"), [(None, True), (True, True), (False, False)] + ("flags", "expected"), + [ + ({}, "monokernel"), + ({"monokernel": False}, None), + ({"specialized_decode": "auto"}, "auto"), + ], ) -def test_specialized_decode_overrides_the_kernel_config(flag, expected): - """``--no-specialized-decode`` turns off a specialized decode that a config - file or the Python API asked for, and an absent flag leaves it alone.""" +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( - specialized_decode=flag, - kernel_config=KernelConfig(specialized_decode=True), + kernel_config=KernelConfig(specialized_decode="monokernel"), + **flags, ).create_engine_config() - assert config.kernel_config.specialized_decode is expected + assert config.kernel_config.specialized_decode == expected @pytest.mark.parametrize( diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py b/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py index 2ccd60dc5c78..a94618cbbada 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py @@ -2,13 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """draft_lm_head.compute_logits, the DSpark draft's base logits on gfx942. -With --specialized-decode on gfx942, the draft's 5 rows of one request get a -sixth row of zeros, so that vLLM's dispatch (rocm_unquantized_gemm) takes the -torch GEMM instead of wvSplitK, and the logits of the first 5 rows come back. -Every other row count, and every call without --specialized-decode, keeps -vLLM's dispatch, which sends up to 5 rows to wvSplitK. The tests count the -wvSplitK calls of vLLM's dispatch on the draft's lm_head shard at TP 4 and -compare the logits with an fp32 GEMM. +With --monokernel on gfx942, the draft's 5 rows of one request get a sixth +row of zeros, so that vLLM's dispatch (rocm_unquantized_gemm) takes the torch +GEMM instead of wvSplitK, and the logits of the first 5 rows come back. Every +other row count, and every call without --monokernel, keeps vLLM's dispatch, +which sends up to 5 rows to wvSplitK. The tests count the wvSplitK calls of +vLLM's dispatch on the draft's lm_head shard at TP 4 and compare the logits +with an fp32 GEMM. """ import pytest @@ -34,8 +34,8 @@ def _skip_unless_gfx942(): def _set_gate(monkeypatch, on: bool): from vllm.models.deepseek_v41.amd import ops - # The model records --specialized-decode when it is created. These tests - # have no model, so they set the recorded flag themselves. + # The model records --monokernel when it is created. These tests have no + # model, so they set the recorded flag themselves. monkeypatch.setattr(ops, "_specialized_decode", on) ops.enabled.cache_clear() @@ -75,10 +75,10 @@ def counted(weight, x, *args, **kwargs): return rows -@pytest.mark.parametrize("specialized_decode", [False, True]) +@pytest.mark.parametrize("monokernel", [False, True]) @pytest.mark.parametrize("rows", [1, 4, 5, 6, 10]) def test_draft_lm_head_takes_5_rows_off_wvsplitk( - monkeypatch, wvsplitk_rows, rows, specialized_decode + monkeypatch, wvsplitk_rows, rows, monokernel ): from vllm.models.deepseek_v41.amd.ops import draft_lm_head @@ -88,16 +88,16 @@ def test_draft_lm_head_takes_5_rows_off_wvsplitk( x = torch.randn(rows, HIDDEN, generator=g, device="cuda") head = _LMHead(w.to(torch.bfloat16)) - _set_gate(monkeypatch, on=specialized_decode) + _set_gate(monkeypatch, on=monokernel) try: got = draft_lm_head.compute_logits(_processor, head, x.to(torch.bfloat16)) torch.accelerator.synchronize() finally: _set_gate(monkeypatch, on=False) - # Only the 5 rows with --specialized-decode leave wvSplitK. Every other - # call keeps vLLM's dispatch, which takes wvSplitK for up to 5 rows. - padded = specialized_decode and rows == draft_lm_head.ROWS + # Only the 5 rows with --monokernel leave wvSplitK. Every other call + # keeps vLLM's dispatch, which takes wvSplitK for up to 5 rows. + padded = monokernel and rows == draft_lm_head.ROWS assert wvsplitk_rows == ([] if padded or rows > 5 else [rows]) assert got.shape == (rows, VOCAB_SHARD) want = x.to(torch.bfloat16).float() @ w.to(torch.bfloat16).float().T diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py index 6bbd553dffc4..cc3ab453e306 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py @@ -35,8 +35,8 @@ def _set_gate(monkeypatch, on: bool): from vllm.models.deepseek_v41.amd import ops - # The model records --specialized-decode when it is created. These tests - # have no model, so they set the recorded flag themselves. + # The model records --monokernel when it is created. These tests have no + # model, so they set the recorded flag themselves. monkeypatch.setattr(ops, "_specialized_decode", on) ops.enabled.cache_clear() diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py b/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py index 32d7ea714f66..bdb3b3f02182 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py @@ -51,8 +51,8 @@ def _on_gfx942() -> bool: def gfx942_step(monkeypatch): """Turns on AITER and the gfx942 decode launches for one test, and makes the workspace that vLLM's logits call takes its output from. The model - records --specialized-decode when it is created, and this test has no - model, so it sets the recorded flag itself.""" + records --monokernel when it is created, and this test has no model, so it + sets the recorded flag itself.""" from vllm._aiter_ops import rocm_aiter_ops from vllm.models.deepseek_v41.amd import ops from vllm.v1.attention.ops.rocm_aiter_mla_sparse import paged_mqa_logits_module diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_topk.py b/tests/kernels/models/amd/test_dsv41_gfx942_topk.py index 241f2835dd6a..148a2c9a2712 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_topk.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_topk.py @@ -41,8 +41,8 @@ def _on_gfx942() -> bool: @pytest.fixture def gfx942_decode(monkeypatch): """Turns on the gfx942 decode helpers for one test. The model records - --specialized-decode when it is created, and this test has no model, so it - sets the recorded flag itself.""" + --monokernel when it is created, and this test has no model, so it sets + the recorded flag itself.""" from vllm.models.deepseek_v41.amd import ops monkeypatch.setattr(ops, "_specialized_decode", True) diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py index 051abaa2dd45..c7b85ddbcc23 100644 --- a/tests/models/common/specialized_decode/reference.py +++ b/tests/models/common/specialized_decode/reference.py @@ -164,7 +164,7 @@ def make_config(**overrides) -> Any: """ config = SimpleNamespace( - kernel_config=SimpleNamespace(specialized_decode=True), + kernel_config=SimpleNamespace(specialized_decode="monokernel"), model_config=SimpleNamespace(dtype=torch.bfloat16, enable_sleep_mode=False), scheduler_config=SimpleNamespace(max_num_seqs=4), parallel_config=SimpleNamespace( diff --git a/tests/models/common/specialized_decode/test_specialized_decode.py b/tests/models/common/specialized_decode/test_specialized_decode.py index 27b7f207aafc..77135db85279 100644 --- a/tests/models/common/specialized_decode/test_specialized_decode.py +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -67,7 +67,7 @@ def positions(width: int) -> torch.Tensor: def test_off_without_the_flag(config): - config.kernel_config.specialized_decode = False + config.kernel_config.specialized_decode = None sizes = list(config.compilation_config.cudagraph_capture_sizes) assert select_decode_spec(config, SPECS) is None assert config.compilation_config.cudagraph_capture_sizes == sizes diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_model.py b/tests/models/minimax_m3/amd_specialized_decode/test_model.py index 9b037aedde30..f34a8d4b57f2 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_model.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Layer- and model-level precision of --specialized-decode on a mini MiniMax M3 at +"""Layer- and model-level precision of --monokernel on a mini MiniMax M3 at the real per-layer shapes, TP4, with random weights (no checkpoint): layers 0-2 dense and stock, 3-4 sparse and fused, the last one ending in the step-end handshake. Engines load dummy weights, so vLLM's real MXFP4 / AITER post-load @@ -12,7 +12,7 @@ residual, post-attention norm, MoE) on captured decode-step inputs, native vs the fused launch; - model: teacher-forced decode (trace replay) with full-vocabulary logprobs, - --specialized-decode eager and under FULL graphs (with frequent epoch resets) vs + --monokernel eager and under FULL graphs (with frequent epoch resets) vs native eager; - fail-stop: a late rank fault under FULL graphs returns no token.""" @@ -334,7 +334,7 @@ def _engine(vllm_runner, model_dir, mono, graphs=False, **extra): kernel_config={ "enable_jit_warmup": False, "moe_backend": "aiter", - "specialized_decode": mono, + "specialized_decode": "monokernel" if mono else None, }, enforce_eager=not graphs, compilation_config={"cudagraph_mode": "FULL_DECODE_ONLY" if graphs else "NONE"}, @@ -392,7 +392,7 @@ def _reset_state(worker): @pytest.mark.parametrize("graphs", [False, True], ids=["eager", "full_graphs"]) def test_model_matches_native(vllm_runner, model_dir, native, graphs): - """Teacher-forced decode with --specialized-decode against native eager. Under + """Teacher-forced decode with --monokernel against native eager. Under FULL graphs (async scheduling) the epoch resets every 3 sampled steps, on every rank at the same steps, and no wait times out.""" with _engine( @@ -506,7 +506,7 @@ def _spec_requests(): def test_eagle3_verify_steps_match_native(vllm_runner, model_dir, draft_dir, capfd): - """EAGLE3 (random draft) with --specialized-decode under FULL graphs: the verify + """EAGLE3 (random draft) with --monokernel under FULL graphs: the verify steps of 1-4 requests (4, 8, 12 padded to 16, and 16 rows) run the fused launches (the device epoch moves on); the target's full-vocabulary logprobs along the generated tokens match native eager teacher-forced on diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py index f2311d6a1aec..1e8030166a2b 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Host side of --specialized-decode: config and weight refusals, per-step routing, +"""Host side of --monokernel: config and weight refusals, per-step routing, the per-forward plan, cache binding and the fail-stop gate.""" from collections import Counter @@ -116,7 +116,9 @@ def _vllm_config(layers=60, **parallel): resolve_indexer_kv_dtype=lambda _default: "fp8" ), compilation_config=SimpleNamespace(mode=CompilationMode.NONE), - kernel_config=SimpleNamespace(moe_backend="aiter", specialized_decode=True), + kernel_config=SimpleNamespace( + moe_backend="aiter", specialized_decode="monokernel" + ), lora_config=None, kv_transfer_config=SimpleNamespace(kv_connector=None), aux_output_config=SimpleNamespace(enable_return_routed_experts=False), diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index cc6027e7adb3..d65afc9a143f 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -294,7 +294,7 @@ def _decode( hf_overrides=TRUNCATED, worker_extension_cls=f"{__name__}.MonoProbe", moe_backend="aiter", - specialized_decode=mono, + specialized_decode="monokernel" if mono else None, language_model_only=True, tokenizer_mode="deepseek_v41", **indexer, diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 03fa1d7f964f..98cd7616dbce 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -205,7 +205,7 @@ def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: def make_config(**overrides) -> Any: config = SimpleNamespace( - kernel_config=SimpleNamespace(specialized_decode=True), + kernel_config=SimpleNamespace(specialized_decode="monokernel"), model_config=SimpleNamespace( enable_sleep_mode=False, hf_config=SimpleNamespace( @@ -512,7 +512,7 @@ def test_other_moe_backend_refuses(monkeypatch): def test_off_without_the_flag(monkeypatch): _gpu(monkeypatch) config = make_config() - config.kernel_config.specialized_decode = False + config.kernel_config.specialized_decode = None assert select_decode_spec(config, SPECS) is None diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py index a6c51c123db9..c91e29f1951b 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -160,7 +160,7 @@ def m3_config(): ), quant_config=None, cache_config=None, - kernel_config=SimpleNamespace(specialized_decode=False), + kernel_config=SimpleNamespace(specialized_decode=None), ) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index b26bcc989b38..ef0f72397854 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -266,6 +266,8 @@ def validate_flashinfer_moe_ep_model( "xpu_woq", ] +SpecializedDecodeBackend = Literal["auto", "monokernel"] + @config class KernelConfig: @@ -293,10 +295,12 @@ class KernelConfig: enable_jit_warmup: bool = True """If True, run JIT compile warmup during kernel warmup.""" - specialized_decode: bool = False - """If True, run decode steps on the model's specialized decode: a kernel - that runs most or all of a decoder layer in a single launch. A deployment - that cannot run it is refused at startup.""" + specialized_decode: SpecializedDecodeBackend | None = None + """Which specialized decode to run decode steps on. `monokernel` runs most + or all of a decoder layer in a single persistent launch, which pays off at + the small batch sizes of low-latency serving; `auto` lets vLLM pick, which + today means the monokernel. None runs vLLM's regular decode. A deployment + that cannot run the selection is refused at startup.""" moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index bf858c409596..6ec7540bb3fd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -86,6 +86,7 @@ LinearBackend, MoEBackend, SparseIndexerTopkBackend, + SpecializedDecodeBackend, ) from vllm.config.load import SafetensorsLoadStrategy from vllm.config.logging import LogLevel @@ -770,7 +771,8 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) - specialized_decode: bool | None = None + specialized_decode: SpecializedDecodeBackend | None = None + monokernel: bool | None = None worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1830,9 +1832,16 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: **kernel_kwargs["enable_flashinfer_autotune"], ) specialized_decode_kwargs = kernel_kwargs["specialized_decode"] - # None, not the config's False, so an absent flag overrides nothing. - specialized_decode_kwargs["default"] = None + specialized_decode_kwargs["nargs"] = "?" + specialized_decode_kwargs["const"] = "auto" kernel_group.add_argument("--specialized-decode", **specialized_decode_kwargs) + kernel_group.add_argument( + "--monokernel", + action=argparse.BooleanOptionalAction, + default=None, + help="Shorthand for --specialized-decode monokernel. " + "--no-monokernel turns off a specialized decode asked for elsewhere.", + ) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -2787,6 +2796,8 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune + if self.monokernel is not None: + kernel_config.specialized_decode = "monokernel" if self.monokernel else None if self.specialized_decode is not None: kernel_config.specialized_decode = self.specialized_decode if self.moe_backend != "auto": diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index d89b13758bb1..aa1059c76c74 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -423,7 +423,7 @@ def forward( def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: """Base logits U_k = lm_head(norm(head_hidden)).""" - # On gfx942 with --specialized-decode, the 5 rows of one request take the + # On gfx942 with --monokernel, the 5 rows of one request take the # torch GEMM instead of wvSplitK (ops/draft_lm_head.py). return draft_lm_head.compute_logits( self.logits_processor, self.lm_head, self.model.norm(hidden_states) @@ -456,7 +456,7 @@ def sample_sequential( head_hidden: torch.Tensor, ) -> bool: """DSparkSpeculator._sample_sequential asks this first. With - --specialized-decode on gfx942 each Markov step is one launch + --monokernel on gfx942 each Markov step is one launch (ops/fused_markov.py), and this returns True. For a step that the fused launches do not take, it returns False, having done nothing, and the speculator runs its own loop.""" diff --git a/vllm/models/deepseek_v41/amd/ops/__init__.py b/vllm/models/deepseek_v41/amd/ops/__init__.py index 8689dc62267b..625b5061ac1f 100644 --- a/vllm/models/deepseek_v41/amd/ops/__init__.py +++ b/vllm/models/deepseek_v41/amd/ops/__init__.py @@ -15,15 +15,14 @@ from vllm.platforms import current_platform -# --specialized-decode (KernelConfig.specialized_decode) of this process's -# model. The helpers run in forward passes, where vLLM's current config is not -# set, so the model records the flag when it is created -# (record_specialized_decode). +# --monokernel (KernelConfig.specialized_decode) of this process's model. The +# helpers run in forward passes, where vLLM's current config is not set, so the +# model records the flag when it is created (record_specialized_decode). _specialized_decode = False def record_specialized_decode(vllm_config) -> None: - """Record whether this process's model runs with --specialized-decode. The + """Record whether this process's model runs with --monokernel. The specialized decode records it when the model is created, before any helper runs.""" global _specialized_decode @@ -33,10 +32,9 @@ def record_specialized_decode(vllm_config) -> None: @functools.cache def enabled() -> bool: - """True with --specialized-decode on a gfx942 GPU. Every helper keeps the - result of vLLM's op it replaces, so each one returns False for a call it - cannot take and vLLM then runs its own op. gfx950 always keeps vLLM's ops - here.""" + """True with --monokernel on a gfx942 GPU. Every helper keeps the result + of vLLM's op it replaces, so each one returns False for a call it cannot + take and vLLM then runs its own op. gfx950 always keeps vLLM's ops here.""" if not _specialized_decode or not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_gfx942 diff --git a/vllm/models/deepseek_v41/amd/ops/draft_lm_head.py b/vllm/models/deepseek_v41/amd/ops/draft_lm_head.py index 1126cce406d0..ff95ac7cc3ed 100644 --- a/vllm/models/deepseek_v41/amd/ops/draft_lm_head.py +++ b/vllm/models/deepseek_v41/amd/ops/draft_lm_head.py @@ -11,7 +11,7 @@ take the torch GEMM, and the logits of the sixth row are dropped. That takes 96 us. Every other GEMM keeps vLLM's dispatch. -It runs on gfx942 with --specialized-decode (``ops.enabled()``). +It runs on gfx942 with --monokernel (``ops.enabled()``). tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py checks it. """ @@ -26,7 +26,7 @@ def compute_logits(logits_processor, lm_head, normed: torch.Tensor): """The base logits ``logits_processor(lm_head, normed)``. On gfx942 with - --specialized-decode, 5 rows run as 6 rows with a row of zeros and come back as + --monokernel, 5 rows run as 6 rows with a row of zeros and come back as 5. The processor returns None on a rank that does not get the gathered logits, and so does this function then.""" if normed.shape[0] != ROWS or not enabled(): diff --git a/vllm/models/deepseek_v41/amd/ops/fused_markov.py b/vllm/models/deepseek_v41/amd/ops/fused_markov.py index c15248ef21bf..72408eeb96b3 100644 --- a/vllm/models/deepseek_v41/amd/ops/fused_markov.py +++ b/vllm/models/deepseek_v41/amd/ops/fused_markov.py @@ -25,7 +25,7 @@ the draft proposes, not the target's output, because the rejection sampler verifies every draft token against the target's own distribution. -It runs on gfx942 with --specialized-decode (``ops.enabled()``). The speculator asks +It runs on gfx942 with --monokernel (``ops.enabled()``). The speculator asks the draft model's ``sample_sequential`` (dspark.py) before its loop, and that calls ``sample_sequential`` here. tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py compares it with vLLM's loop. diff --git a/vllm/models/deepseek_v41/amd/ops/sparse_indexer.py b/vllm/models/deepseek_v41/amd/ops/sparse_indexer.py index 2c5d831ab31d..7d737d8f79ac 100644 --- a/vllm/models/deepseek_v41/amd/ops/sparse_indexer.py +++ b/vllm/models/deepseek_v41/amd/ops/sparse_indexer.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The sparse indexer layer of DeepSeek-V4.1 with --specialized-decode on gfx942. +"""The sparse indexer layer of DeepSeek-V4.1 with --monokernel on gfx942. The layer runs vLLM's ROCm indexer op, rocm_aiter_sparse_attn_indexer, on every step except a step of only decode rows. Such a step runs the decode path of that op here, with the gfx942 launches of topk.py in front of vLLM's logits, candidate selection, candidate mask and top-k. Each launch returns False for a call that it cannot take, and then vLLM's op at that place runs. -The model takes this layer only with --specialized-decode on gfx942 +The model takes this layer only with --monokernel on gfx942 (DeepseekV41ROCMAiterMLAAttention._indexer_cls in rocm.py). """ @@ -98,7 +98,7 @@ def gfx942_sparse_attn_indexer( """rocm_aiter_sparse_attn_indexer with the same arguments, whose steps of only decode rows run _decode_step. The profiling run, prefill and mixed steps, a layer that writes its own K cache and a process without - --specialized-decode on gfx942 run vLLM's op unchanged. + --monokernel on gfx942 run vLLM's op unchanged. Like vLLM's op, this is a break of a breakable CUDA graph capture. Its Python runs again at every replay, so the choice between the two paths diff --git a/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py b/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py index f38ee972af61..7fd6e277de14 100644 --- a/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py +++ b/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py @@ -22,7 +22,7 @@ ComputeDSparkNoncausalSWAIndicesKernel in place of the window kernel. That is 3 more builds of 8 launches after every target step. -On gfx942 with --specialized-decode (``ops.enabled()``), decode +On gfx942 with --monokernel (``ops.enabled()``), decode batches of up to MAX_ROWS rows, causal or the draft block, use it. Other batches take the rest of DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build. tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py compares it with the vLLM ops diff --git a/vllm/models/deepseek_v41/amd/ops/topk_pack.py b/vllm/models/deepseek_v41/amd/ops/topk_pack.py index 7315454aa038..09c435e26e03 100644 --- a/vllm/models/deepseek_v41/amd/ops/topk_pack.py +++ b/vllm/models/deepseek_v41/amd/ops/topk_pack.py @@ -17,7 +17,7 @@ ragged entries past indptr[-1], which _compute_global_topk_ragged_indices_and_indptr_old also leaves uninitialized. -On gfx942 with --specialized-decode (``ops.enabled()``), batches +On gfx942 with --monokernel (``ops.enabled()``), batches of up to MAX_ROWS rows use it. tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py compares it with _compute_global_topk_ragged_indices_and_indptr_old. """ diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 6d376c1d8682..ef702e62a046 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -211,7 +211,7 @@ def compute_global_topk_ragged_indices_and_indptr( block_size: int, is_valid_token: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # On gfx942 with --specialized-decode, a decode batch of up to 16 + # On gfx942 with --monokernel, a decode batch of up to 16 # rows is packed with one launch instead of 4 (topk_pack.py). Other # batches return None there and take the function below. from vllm.models.deepseek_v41.amd.ops import topk_pack @@ -380,7 +380,7 @@ def build( fast_build: bool = False, replay_start: torch.Tensor | None = None, ) -> DeepseekV41ROCMAiterSparseSWAMetadata: - # On gfx942 with --specialized-decode, a decode batch gets this + # On gfx942 with --monokernel, a decode batch gets this # metadata from one launch instead of 8 launches per SWA group # (swa_decode_meta.py). Other batches return None there and take # the build below. @@ -579,13 +579,13 @@ def forward_q( class DeepseekV41Gfx942Indexer(DeepseekV4Indexer): """vLLM's FP8 indexer, whose layers run the gfx942 top-512 launches of ops/topk.py on decode steps (ops/sparse_indexer.py). The attention below - takes it only with --specialized-decode on gfx942.""" + takes it only with --monokernel on gfx942.""" attn_cls = Gfx942SparseAttnIndexer def _gfx942_specialized_decode() -> bool: - """True with --specialized-decode on a gfx942 GPU, read from the config of the + """True with --monokernel on a gfx942 GPU, read from the config of the model that is being built. ops.enabled() cannot answer here, because the mono decode layers record the flag only after their attention is built.""" from vllm.platforms.rocm import on_gfx942 diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py index 2fc7c032f0b6..b2fb5a89b576 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DeepSeek-V4.1 specialized decode: a FlyDSL decode launch a layer over the -stock backbone, enabled by ``--specialized-decode``. What the deployment has to +stock backbone, enabled by ``--monokernel``. What the deployment has to satisfy to run it is declared in ``..spec``.""" diff --git a/vllm/models/minimax_m3/amd/specialized_decode/README.md b/vllm/models/minimax_m3/amd/specialized_decode/README.md index d6a8e9da2edf..0bbdfdb12ff3 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/README.md +++ b/vllm/models/minimax_m3/amd/specialized_decode/README.md @@ -1,6 +1,6 @@ # MiniMax M3 specialized decode -`--specialized-decode` on MI355X/gfx950 runs each sparse layer (3-59) as vLLM's input +`--monokernel` on MI355X/gfx950 runs each sparse layer (3-59) as vLLM's input norm and attention followed by one persistent FlyDSL launch for the rest of the layer (`mono/ffn.py`): the attention TP all-reduce and residual add, the Gemma post-attention norm, the FP32 router and top-4, the shared and routed experts, diff --git a/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py b/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py index 60185712e3df..80a83157ef70 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py +++ b/vllm/models/minimax_m3/amd/specialized_decode/runtime/weights.py @@ -42,7 +42,7 @@ def effective_activation_mode(experts) -> EffectiveActivationMode: def _require(condition: bool, message: str) -> None: if not condition: - raise ValueError(f"--specialized-decode {message}.") + raise ValueError(f"--monokernel {message}.") def _tensor( From 3f657f2fc49f56db0278dfd442d98e619bcd9395 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 22 Oct 2026 21:11:54 -0500 Subject: [PATCH 75/79] fix: add kv cache stride assert Signed-off-by: simondanielsson --- tests/models/test_dsv41_specialized_decode.py | 52 +++++++++++++++++-- .../amd/specialized_decode/runtime/decode.py | 21 ++++++-- .../amd/specialized_decode/spec.py | 14 +++-- 3 files changed, 75 insertions(+), 12 deletions(-) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index 98cd7616dbce..cc587496074c 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -22,6 +22,7 @@ LAYERS = 6 HIDDEN = 16 BLOCKS = 256 +BLOCK = 32 # KV records a cache block holds class _StubRunner: @@ -125,7 +126,9 @@ def _attn(index: int, **overrides) -> Any: compress_ratio=1, kv_mxfp8=False, kv_cache_dtype="fp8_ds_mla", - swa_cache_layer=SimpleNamespace(prefix=f"layers.{index}.swa", kv_cache=None), + swa_cache_layer=SimpleNamespace( + prefix=f"layers.{index}.swa", kv_cache=None, block_size=BLOCK + ), compressed_cache_prefix=f"layers.{index}.comp", topk_indices_buffer=torch.zeros(64, 8, dtype=torch.int32), fused_wqa_wkv=_proj(), @@ -199,13 +202,14 @@ def epilogue(self, hidden, *rest, replay=True): return hidden -def make_cache(blocks: int = 4, block: int = 32) -> torch.Tensor: +def make_cache(blocks: int = 4, block: int = BLOCK) -> torch.Tensor: return _bytes(blocks, block * sd.RECORD) def make_config(**overrides) -> Any: config = SimpleNamespace( kernel_config=SimpleNamespace(specialized_decode="monokernel"), + cache_config=SimpleNamespace(block_size=BLOCK), model_config=SimpleNamespace( enable_sleep_mode=False, hf_config=SimpleNamespace( @@ -254,7 +258,7 @@ def _swa(**overrides) -> Any: decode_swa_lens=torch.zeros(8, dtype=torch.int32), token_to_req_indices=torch.zeros(8, dtype=torch.int32), slot_mapping=torch.zeros(8, dtype=torch.int64), - block_size=32, + block_size=BLOCK, ) return SimpleNamespace(**(fields | overrides)) @@ -264,7 +268,7 @@ def _metadata(model, **overrides) -> dict: for layer in model.layers: meta[layer.attn.swa_cache_layer.prefix] = _swa(**overrides) meta[layer.attn.compressed_cache_prefix] = SimpleNamespace( - block_size=32, block_table=torch.zeros(8, 4, dtype=torch.int32) + block_size=BLOCK, block_table=torch.zeros(8, 4, dtype=torch.int32) ) return meta @@ -509,6 +513,46 @@ def test_other_moe_backend_refuses(monkeypatch): SpecializedDecode(model, config, select_decode_spec(config, SPECS)) +def test_cache_whose_blocks_overlap_refuses(monkeypatch): + """A cache whose block stride covers one record but not a whole block puts + every block after the first inside its predecessor. Checking the stride + against RECORD alone passes it and the kernels then read the wrong KV.""" + decode, _ = _build(monkeypatch) + for name in decode.kernel.cache_layers: + narrow = _bytes(4, sd.RECORD) + reasons = decode.spec.refuse_caches( + decode.vllm_config, decode.kernel, {name: narrow} + ) + assert reasons == [ + f"{name}: block stride {sd.RECORD} under {BLOCK * sd.RECORD}" + ] + + +def test_compressed_cache_blocks_hold_a_record_a_ratio_of_tokens(monkeypatch): + """A compressed cache keeps one record per ``compress_ratio`` tokens, so + its blocks are that much shorter than the SWA cache's.""" + decode, _ = _build(monkeypatch, compress_ratio=2) + plan = _plan(decode, 1) + assert decode.kernel.block_records[plan.comp_key] == BLOCK // 2 + assert decode.kernel.block_records[plan.swa_key] == BLOCK + assert not decode.spec.refuse_caches( + decode.vllm_config, + decode.kernel, + {plan.comp_key: make_cache(block=BLOCK // 2)}, + ) + + +def test_record_view_refuses_blocks_wider_than_the_stride(): + """``refuse_caches`` only has the spec's block size; the step brings its + own, so the view settles it there.""" + from vllm.models.deepseek_v41.amd.specialized_decode.runtime import decode as rt + + kv = make_cache(block=BLOCK) + assert rt._record_view(kv, BLOCK).shape == (4, BLOCK, sd.RECORD) + with pytest.raises(AssertionError): + rt._record_view(kv, BLOCK * 2) + + def test_off_without_the_flag(monkeypatch): _gpu(monkeypatch) config = make_config() diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py index 11d48601e6a7..54d01c610fc1 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py @@ -36,8 +36,14 @@ def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's - data rows, then its scale words: the kernels index both from it).""" + data rows, then its scale words: the kernels index both from it). + + A block's stride has to span every record in it, or the next block starts + inside this one. ``_refuse_caches`` screens this at startup from the KV + cache spec; here the step's own block size settles it. + """ kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + assert kv.stride(0) >= block * RECORD, (tuple(kv.shape), kv.stride(), block) return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) @@ -138,8 +144,10 @@ def __init__( self.swa_keys: list[str] = [] self.comp_keys: list[str] = [] self.plans: dict[int, LayerPlan] = {} + self.block_records: dict[str, int] = {} tp_size = vllm_config.parallel_config.tensor_parallel_size + block_size = vllm_config.cache_config.block_size for idx in range(self.embed_layer, self.layer_end): layer = model.layers[idx] plan = self._plan_layer(layer, tp_size, first=idx == self.embed_layer) @@ -148,12 +156,15 @@ def __init__( self.plans[layer.attn.layer_id] = plan if plan.ffn_only: continue - for key, keys in ( - (plan.swa_key, self.swa_keys), - (plan.comp_key, self.comp_keys), + for key, keys, records in ( + (plan.swa_key, self.swa_keys, layer.attn.swa_cache_layer.block_size), + (plan.comp_key, self.comp_keys, block_size // (plan.ratio or 1)), ): - if key is not None and key not in keys: + if key is None: + continue + if key not in keys: keys.append(key) + self.block_records[key] = max(self.block_records.get(key, 0), records) self.cache_layers = tuple(sorted(set(self.swa_keys) | set(self.comp_keys))) def _plan_layer(self, layer, tp_size: int, first: bool) -> LayerPlan | None: diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py index d252ff71effc..e817b484c708 100644 --- a/vllm/models/deepseek_v41/amd/specialized_decode/spec.py +++ b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py @@ -164,7 +164,13 @@ def _refuse_model(model) -> list[str]: def _refuse_caches(vllm_config, decode: "LayerLaunches", caches) -> list[str]: """Every reason the kernels cannot read these caches: they index each as - [blocks, block, RECORD] bytes.""" + [blocks, block, RECORD] bytes, so a block's stride has to span every record + the block holds, or the next block starts inside it. + + ``LayerLaunches.block_records`` holds those counts as the KV cache spec + asks for them, which the runner can only scale up, so this is the floor; + ``_record_view`` settles each step against the block size it is given. + """ reasons = [] for name, cache in caches.items(): if not isinstance(cache, torch.Tensor): @@ -174,8 +180,10 @@ def _refuse_caches(vllm_config, decode: "LayerLaunches", caches) -> list[str]: reasons.append(f"{name}: {cache.dtype} records, needs bytes") if cache.dim() < 2: reasons.append(f"{name}: {cache.dim()} dims, needs [blocks, block, bytes]") - elif cache.stride(0) < RECORD: - reasons.append(f"{name}: block stride {cache.stride(0)} under {RECORD}") + continue + need = decode.block_records.get(name, 1) * RECORD + if cache.stride(0) < need: + reasons.append(f"{name}: block stride {cache.stride(0)} under {need}") return reasons From f3a851efab285c78f2547b9ae45503b40bdae7f0 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 23 Oct 2026 08:46:15 -0500 Subject: [PATCH 76/79] fix: failing tests Signed-off-by: simondanielsson --- .../models/amd/test_dsv41_gfx942_fused_markov.py | 11 ++++++++++- .../specialized_decode/test_specialized_decode.py | 2 +- tests/models/test_dsv41_specialized_decode.py | 4 ++-- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py index cc3ab453e306..f7b91ba0ba1e 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py @@ -25,8 +25,17 @@ from vllm.platforms import current_platform + +def _on_gfx942() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx942 + + return on_gfx942() + + pytestmark = pytest.mark.skipif( - not current_platform.is_rocm(), reason="the fused launch is for ROCm" + not _on_gfx942(), reason="the fused launch is built for gfx942" ) VOCAB, RANK, N_SPEC, MAX_REQS = 129280, 256, 5, 4 diff --git a/tests/models/common/specialized_decode/test_specialized_decode.py b/tests/models/common/specialized_decode/test_specialized_decode.py index 77135db85279..18a9302febb3 100644 --- a/tests/models/common/specialized_decode/test_specialized_decode.py +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -274,7 +274,7 @@ def test_a_spec_that_pads_runs_a_short_step_at_the_next_width(config, model): spec = replace(reference.REFERENCE, pad_steps=True) decode = SpecializedDecode(model, config, spec) decode.bind_caches() - assert decode.widths == (1, 2, 4, 8) + assert decode.widths == reference.widths(config) with forward(config, CUDAGraphMode.FULL, {"prefill": False}): step = decode.plan(positions(3)) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py index cc587496074c..d59ed1eb4676 100644 --- a/tests/models/test_dsv41_specialized_decode.py +++ b/tests/models/test_dsv41_specialized_decode.py @@ -47,8 +47,8 @@ def first_pre_mix(self, rows): def forward(self, *args, **kwargs): return ("whole", args, kwargs) - def ffn(self, *args): - return ("ffn", args) + def ffn(self, *args, **kwargs): + return ("ffn", args, kwargs) _MONO = "vllm.models.deepseek_v41.amd.mono" From 8dd08443c8ddffb1757ba08e4928f0375d491bde Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 23 Oct 2026 08:53:11 -0500 Subject: [PATCH 77/79] docs: add design doc Signed-off-by: simondanielsson --- docs/design/specialized_decode.md | 337 ++++++++++++++++++++++++++++++ 1 file changed, 337 insertions(+) create mode 100644 docs/design/specialized_decode.md diff --git a/docs/design/specialized_decode.md b/docs/design/specialized_decode.md new file mode 100644 index 000000000000..befc36655f27 --- /dev/null +++ b/docs/design/specialized_decode.md @@ -0,0 +1,337 @@ +# Specialized Decode Kernels + +A *specialized decode kernel* replaces part or all of a model's decode step with a +single hand-written launch, which pays off at the small batch sizes of low-latency +serving. The monokernel (FlyDSL, persistent, AMD) is one such kernel; Kimi K3's +CuTeDSL executor is another. This document describes the shared contract they are +written against, which lives in `vllm/models/common/specialized_decode/`. + +The contract exists to (1) improve standardization and (2) limit the lines changed in +model files as more specialized decode kernels are added. The kernels themselves will +naturally remain difficult to review thoroughly, but model code should stay easy to +review. + +## Goal + +**An author of a specialized decode kernel should write:** + +1. A FlyDSL/CuTeDSL implementation and its utilities, in `vllm/models///` +2. The conditions under which it can run +3. Minimal glue connecting it to the model code + +**The author should not have to write:** + +1. Another new way to glue the kernel into the model code +2. JIT/warm-up registration +3. Basic config checks during start up, re-used by many specialized kernels anyway + (parallelism, spec decode, graph mode) +4. Checks of whether to run the specialized kernel or regular vLLM code at every step +5. TP rank agreement of the above +6. Updating of CUDA graph capture sizes +7. Kernel selection when one model has more than one specialized kernel + +The shared library implements the latter points. + +## How + +Declare when the specialized kernel can run: + +```python +# refuse causes crash, skip causes fallback to regular vllm execution + +@dataclass(frozen=True) +class DecodeSpec(Generic[KernelT]): + name: str + supports: frozenset[Feature] + widths: Callable[[VllmConfig], tuple[int, ...]] + refuse_config: Callable[[VllmConfig], list[str]] + refuse_model: Callable[[Any], list[str]] + refuse_caches: Callable[[VllmConfig, KernelT, Mapping[str, Any]], list[str]] + skip_step: Callable[[VllmConfig, KernelT, int], str | None] + build: Callable[[VllmConfig, Any, tuple[int, ...]], KernelT] + pad_steps: bool = False + set_capture_sizes: bool = True +``` + +The two flags cover the two ways a kernel can relate its built widths to the engine's +batch sizes. The default is exact: the served widths become the CUDA graph capture +sizes, and a step of any other row count falls back. With `pad_steps=True` a short step +instead launches at the next width up, which is what `Step.kernel_width` carries, and +such a spec usually also sets `set_capture_sizes=False` so that the deployment keeps the +sizes it asked for. GLM-5.2 wants the second shape: built at `(1, 2, 4, 5, 6, 8)`, it +serves a batch of 3 on the size-4 launch rather than refusing it. + +`KernelT` is a stand-in for the entire model forward, implementing + +```python +class DecodeImpl(Protocol): + cache_layers: tuple[str, ...] + + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: ... + def bind_caches(self, caches: Mapping[str, Any]) -> None: ... +``` + +A kernel that compiles lazily, from a key per launch, can additionally implement + +```python +@runtime_checkable +class DecodeWarmable(Protocol): + def warmup_keys(self) -> list[Any]: ... + def compile(self, key: Any) -> None: ... +``` + +and the dispatcher registers it with vLLM's JIT warm-up, so the first decode step does +not pay the compile. This is optional on purpose. A kernel that compiles eagerly while +it is built, as Kimi K3 does in its provider constructors and again in `bind_caches`, +implements neither method and the dispatcher skips registration. It is also the only +part of the contract with a placement requirement: vLLM's warm-up registry is active +only while the model loads, so a kernel that wants automatic warm-up has to be built in +`process_weights_after_loading`. Building later (Kimi K3 builds in its `ModelState`, +because EAGLE3 aux layers are configured after the loader returns) is fine, it just +means warming up eagerly instead. + +Once the kernel author has written their CuTeDSL/FlyDSL kernel, the two remaining points + +> (2) The conditions under which it can run +> +> (3) Minimal glue connecting it to the model code + +are implemented through a `DecodeSpec` and by wiring the kernel code through a +`DecodeImpl`. The remainder is taken care of by the shared `SpecializedDecode` utility. + +!!! note + There is deliberately no further constraint on how one implements `run`. There is a + big benefit in letting each model decide exactly how to structure this. We have + already seen many different ways of structuring the underlying FlyDSL/CuTeDSL + kernels. + +### Model-side code + +The model-side code looks like this: + +```python +class MyForCausalLM(nn.Module): + def __init__(self, *, vllm_config, prefix=""): + self.decode_spec = select_decode_spec(vllm_config, SPECS) + self.specialized_decode: SpecializedDecode[MyKernel] | None = None + + def process_weights_after_loading(self): + if self.decode_spec is not None: + self.specialized_decode = SpecializedDecode( + self.model, self.vllm_config, self.decode_spec + ) + + def forward(self, input_ids, positions, intermediate_tensors=None, + inputs_embeds=None, **kwargs): + decode = self.specialized_decode + if decode is not None: + step = decode.plan(positions, intermediate_tensors) + if step is not None: + return decode.run(step, input_ids, inputs_embeds, **kwargs) + return self.model(input_ids, positions, intermediate_tensors, + inputs_embeds, **kwargs) +``` + +plus one hook to prepare the KV cache + +```python +class MyModelState(DefaultModelState): + def initialize_kv_cache(self, kv_cache_config, block_tables): + super().initialize_kv_cache(kv_cache_config, block_tables) + if self.model.specialized_decode is not None: + self.model.specialized_decode.bind_caches() +``` + +On the model side nothing else needs to be touched. + +### Directory layout + +"Specialized decode" is the contract. "Monokernel" is one implementation of it (FlyDSL, +persistent kernel, AMD). Only the contract and the host-side glue around it are shared: + +| Path | Holds | +| --- | --- | +| `vllm/models/common/specialized_decode/` | The contract: spec, guards, dispatch. Plain Python, no kernel code and no vendor runtime in its import graph, so it loads anywhere. | +| `vllm/models///specialized_decode/spec.py` | One model's `DecodeSpec`: what the deployment has to satisfy. No kernel backend in its import graph, so selection runs before anything vendor-specific loads. | +| `vllm/models///specialized_decode/runtime/` | That model's `DecodeImpl`, imported lazily from `spec.build`, with its own kernels beneath. | + +Every kernel, including the wave and math primitives, the tagged mailbox, FP8/MX +arithmetic, the execution model and the symmetric peer buffers, stays inside the model +that owns it, for example at `vllm/models/deepseek_v41/amd/mono/`. Kimi K3's `cutedsl/` +and GLM-5.2's `flydsl/` sit in the same place for their models. + +Cross-model-family sharing of kernel code is deliberately out of scope. It needs +validation that has not been done (DeepSeek against Kimi, in particular), and promoting +those files would make every model's kernels move together. Keeping them model-local +costs some duplication and keeps the blast radius of a contract change at the contract. +A shared `/` directory can be added once two families have actually been +cross-checked. + +One consequence worth calling out: a model's JIT cache key is typically seeded by a +digest of every `.py` under that model's kernel package, so an edit to one model's +kernels rebuilds nothing of another's. + +## Adding a second kernel for the same model + +Say DeepSeek-V4.1 gets a gfx942 kernel next to the gfx950 one. Write the FlyDSL kernel, +then: + +```python +# vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode_gfx942.py + +class DSV41Gfx942Kernel(DSV41Kernel): # reuse what you want + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + ... + return self.runner.run_gfx942_kernel(...) + + def compile(self, key): + width, idx = key + self.runner.kernels(width, self.plans[idx]) +``` + +```python +# vllm/models/deepseek_v41/amd/specialized_decode/spec.py + +def _refuse_config_gfx942(vllm_config) -> list[str]: + reasons = [] + if not current_platform.is_gfx942(): + reasons.append("needs a gfx942 GPU") + ... + return reasons + + +def _build_gfx942(vllm_config, model, widths): # the kernel import stays lazy + from .runtime.decode_gfx942 import DSV41Gfx942Kernel + + return DSV41Gfx942Kernel(vllm_config, model, widths) + + +DECODE_SPEC_GFX942 = DecodeSpec["DSV41Gfx942Kernel"]( + name="dsv41-cdna3", + supports=frozenset({Feature.PIPELINE_PARALLEL}), # no spec decode here + widths=lambda cfg: tuple(range(1, 17)), # fewer rows than gfx950 + refuse_config=_refuse_config_gfx942, + refuse_model=_refuse_model, # shared with gfx950 + refuse_caches=_refuse_caches, + skip_step=_skip_step, + build=_build_gfx942, +) + +SPECS = (DECODE_SPEC_GFX950, DECODE_SPEC_GFX942) +``` + +Then one line in the model: + +```python +self.decode_spec = select_decode_spec(vllm_config, SPECS) +``` + +The dispatcher takes the first compatible spec. Specializing on something other than +arch works the same way. A hyper-specialized batch-size 1 kernel can just set +`widths=lambda cfg: (1,)` and be placed first in the `SPECS` list. + +## A kernel that only covers part of a layer + +Sometimes a layer cannot be taken whole: DeepSeek-V4.1 has layers where only the FFN +seam is worth a launch and vLLM's attention has to run. The library has no concept of +this, and does not need one, because the kernel owns the layer loop. It decides per +layer whether to launch the whole thing or to run vLLM's layer with the seam injected: + +```python +class DSV41Kernel: + def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: + ... + for idx in range(self.embed_layer, self.layer_end): + layer, plan = model.layers[idx], self.plans.get(...) + ffn = out = None + if plan.ffn_only: + ffn = partial(self.launch_ffn, plan) # a seam, not a layer + else: + out = self.launch_whole(plan, hidden, positions, ...) + if out is None: + out = layer(hidden, positions, ..., decode_ffn=ffn) + hidden, residual, post_mix, res_mix, pre_mix = out + ... +``` + +`decode_ffn` is an argument on DeepSeek-V4.1's own decoder layer, not a library concept, +so a model that needs a seam adds one and a model that does not never hears about it. +The same shape covers a whole-layer launch that declines at runtime: `launch_whole` +returns None and vLLM's layer runs instead, which is also the per-layer fallback (a +layer with no plan simply runs as usual). + +The build side follows: `warmup_keys` returns a `(width, "ffn")` key for the seam layers +and a whole-layer key for the rest, and `compile` dispatches on it. + +## Worked example: a speculator's draft layers + +A specialized decode can also be put on a speculator's draft head. A model-level kernel +binds to one backbone, and the draft head is the speculator's, not the target's, so the +draft needs its own `DecodeSpec` rather than falling out for free. DeepSeek-V4.1 gives +it one, `DRAFT_DECODE_SPEC`, and the three model touch points go on +`DSparkDeepseekV4ForCausalLM` exactly as they do on the target. + +The two kernels share `LayerLaunches`, which settles everything at load (the plans, the +runner, the warm-up keys); each subclass owns only its layer loop and its routing. The +draft differs in four ways, all of them data: its own routed experts and top-k +(`dspark_n_routed_experts` / `dspark_num_experts_per_tok`, 128 / 3 on V4.1-Flash), one +fp4 ug group instead of two, SWA rows that are a block's non-causal window rather than a +causal decode token's (`get_dspark_swa_index_width`), and a first layer that is handed +the embedding already broadcast to the hc copies. A draft layer the kernels cannot take +whole runs entirely as vLLM's: there is no FFN-only launch for the draft. + +One lifecycle detail: `ModelState` reaches only the target causal LM, so the draft heads +built in this process register themselves and +`DeepseekV41ROCMModelState.initialize_kv_cache` drains the registry, checking the +draft's caches alongside the target's for both the profiling cache and the real one. + +## Worked example: a second GPU architecture + +Supporting a second architecture (DeepSeek-V4.1 on MI300X and MI325X next to MI355X) is +entirely model-specific and lands inside `vllm/models/deepseek_v41/amd/`, leaving the +contract untouched: CDNA3 joins `CDNA_VERSIONS`, the row cap narrows to 42 on gfx942 (a +64 KB LDS workgroup), and the ARCH-vs-GPU refusal and the gfx950-only MoE backend check +go into `_refuse_model`. The per-layer differences (no fused seam requirement, TP 4 +only, vLLM's FP8 indexer instead of the MXFP4 one) sit in `_skip_layer` and `_ffn_only` +next to the checks they qualify. + +Two pieces need a different shape: + +**The weight copies.** The gfx942 kernels read their own re-encoded copies of a layer's +weights, made at load by wrapping each `quant_method.process_weights_after_loading`, so +that they count as model memory when vLLM sizes the KV cache. Those per-layer quant +hooks run *before* the model-level `process_weights_after_loading`, which is where +`SpecializedDecode` is built, so the wrapping cannot live at the build point. It happens +at model construction instead, from `prepare_weights`, called next to +`select_decode_spec` in the causal LM's `__init__`. That function also records the flag +the gfx942 decode helpers read and builds their top-k hooks. The copy machinery is one +file, `runtime/weights942.py`. + +**The draft on gfx942.** There the draft layers take the FFN launch after vLLM's +attention, not a whole-layer launch, so `_plan_layer` returns an FFN plan for them and +the draft's layer loop grows the same `decode_ffn` injection the target's has. +`DSV41DraftKernel` owns its own runner, so on gfx942 it is built at the target's fp4 ug +groups rather than the draft's one, which keeps the kernel configuration that was +validated there. The cost is one extra scratch buffer. + +One mechanical follow-on: `runner.ffn_kernel` is keyed by `(tokens, experts, topk)`, so +the warm-up keys carry the routing. + +## The flag + +A deployment asks for a specialized decode with `--specialized-decode`, whose value is +the backend to run on: + +- `--specialized-decode monokernel` runs most or all of a decoder layer in a single + persistent launch. +- `--specialized-decode` with no value, or `--specialized-decode auto`, lets vLLM pick, + which today means the monokernel. +- Leaving it unset runs vLLM's regular decode. An unknown backend is rejected by argparse + at startup, and a deployment that cannot run the selected backend is refused at + startup rather than falling back silently. + +`--monokernel` is a CLI shorthand for `--specialized-decode monokernel`, and +`--no-monokernel` turns off a specialized decode asked for in a config file or through +the Python API. The underlying config field is +`KernelConfig.specialized_decode: Literal["auto", "monokernel"] | None`, so a future +implementation gets a name rather than a second boolean. From 39047ad595fbc193a384233a47606cec593f6305 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 23 Oct 2026 11:26:21 -0500 Subject: [PATCH 78/79] chore: refine docs Signed-off-by: simondanielsson --- docs/design/specialized_decode.md | 429 ++++++++++++++---------------- 1 file changed, 203 insertions(+), 226 deletions(-) diff --git a/docs/design/specialized_decode.md b/docs/design/specialized_decode.md index befc36655f27..ec68777b61f7 100644 --- a/docs/design/specialized_decode.md +++ b/docs/design/specialized_decode.md @@ -1,67 +1,82 @@ # Specialized Decode Kernels -A *specialized decode kernel* replaces part or all of a model's decode step with a -single hand-written launch, which pays off at the small batch sizes of low-latency -serving. The monokernel (FlyDSL, persistent, AMD) is one such kernel; Kimi K3's -CuTeDSL executor is another. This document describes the shared contract they are -written against, which lives in `vllm/models/common/specialized_decode/`. +## Motivation -The contract exists to (1) improve standardization and (2) limit the lines changed in -model files as more specialized decode kernels are added. The kernels themselves will -naturally remain difficult to review thoroughly, but model code should stay easy to -review. +A **specialized decode kernel** replaces part or all of a model's decode step with a single hand-written launch. At the small batch sizes of low-latency serving, the regular path is dominated by launch overhead, so one persistent kernel can be much faster. The monokernel (FlyDSL, AMD) is one. Kimi K3's CuTeDSL executor is another. More are coming. -## Goal +Each one used to bring its own glue into model files. `vllm/models/common/specialized_decode/` replaces that glue with one contract, so model code remains easy to review as the number of kernels grows. -**An author of a specialized decode kernel should write:** +An author writes the specialized kernel, a `DecodeSpec` saying under what conditions it can run, and a `DecodeImpl` dispatcher to connect the kernel to the model code. The following is then handled by the library: -1. A FlyDSL/CuTeDSL implementation and its utilities, in `vllm/models///` -2. The conditions under which it can run -3. Minimal glue connecting it to the model code +- JIT warm-up registration +- the config and model checks +- per-forward choice between the kernel and vLLM's path +- TP rank agreement on all of it +- setting CUDA graph capture sizes +- and picking between several kernels for one model, for instance based on hardware architecture. -**The author should not have to write:** +Key design principles: -1. Another new way to glue the kernel into the model code -2. JIT/warm-up registration -3. Basic config checks during start up, re-used by many specialized kernels anyway - (parallelism, spec decode, graph mode) -4. Checks of whether to run the specialized kernel or regular vLLM code at every step -5. TP rank agreement of the above -6. Updating of CUDA graph capture sizes -7. Kernel selection when one model has more than one specialized kernel +- **The kernel owns its layer loop**: the library has no opinion about layers, so a kernel may replace one layer, two, or the whole backbone +- **Default deny**: an feature the spec does not list refuses the deployment at startup +- **Refusals carry reasons**: every check returns strings, which are logged and raised +- **Three touch points in the model**: construct, build, forward. Nothing in the decoder layer +- **No kernel code in the contract**: the library is plain Python and imports no vendor runtime, so it loads on any platform -The shared library implements the latter points. +## Quick Overview -## How +### Declaring a Spec -Declare when the specialized kernel can run: +A spec says when the kernel can run and how to build it. It holds no kernel code, so it is evaluated before anything vendor-specific loads: ```python -# refuse causes crash, skip causes fallback to regular vllm execution - -@dataclass(frozen=True) -class DecodeSpec(Generic[KernelT]): - name: str - supports: frozenset[Feature] - widths: Callable[[VllmConfig], tuple[int, ...]] - refuse_config: Callable[[VllmConfig], list[str]] - refuse_model: Callable[[Any], list[str]] - refuse_caches: Callable[[VllmConfig, KernelT, Mapping[str, Any]], list[str]] - skip_step: Callable[[VllmConfig, KernelT, int], str | None] - build: Callable[[VllmConfig, Any, tuple[int, ...]], KernelT] - pad_steps: bool = False - set_capture_sizes: bool = True +# vllm/models/deepseek_v41/amd/specialized_decode/spec.py +from vllm.models.common.specialized_decode.spec import DecodeSpec, Feature + + +def _refuse_config(vllm_config) -> list[str]: + reasons = [] + if vllm_config.parallel_config.tensor_parallel_size not in (2, 4): + reasons.append("needs tensor parallel size 2 or 4") + return reasons + + +def _build(vllm_config, model, widths): + from .runtime.decode import DSV41Kernel # the kernel import stays lazy + + return DSV41Kernel(vllm_config, model, widths) + + +DECODE_SPEC = DecodeSpec["DSV41Kernel"]( + name="dsv41-cdna4", + supports=frozenset({Feature.PIPELINE_PARALLEL, Feature.SPECULATIVE_DECODING}), + widths=lambda cfg: tuple(range(1, 49)), + refuse_config=_refuse_config, + refuse_model=_refuse_model, + refuse_caches=_refuse_caches, + skip_step=_skip_step, + build=_build, +) ``` -The two flags cover the two ways a kernel can relate its built widths to the engine's -batch sizes. The default is exact: the served widths become the CUDA graph capture -sizes, and a step of any other row count falls back. With `pad_steps=True` a short step -instead launches at the next width up, which is what `Step.kernel_width` carries, and -such a spec usually also sets `set_capture_sizes=False` so that the deployment keeps the -sizes it asked for. GLM-5.2 wants the second shape: built at `(1, 2, 4, 5, 6, 8)`, it -serves a batch of 3 on the size-4 launch rather than refusing it. +| Field | Meaning | +| --- | --- | +| `name` | Log and metric name. | +| `supports` | Engine features this spec serves. Default deny. | +| `widths` | The step sizes a launch serves, read from the config. | +| `refuse_config` | Why the config cannot run it. Checked before load. | +| `refuse_model` | Why the loaded model cannot run it. Checked after load. | +| `refuse_caches` | Why the KV caches cannot be read. Checked when they are bound. | +| `skip_step` | Why *this* step should run vLLM's path instead. | +| `build` | Builds this rank's kernel. | +| `pad_steps` | Whether a short step launches at the next width up. Off by default. | +| `set_capture_sizes` | Whether the widths become the capture sizes. On by default. | + +`widths` is a function of the config, so a spec can read the GPU. One kernel may take 48 rows on gfx950 but only 42 on gfx942, where the workgroup has less LDS. -`KernelT` is a stand-in for the entire model forward, implementing +### Implementing the Kernel + +The kernel stands in for the entire model forward: ```python class DecodeImpl(Protocol): @@ -71,7 +86,9 @@ class DecodeImpl(Protocol): def bind_caches(self, caches: Mapping[str, Any]) -> None: ... ``` -A kernel that compiles lazily, from a key per launch, can additionally implement +`run` takes and returns exactly what the model's own forward does. `step` carries the row count, the width to launch at, and the positions. There is deliberately no further constraint on how `run` is written: FlyDSL and CuTeDSL kernels are structured very differently, and forcing one shape would cost more than it saves. + +A kernel that compiles lazily, from one key per launch, can also implement `DecodeWarmable`: ```python @runtime_checkable @@ -80,34 +97,14 @@ class DecodeWarmable(Protocol): def compile(self, key: Any) -> None: ... ``` -and the dispatcher registers it with vLLM's JIT warm-up, so the first decode step does -not pay the compile. This is optional on purpose. A kernel that compiles eagerly while -it is built, as Kimi K3 does in its provider constructors and again in `bind_caches`, -implements neither method and the dispatcher skips registration. It is also the only -part of the contract with a placement requirement: vLLM's warm-up registry is active -only while the model loads, so a kernel that wants automatic warm-up has to be built in -`process_weights_after_loading`. Building later (Kimi K3 builds in its `ModelState`, -because EAGLE3 aux layers are configured after the loader returns) is fine, it just -means warming up eagerly instead. - -Once the kernel author has written their CuTeDSL/FlyDSL kernel, the two remaining points - -> (2) The conditions under which it can run -> -> (3) Minimal glue connecting it to the model code - -are implemented through a `DecodeSpec` and by wiring the kernel code through a -`DecodeImpl`. The remainder is taken care of by the shared `SpecializedDecode` utility. +The dispatcher then registers it with vLLM's JIT warm-up, so the first decode step does not pay the compile. This is optional. Kimi K3 compiles eagerly while it builds, implements neither method, and hence registration is skipped. !!! note - There is deliberately no further constraint on how one implements `run`. There is a - big benefit in letting each model decide exactly how to structure this. We have - already seen many different ways of structuring the underlying FlyDSL/CuTeDSL - kernels. + Warm-up is the one part of the contract with a placement requirement. vLLM's warm-up registry is only active while the model loads, so a kernel that wants automatic warm-up must be built in `process_weights_after_loading`. Building later is fine, it just means warming up eagerly. -### Model-side code +### Wiring It Into the Model -The model-side code looks like this: +Three touch points, all on the causal LM: ```python class MyForCausalLM(nn.Module): @@ -132,7 +129,7 @@ class MyForCausalLM(nn.Module): inputs_embeds, **kwargs) ``` -plus one hook to prepare the KV cache +Plus one hook, so the kernel sees the KV caches once the runner has bound them: ```python class MyModelState(DefaultModelState): @@ -142,75 +139,90 @@ class MyModelState(DefaultModelState): self.model.specialized_decode.bind_caches() ``` -On the model side nothing else needs to be touched. +Nothing else changes. In particular, the decoder layer does not. -### Directory layout +!!! warning + The model must not carry `@support_torch_compile`. A specialized decode replaces the layer loop, so a `torch.compile`'d model is refused with a message saying so. -"Specialized decode" is the contract. "Monokernel" is one implementation of it (FlyDSL, -persistent kernel, AMD). Only the contract and the host-side glue around it are shared: +### Enabling It -| Path | Holds | -| --- | --- | -| `vllm/models/common/specialized_decode/` | The contract: spec, guards, dispatch. Plain Python, no kernel code and no vendor runtime in its import graph, so it loads anywhere. | -| `vllm/models///specialized_decode/spec.py` | One model's `DecodeSpec`: what the deployment has to satisfy. No kernel backend in its import graph, so selection runs before anything vendor-specific loads. | -| `vllm/models///specialized_decode/runtime/` | That model's `DecodeImpl`, imported lazily from `spec.build`, with its own kernels beneath. | +```bash +vllm serve deepseek-ai/DeepSeek-V4.1 --specialized-decode monokernel +vllm serve deepseek-ai/DeepSeek-V4.1 --specialized-decode # auto: vLLM picks +vllm serve deepseek-ai/DeepSeek-V4.1 --monokernel # shorthand, will be deprecated +``` -Every kernel, including the wave and math primitives, the tagged mailbox, FP8/MX -arithmetic, the execution model and the symmetric peer buffers, stays inside the model -that owns it, for example at `vllm/models/deepseek_v41/amd/mono/`. Kimi K3's `cutedsl/` -and GLM-5.2's `flydsl/` sit in the same place for their models. +## Lifecycle + +```text +Model.__init__ + ↓ +[select_decode_spec] → None if the flag is off + → the first spec whose refuse_config() is empty + → raises DecodeUnsupported if none fits + → optionally sets the CUDA graph capture sizes to the widths + ↓ +process_weights_after_loading + ↓ +[SpecializedDecode(...)] → refuse_model(), then agree across the TP group + → spec.build() → the kernel + → register warmup_keys()/compile() with the JIT warm-up + ↓ +ModelState.initialize_kv_cache (runs twice: profiling cache, then the real one) + ↓ +[bind_caches] → refuse_caches(), if ok then kernel.bind_caches() + ↓ +forward, every step + ↓ +[plan] → None (run vLLM's path) if: unserved width, + no forward context, piecewise CUDA graph, caches not bound, + or skip_step() gave a reason + ↓ +[run] → kernel.run(step, input_ids, inputs_embeds, **kwargs) +``` -Cross-model-family sharing of kernel code is deliberately out of scope. It needs -validation that has not been done (DeepSeek against Kimi, in particular), and promoting -those files would make every model's kernels move together. Keeping them model-local -costs some duplication and keeps the blast radius of a contract change at the contract. -A shared `/` directory can be added once two families have actually been -cross-checked. +## Core Concepts -One consequence worth calling out: a model's JIT cache key is typically seeded by a -digest of every `.py` under that model's kernel package, so an edit to one model's -kernels rebuilds nothing of another's. +### Refuse Versus Skip -## Adding a second kernel for the same model +**Refuse** means the deployment is wrong. It raises `DecodeUnsupported` and the engine does not start. Use it for anything a request cannot change. Example: the parallelism, the GPU, the quantization backend, the cache layout. A silent fallback here would hide a misconfiguration behind a quiet performance loss. -Say DeepSeek-V4.1 gets a gfx942 kernel next to the gfx950 one. Write the FlyDSL kernel, -then: +**Skip** means this one step is not for the kernel. `skip_step` returns a reason, `plan()` returns None, and the model's own forward runs. Use it for anything that varies step to step: a prefill in the batch, an unserved row count, missing metadata. -```python -# vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode_gfx942.py +Every rank must reach the same verdict, or the ranks that took the kernel will wait on collectives the others never issue. The library handles this. -class DSV41Gfx942Kernel(DSV41Kernel): # reuse what you want - def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - ... - return self.runner.run_gfx942_kernel(...) +### Features - def compile(self, key): - width, idx = key - self.runner.kernels(width, self.plans[idx]) -``` +`supports` is default deny. The library works out which features the config turned on, subtracts the spec's set, and refuses whatever is left, naming it. `Feature` covers the parallelism modes, LoRA, KV connectors, EPLB, sleep mode, CPU weight offload, weight updates, routed expert capture and ubatching. -```python -# vllm/models/deepseek_v41/amd/specialized_decode/spec.py +Two checks are enabled by default. -def _refuse_config_gfx942(vllm_config) -> list[str]: - reasons = [] - if not current_platform.is_gfx942(): - reasons.append("needs a gfx942 GPU") - ... - return reasons +1. the engine must be on Model Runner V2 +1. the CUDA graph mode is `FULL` +### Widths and Capture Sizes -def _build_gfx942(vllm_config, model, widths): # the kernel import stays lazy - from .runtime.decode_gfx942 import DSV41Gfx942Kernel +A kernel is built for a fixed set of row counts, and a step of any other size cannot launch. There are two ways to handle that: - return DSV41Gfx942Kernel(vllm_config, model, widths) +- **Exact (default).** `select_decode_spec` amends the CUDA graph capture sizes to the served widths, keeping any size larger than the widest launch. A step off the list falls back to vLLM's path. +- **Padded.** With `pad_steps=True` a short step launches at the next width up and the kernel owns the padding. Such a spec usually also sets `set_capture_sizes=False`, so the deployment keeps the capture sizes it asked for. A kernel built at `(1, 2, 4, 5, 6, 8)` then serves a batch of 3 on the size-4 launch. +### Cache Binding +`cache_layers` names the attention layers whose KV caches the kernel reads. `bind_caches()` looks each up in the static forward context, runs `spec.refuse_caches`, and only then hands them to the kernel, which therefore never sees a cache its spec has not accepted. The hook fires twice during startup, once for the profiling cache and once for the real one. `plan()` refuses to run before the first call, so a model that forgets the hook falls back to vLLM's path rather than reading unbound memory. + +## Recipes + +### Two Kernels for One Model + +Write the second kernel, then add a second spec and list it in preference order: + +```python DECODE_SPEC_GFX942 = DecodeSpec["DSV41Gfx942Kernel"]( name="dsv41-cdna3", supports=frozenset({Feature.PIPELINE_PARALLEL}), # no spec decode here widths=lambda cfg: tuple(range(1, 17)), # fewer rows than gfx950 - refuse_config=_refuse_config_gfx942, + refuse_config=_refuse_config_gfx942, # "needs a gfx942 GPU" refuse_model=_refuse_model, # shared with gfx950 refuse_caches=_refuse_caches, skip_step=_skip_step, @@ -220,118 +232,83 @@ DECODE_SPEC_GFX942 = DecodeSpec["DSV41Gfx942Kernel"]( SPECS = (DECODE_SPEC_GFX950, DECODE_SPEC_GFX942) ``` -Then one line in the model: +The model already passes `SPECS` to `select_decode_spec`, so nothing else changes. The dispatcher takes the first spec that accepts the config. Specializing on anything other than the architecture works the same way: a kernel that only handles batch size 1 sets `widths=lambda cfg: (1,)` and goes first in the list. + +### A Kernel That Covers Part of a Layer + +Sometimes we want to write a specialized kernel for part of a layer. For instance, DeepSeek-V4.1 has layers where only the FFN seam uses a specialized kernel but attention runs in normal vLLM code. The library needs no concept of this, because the kernel owns the layer loop and decides per layer: ```python -self.decode_spec = select_decode_spec(vllm_config, SPECS) +# in DecodeImpl.run: +for idx in range(self.embed_layer, self.layer_end): + layer, plan = model.layers[idx], self.plans.get(...) + ffn = out = None + if plan.ffn_only: + ffn = partial(self.launch_ffn, plan) # a seam, not a layer + else: + out = self.launch_whole(plan, hidden, positions, ...) + if out is None: + # inject the specialized kernel into the vLLM execution path + out = layer(hidden, positions, ..., decode_ffn=ffn) ``` -The dispatcher takes the first compatible spec. Specializing on something other than -arch works the same way. A hyper-specialized batch-size 1 kernel can just set -`widths=lambda cfg: (1,)` and be placed first in the `SPECS` list. +The decoder layer can then decide where to use `decode_ffn`. + +### A Kernel on a Speculator's Draft Head + +A drafter needs its own spec, and additionally needs to plug into the target model's KV cache binding hook. + +The draft's causal LM carries the same three touch points as the target: + +```python +# vllm/models/deepseek_v41/amd/dspark.py +class DSparkDeepseekV4ForCausalLM(nn.Module): + def __init__(self, *, vllm_config, prefix=""): + self.specialized_decode_spec = select_decode_spec( + vllm_config, DRAFT_DECODE_SPEC + ) + self.specialized_decode: SpecializedDecode[DSV41DraftKernel] | None = None + self.model = DSparkDeepseekV4Model(vllm_config=vllm_config, ...) + + def process_weights_after_loading(self): + if self.specialized_decode_spec is not None and self.specialized_decode is None: + self.specialized_decode = SpecializedDecode( + self.model, self.vllm_config, self.specialized_decode_spec + ) + # register the specialized drafter + _DRAFT_DECODES.append(self.specialized_decode) -## A kernel that only covers part of a layer + def forward(self, input_ids, positions, inputs_embeds=None): + if self.specialized_decode is not None: + step = self.specialized_decode.plan(positions) + if step is not None: + return self.specialized_decode.run(step, input_ids, inputs_embeds) + return self.model(input_ids, positions, inputs_embeds) +``` -Sometimes a layer cannot be taken whole: DeepSeek-V4.1 has layers where only the FFN -seam is worth a launch and vLLM's attention has to run. The library has no concept of -this, and does not need one, because the kernel owns the layer loop. It decides per -layer whether to launch the whole thing or to run vLLM's layer with the seam injected: +Since the drafter doesn't have its own `ModelState`, we leverage the target model's model state to bind the drafter's caches. We do this by amending the target's `initialize_kv_cache` to also bind the cache of any registered drafter: ```python -class DSV41Kernel: - def run(self, step: Step, input_ids, inputs_embeds, **kwargs) -> Any: - ... - for idx in range(self.embed_layer, self.layer_end): - layer, plan = model.layers[idx], self.plans.get(...) - ffn = out = None - if plan.ffn_only: - ffn = partial(self.launch_ffn, plan) # a seam, not a layer - else: - out = self.launch_whole(plan, hidden, positions, ...) - if out is None: - out = layer(hidden, positions, ..., decode_ffn=ffn) - hidden, residual, post_mix, res_mix, pre_mix = out - ... +# vllm/models/deepseek_v41/amd/dspark.py +_DRAFT_DECODES: list[SpecializedDecode] = [] + + +def bind_draft_decode_caches() -> None: + for decode in _DRAFT_DECODES: + decode.bind_caches() + + +# vllm/models/deepseek_v41/amd/model_state.py +class DeepseekV41ROCMModelState(DeepseekV41ModelState): + def initialize_kv_cache(self, kv_cache_config, block_tables): + super().initialize_kv_cache(kv_cache_config, block_tables) + if self.model.specialized_decode is not None: + self.model.specialized_decode.bind_caches() + # NEW + bind_draft_decode_caches() ``` -`decode_ffn` is an argument on DeepSeek-V4.1's own decoder layer, not a library concept, -so a model that needs a seam adds one and a model that does not never hears about it. -The same shape covers a whole-layer launch that declines at runtime: `launch_whole` -returns None and vLLM's layer runs instead, which is also the per-layer fallback (a -layer with no plan simply runs as usual). - -The build side follows: `warmup_keys` returns a `(width, "ffn")` key for the seam layers -and a whole-layer key for the rest, and `compile` dispatches on it. - -## Worked example: a speculator's draft layers - -A specialized decode can also be put on a speculator's draft head. A model-level kernel -binds to one backbone, and the draft head is the speculator's, not the target's, so the -draft needs its own `DecodeSpec` rather than falling out for free. DeepSeek-V4.1 gives -it one, `DRAFT_DECODE_SPEC`, and the three model touch points go on -`DSparkDeepseekV4ForCausalLM` exactly as they do on the target. - -The two kernels share `LayerLaunches`, which settles everything at load (the plans, the -runner, the warm-up keys); each subclass owns only its layer loop and its routing. The -draft differs in four ways, all of them data: its own routed experts and top-k -(`dspark_n_routed_experts` / `dspark_num_experts_per_tok`, 128 / 3 on V4.1-Flash), one -fp4 ug group instead of two, SWA rows that are a block's non-causal window rather than a -causal decode token's (`get_dspark_swa_index_width`), and a first layer that is handed -the embedding already broadcast to the hc copies. A draft layer the kernels cannot take -whole runs entirely as vLLM's: there is no FFN-only launch for the draft. - -One lifecycle detail: `ModelState` reaches only the target causal LM, so the draft heads -built in this process register themselves and -`DeepseekV41ROCMModelState.initialize_kv_cache` drains the registry, checking the -draft's caches alongside the target's for both the profiling cache and the real one. - -## Worked example: a second GPU architecture - -Supporting a second architecture (DeepSeek-V4.1 on MI300X and MI325X next to MI355X) is -entirely model-specific and lands inside `vllm/models/deepseek_v41/amd/`, leaving the -contract untouched: CDNA3 joins `CDNA_VERSIONS`, the row cap narrows to 42 on gfx942 (a -64 KB LDS workgroup), and the ARCH-vs-GPU refusal and the gfx950-only MoE backend check -go into `_refuse_model`. The per-layer differences (no fused seam requirement, TP 4 -only, vLLM's FP8 indexer instead of the MXFP4 one) sit in `_skip_layer` and `_ffn_only` -next to the checks they qualify. - -Two pieces need a different shape: - -**The weight copies.** The gfx942 kernels read their own re-encoded copies of a layer's -weights, made at load by wrapping each `quant_method.process_weights_after_loading`, so -that they count as model memory when vLLM sizes the KV cache. Those per-layer quant -hooks run *before* the model-level `process_weights_after_loading`, which is where -`SpecializedDecode` is built, so the wrapping cannot live at the build point. It happens -at model construction instead, from `prepare_weights`, called next to -`select_decode_spec` in the causal LM's `__init__`. That function also records the flag -the gfx942 decode helpers read and builds their top-k hooks. The copy machinery is one -file, `runtime/weights942.py`. - -**The draft on gfx942.** There the draft layers take the FFN launch after vLLM's -attention, not a whole-layer launch, so `_plan_layer` returns an FFN plan for them and -the draft's layer loop grows the same `decode_ffn` injection the target's has. -`DSV41DraftKernel` owns its own runner, so on gfx942 it is built at the target's fp4 ug -groups rather than the draft's one, which keeps the kernel configuration that was -validated there. The cost is one extra scratch buffer. - -One mechanical follow-on: `runner.ffn_kernel` is keyed by `(tokens, experts, topk)`, so -the warm-up keys carry the routing. - -## The flag - -A deployment asks for a specialized decode with `--specialized-decode`, whose value is -the backend to run on: - -- `--specialized-decode monokernel` runs most or all of a decoder layer in a single - persistent launch. -- `--specialized-decode` with no value, or `--specialized-decode auto`, lets vLLM pick, - which today means the monokernel. -- Leaving it unset runs vLLM's regular decode. An unknown backend is rejected by argparse - at startup, and a deployment that cannot run the selected backend is refused at - startup rather than falling back silently. - -`--monokernel` is a CLI shorthand for `--specialized-decode monokernel`, and -`--no-monokernel` turns off a specialized decode asked for in a config file or through -the Python API. The underlying config field is -`KernelConfig.specialized_decode: Literal["auto", "monokernel"] | None`, so a future -implementation gets a name rather than a second boolean. +## See Also + +- [CUDA Graphs](cuda_graphs.md), the capture sizes a spec amends +- [Model Runner V2](model_runner_v2.md), the runner a specialized decode requires From b426bf7bea3c5f8dc185a0655bf781f68ff2615c Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 29 Oct 2026 09:27:54 -0500 Subject: [PATCH 79/79] fix: remove --specialized-decode option without specying variant or auto Signed-off-by: simondanielsson --- docs/design/specialized_decode.md | 4 ++-- tests/engine/test_arg_utils.py | 4 +--- vllm/engine/arg_utils.py | 7 +++---- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/docs/design/specialized_decode.md b/docs/design/specialized_decode.md index ec68777b61f7..c4dafb2cf402 100644 --- a/docs/design/specialized_decode.md +++ b/docs/design/specialized_decode.md @@ -148,8 +148,8 @@ Nothing else changes. In particular, the decoder layer does not. ```bash vllm serve deepseek-ai/DeepSeek-V4.1 --specialized-decode monokernel -vllm serve deepseek-ai/DeepSeek-V4.1 --specialized-decode # auto: vLLM picks -vllm serve deepseek-ai/DeepSeek-V4.1 --monokernel # shorthand, will be deprecated +vllm serve deepseek-ai/DeepSeek-V4.1 --specialized-decode auto # vLLM picks +vllm serve deepseek-ai/DeepSeek-V4.1 --monokernel # shorthand, will be deprecated ``` ## Lifecycle diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 86c83a416f2c..ebabc32210be 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -852,7 +852,6 @@ def test_ir_op_priority_deprecation_warning(caplog_vllm, disable_log_dedup): ("argv", "expected"), [ ([], None), - (["--specialized-decode"], "auto"), (["--specialized-decode", "auto"], "auto"), (["--specialized-decode", "monokernel"], "monokernel"), (["--monokernel"], "monokernel"), @@ -863,8 +862,7 @@ def test_ir_op_priority_deprecation_warning(caplog_vllm, disable_log_dedup): ], ) def test_specialized_decode_cli(argv, expected): - """``--monokernel`` is shorthand for ``--specialized-decode monokernel``, - and a bare ``--specialized-decode`` means ``auto``.""" + """``--monokernel`` is shorthand for ``--specialized-decode monokernel``.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = EngineArgs.from_cli_args(parser.parse_args(argv)) config = args.create_engine_config() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b8516ad5266b..9f85abe5968b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1846,10 +1846,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) - specialized_decode_kwargs = kernel_kwargs["specialized_decode"] - specialized_decode_kwargs["nargs"] = "?" - specialized_decode_kwargs["const"] = "auto" - kernel_group.add_argument("--specialized-decode", **specialized_decode_kwargs) + kernel_group.add_argument( + "--specialized-decode", **kernel_kwargs["specialized_decode"] + ) kernel_group.add_argument( "--monokernel", action=argparse.BooleanOptionalAction,