diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 0fcd068a2f08..36dae49f6f52 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -76,6 +76,7 @@ steps: - tests/models/test_deepseek_v41_pipeline.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py + - tests/models/common/specialized_decode/ - tests/models/test_deepseek_v4_vl_stream_weights.py - tests/models/test_muse_spark_rocm_loading.py - tests/models/test_muse_spark_logits_local.py @@ -97,6 +98,7 @@ steps: models/test_muse_spark_rocm_loading.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py + models/common/specialized_decode mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -119,7 +121,7 @@ steps: - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - - tests/models/test_dsv41_mono_decode.py + - tests/models/test_dsv41_specialized_decode.py - tests/models/test_muse_spark_rocm_loading.py - tests/models/test_dsv4_mono_decode.py - tests/models/test_deepseek_v4_vl_stream_weights.py @@ -127,6 +129,8 @@ steps: - vllm/models/muse_spark/ - tests/models/test_muse_spark_logits_local.py - tests/models/test_muse_spark_rocm_shared_rope.py + - vllm/models/common/specialized_decode/ + - tests/models/common/specialized_decode/ commands: - >- pytest -v -s @@ -142,17 +146,18 @@ steps: models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - models/test_dsv41_mono_decode.py + models/test_dsv41_specialized_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py + models/common/specialized_decode - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/hy_v4/test_triton_ihc.py - pytest -v -s models/test_dsv4_mono_decode.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: 45 device: mi355_4 diff --git a/docs/design/specialized_decode.md b/docs/design/specialized_decode.md new file mode 100644 index 000000000000..b69b5a6acaa8 --- /dev/null +++ b/docs/design/specialized_decode.md @@ -0,0 +1,314 @@ +# Specialized Decode Kernels + +## Motivation + +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. + +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 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: + +- 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. + +Key design principles: + +- **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 + +## Quick Overview + +### Declaring a Spec + +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 +# 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, 8): + reasons.append("needs tensor parallel size 2, 4 or 8") + 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, +) +``` + +| 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. + +### Implementing the Kernel + +The kernel stands in for the entire model forward: + +```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: ... +``` + +`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 +class DecodeWarmable(Protocol): + def warmup_keys(self) -> list[Any]: ... + def compile(self, key: Any) -> None: ... +``` + +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 + 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. + +### Wiring It Into the Model + +Three touch points, all on the causal LM: + +```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, so the kernel sees the KV caches once the runner has bound them: + +```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() +``` + +Nothing else changes. In particular, the decoder layer does not. + +!!! 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. + +### Enabling It + +```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 +``` + +## 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) +``` + +## Core Concepts + +### Refuse Versus Skip + +**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. + +**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. + +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. + +### Features + +`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. + +Two checks are enabled by default. + +1. the engine must be on Model Runner V2 +1. the CUDA graph mode is `FULL` + +### Widths and Capture Sizes + +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: + +- **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, # "needs a gfx942 GPU" + 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) +``` + +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 +# 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 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) + + 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) +``` + +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 +# 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() +``` + +## 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 diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index c183b32a8d32..ebabc32210be 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -848,6 +848,53 @@ def test_ir_op_priority_deprecation_warning(caplog_vllm, disable_log_dedup): assert msg in caplog_vllm.text +@pytest.mark.parametrize( + ("argv", "expected"), + [ + ([], None), + (["--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``.""" + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = EngineArgs.from_cli_args(parser.parse_args(argv)) + config = args.create_engine_config() + assert config.kernel_config.specialized_decode == expected + + +def test_specialized_decode_rejects_an_unknown_backend(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + with pytest.raises(SystemExit): + parser.parse_args(["--specialized-decode", "flydsl"]) + + +@pytest.mark.parametrize( + ("flags", "expected"), + [ + ({}, "monokernel"), + ({"monokernel": False}, None), + ({"specialized_decode": "auto"}, "auto"), + ], +) +def test_specialized_decode_overrides_the_kernel_config(flags, expected): + """``--no-monokernel`` turns off a specialized decode that a config file or + the Python API asked for, and absent flags leave it alone.""" + from vllm.config.kernel import KernelConfig + + config = EngineArgs( + kernel_config=KernelConfig(specialized_decode="monokernel"), + **flags, + ).create_engine_config() + assert config.kernel_config.specialized_decode == expected + + @pytest.mark.parametrize( ("input_json", "expected_json"), [ 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 752292c124fc..2ac442b00e63 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 @@ -38,7 +38,7 @@ def _set_gate(monkeypatch, on: bool): # The model records --monokernel when it is created. These tests have no # model, so they set the recorded flag themselves. - monkeypatch.setattr(ops, "_monokernel", on) + monkeypatch.setattr(ops, "_specialized_decode", on) ops.enabled.cache_clear() 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 892987dfd7f3..11c3a7538f22 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 @@ -37,7 +46,7 @@ def _set_gate(monkeypatch, on: bool): # The model records --monokernel when it is created. These tests have no # model, so they set the recorded flag themselves. - monkeypatch.setattr(ops, "_monokernel", on) + monkeypatch.setattr(ops, "_specialized_decode", on) ops.enabled.cache_clear() 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 38733b183e6e..bdb3b3f02182 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py @@ -64,7 +64,7 @@ def gfx942_step(monkeypatch): if paged_mqa_logits_module() is None: pytest.skip("AITER's paged MQA logits module is not installed") monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) - monkeypatch.setattr(ops, "_monokernel", True) + monkeypatch.setattr(ops, "_specialized_decode", True) ops.enabled.cache_clear() if not is_workspace_manager_initialized(): init_workspace_manager(torch.device("cuda")) diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_topk.py b/tests/kernels/models/amd/test_dsv41_gfx942_topk.py index 6964f5ca5d0a..148a2c9a2712 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_topk.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_topk.py @@ -45,7 +45,7 @@ def gfx942_decode(monkeypatch): the recorded flag itself.""" from vllm.models.deepseek_v41.amd import ops - monkeypatch.setattr(ops, "_monokernel", True) + monkeypatch.setattr(ops, "_specialized_decode", True) ops.enabled.cache_clear() yield ops.enabled.cache_clear() 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..d712451fad0a --- /dev/null +++ b/tests/models/common/specialized_decode/kit.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""What a model's ``DecodeSpec`` must satisfy.""" + +import subprocess +import sys +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any + +import pytest + +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import ForwardContext, override_forward_context +from vllm.models.common.specialized_decode.dispatch import SpecializedDecode +from vllm.models.common.specialized_decode.guards import select_decode_spec + +KERNEL_BACKENDS = ("aiter", "flydsl") +"""Backends a spec module must not pull in of its own: it is imported at config +time, before anything has decided the spec can run. The package holding it is +the model's own hardware entry point, which already pulls its platform's +backend, so only what the spec module adds on top of that package counts.""" + + +@contextmanager +def forward(config, mode: CUDAGraphMode, attn_metadata: Any) -> Iterator[None]: + """One forward pass's context, as a runner would set it.""" + context = ForwardContext( + no_compile_layers=config.compilation_config.static_forward_context, + attn_metadata=attn_metadata, + slot_mapping={}, + cudagraph_runtime_mode=mode, + ) + with override_forward_context(context): + yield + + +class DecodeConformance: + """What every model's ``DecodeSpec`` must satisfy. Subclass it in the + model's test module and fill in the hooks below.""" + + specs: tuple[Any, ...] + """The model's variants, as it passes them to ``select_decode_spec``.""" + + specs_module: str + """Where ``specs`` is defined, for the import test to load on its own.""" + + def make_config(self, **overrides) -> Any: + """A config this spec serves, with ``overrides`` applied on top.""" + raise NotImplementedError + + def make_model(self, config) -> Any: + """A loaded model this spec serves, with its KV caches bound into + ``config.compilation_config.static_forward_context``.""" + raise NotImplementedError + + def bad_caches(self, config, model) -> dict[str, Any]: + """Caches that break at least two of ``refuse_caches``'s checks.""" + raise NotImplementedError + + @pytest.fixture + def config(self): + return self.make_config() + + @pytest.fixture + def model(self, config): + return self.make_model(config) + + @pytest.fixture + def decode(self, config, model): + return self.make_decode(config, model) + + def make_decode(self, config, model=None, bind: bool = True): + """This rank's ``SpecializedDecode``, built the way a model builds it: + the variant picked once by ``select_decode_spec``, then the dispatch + over the loaded model, then the caches the runner binds.""" + model = self.make_model(config) if model is None else model + spec = select_decode_spec(config, self.specs) + assert spec is not None + decode = SpecializedDecode(model, config, spec) + if bind: + decode.bind_caches() + return decode + + def test_spec_module_imports_no_kernel_backend(self): + module = self.specs_module + package = module.rpartition(".")[0] + source = ( + "import sys; sys.path.insert(0, '.'); import importlib;" + f"importlib.import_module({package!r});" + "loaded = set(sys.modules);" + f"importlib.import_module({module!r});" + "print('\\n'.join(set(sys.modules) - loaded))" + ) + loaded = subprocess.run( + [sys.executable, "-c", source], + capture_output=True, + text=True, + check=True, + ).stdout.split() + pulled = {n.split(".")[0] for n in loaded} & set(KERNEL_BACKENDS) + assert not pulled, f"{module} imports {sorted(pulled)}" + + def test_refuse_caches_reports_every_reason(self, config, model, decode): + reasons = decode.spec.refuse_caches( + decode.vllm_config, decode.kernel, self.bad_caches(config, model) + ) + assert len(reasons) >= 2, reasons + + def test_build_does_not_move_weights(self, config): + model = self.make_model(config) + before = {n: p.data_ptr() for n, p in model.named_parameters()} + self.make_decode(config, model) + after = {n: p.data_ptr() for n, p in model.named_parameters()} + assert before == after diff --git a/tests/models/common/specialized_decode/reference.py b/tests/models/common/specialized_decode/reference.py new file mode 100644 index 000000000000..c7b85ddbcc23 --- /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.kernel_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="monokernel"), + 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..18a9302febb3 --- /dev/null +++ b/tests/models/common/specialized_decode/test_specialized_decode.py @@ -0,0 +1,333 @@ +# 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.""" + +import importlib +import sys +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 = 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 + + +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] + + +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 + + +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 == 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 == reference.widths(config) + + 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"] +) +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)) + + +def test_the_contract_imports_without_a_vendor_runtime(monkeypatch): + """The library loads where neither FlyDSL nor aiter is installed, so a + 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) + 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/minimax_m3/amd_specialized_decode/test_model.py b/tests/models/minimax_m3/amd_specialized_decode/test_model.py index a84223241854..f34a8d4b57f2 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_model.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_model.py @@ -334,7 +334,7 @@ def _engine(vllm_runner, model_dir, mono, graphs=False, **extra): kernel_config={ "enable_jit_warmup": False, "moe_backend": "aiter", - "monokernel": mono, + "specialized_decode": "monokernel" if mono else None, }, enforce_eager=not graphs, compilation_config={"cudagraph_mode": "FULL_DECODE_ONLY" if graphs else "NONE"}, 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 27d2aa2b70b4..18a9cfff4142 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py @@ -617,6 +617,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 @@ -656,7 +657,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/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py index 57863815b88e..873bc23f4ae7 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py @@ -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", monokernel=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_decode.py b/tests/models/test_dsv41_mono_decode.py deleted file mode 100644 index 9831014f8492..000000000000 --- a/tests/models/test_dsv41_mono_decode.py +++ /dev/null @@ -1,831 +0,0 @@ -# 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 (MAX_ROWS_GFX942 on gfx942) 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.""" - -import sys -from collections import defaultdict -from types import ModuleType, SimpleNamespace - -import pytest -import torch - -from vllm.config import CUDAGraphMode -from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -from vllm.models.deepseek_v41.amd import mono_decode as md - -M = 6 - - -class _Runner: - def __init__(self): - self.calls = [] - self.draft = None - - def forward(self, *args, **kwargs): - self.calls.append((args, {**kwargs, "draft_runner": self.draft})) - return ("mono outputs",) - - def first_pre_mix(self, rows): - return torch.zeros(rows, 4) - - def ffn(self, *args, **kwargs): - self.calls.append((args, kwargs)) - return ("ffn outputs",) - - -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), - compress_ratio=1, - indexer=indexer, - _compressed_kv_cache=cache, - ) - return SimpleNamespace(attn=attn, engram=None, hc_mult=4) - - -def _swa(**overrides): - 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), - 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), - ) - - -@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 _run( - mode=CUDAGraphMode.FULL, - metadata=None, - ffn_only=False, - entry=None, - layer=None, - max_rows=md.MAX_ROWS, - mono=None, - **inputs, - ): - # The row limit depends on the GPU (MAX_ROWS_GFX942 on gfx942), so the - # test sets it and does not ask the GPU of the machine it runs on. - monkeypatch.setattr(md, "_max_rows", lambda: max_rows) - 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) - - def pick(device, draft=None): - runner.draft = draft - return runner - - monkeypatch.setattr(md, "_mono_runner", pick) - monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) - args = {**_inputs(), **inputs} - mono = mono or 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 - - -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"] - - -DRAFT_WIDTH = 192 # the DSpark draft's non-causal SWA rows: window 128 + 5 drafts - - -def _draft(**overrides): - """A DSpark draft layer's mono path, and the metadata its step carries.""" - layer = _layer(comp=None) - layer.attn.compress_ratio = 0 - swa = _swa( - decode_swa_width=DRAFT_WIDTH, - decode_swa_indices=torch.zeros(64, 1, DRAFT_WIDTH, dtype=torch.int32), - ) - mono = md.MonoDecodeLayer(draft=True, experts=128, topk=3, swa_width=DRAFT_WIDTH) - return dict(layer=layer, mono=mono, metadata={"swa": swa}, **overrides) - - -def test_draft_layer_runs_mono_on_its_own_runner(run): - out, calls = run(**_draft()) - assert out == ("mono outputs",) and len(calls) == 1 - assert calls[0][0][9].stride(0) == DRAFT_WIDTH # the non-causal SWA rows - assert calls[0][1]["draft_runner"] is not None # not the target's runner - - -def test_draft_first_layer_takes_the_broadcast_embedding(run): - """The draft hands its first layer the embedding already repeated to the - hc copies: K1 takes it as the settled residual.""" - x = torch.zeros(M, 4, 5120, dtype=torch.bfloat16) - out, calls = run( - **_draft(residual=None, post_mix=None, res_mix=None, pre_mix=None, x=x) - ) - assert out == ("mono outputs",) - args, kwargs = calls[0] - assert args[1].shape == (M, 5120) and args[2] is x and not kwargs["fold"] - - -def test_draft_and_target_layers_take_only_their_own_swa_rows(run): - causal = {"swa": _swa(), "comp": SimpleNamespace(block_size=128)} - out, calls = run(**{**_draft(), "metadata": causal}) - assert out is None and not calls - target = {"swa": _draft()["metadata"]["swa"], "comp": causal["comp"]} - out, calls = run(metadata=target) - assert out is None and not calls - - -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 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 - - -@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 _linear(): - """A dense linear whose weights vLLM processes at load.""" - return SimpleNamespace( - weight=object(), - weight_scale=object(), - quant_method=SimpleNamespace(process_weights_after_loading=lambda layer: None), - ) - - -# The routed experts and picks a token of the target's MoE layers and of the -# DSpark draft's 3 layers (dspark_n_routed_experts, dspark_num_experts_per_tok). -TARGET_ROUTING = (384, 6) -DRAFT_ROUTING = (128, 3) - - -def _decoder_layer(engram=None, routing=TARGET_ROUTING, **attn): - attn = ( - dict( - layer_id=21, - compressor=None, - indexer=None, - compress_ratio=1, - kv_mxfp8=False, - kv_cache_dtype="fp8_ds_mla", - ) - | {name: _linear() for name in md.ATTN_LINEARS_942} - | attn - ) - backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 - # On gfx942, create wraps process_weights_after_loading to copy the experts. - quant_method = SimpleNamespace( - mxfp4_backend=backend, process_weights_after_loading=lambda layer: None - ) - ffn = SimpleNamespace( - experts=SimpleNamespace( - routed_experts=SimpleNamespace(quant_method=quant_method) - ), - shared_experts=SimpleNamespace(gate_up_proj=_linear(), down_proj=_linear()), - gate=SimpleNamespace(tid2eid=None), - n_routed_experts=routing[0], - n_activated_experts=routing[1], - 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, - ) - - -def _indexer(): - from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer - - return object.__new__(DeepseekV41RocmMxfp4Indexer) - - -@pytest.mark.parametrize( - "layer, path", - [ - (_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), - ], - ids=[ - "standard", - "window", - "engram", - "kv-source", - "index-source", - "compressor-only", - "other-indexer", - "kv", - "draft", - ], -) -def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): - """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 - - -@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 - assert calls[0][1] == {"topk": TARGET_ROUTING[1]} - - -def test_ffn_layer_never_runs_whole(run): - out, calls = run(ffn_only=True, entry="layer") - assert out is None and not calls - - -@pytest.mark.parametrize( - "case", - [ - 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)), - ], - ids=["piecewise", "prefill", "tokens-past-rows", "no-metadata", "rows"], -) -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 _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, - dspark_n_routed_experts=DRAFT_ROUTING[0], - dspark_num_experts_per_tok=DRAFT_ROUTING[1], - ) - ), - parallel_config=SimpleNamespace(**parallel), - kernel_config=SimpleNamespace(monokernel=monokernel), - ) - - -def _deployment(monkeypatch, cdna=4, cus=256, tp=2, target=None): - """A deployment the kernels take (the runner module stubbed: no FlyDSL). - ``target`` is the architecture FlyDSL compiles for, by default the GPU's.""" - import vllm.platforms.rocm as rocm - - 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: tp) - monkeypatch.setattr(md, "_compute_units", lambda: cus) - gpu = "gfx942" if cdna == 3 else "gfx950" - monkeypatch.setattr(md, "_gpu_arch", lambda: gpu) - arch = ModuleType("vllm.models.deepseek_v41.amd.mono.common.arch") - arch.target_arch = lambda: target or gpu # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, arch.__name__, arch) - # On gfx942, create builds the top-k extension with hipcc. These tests do - # not need it, so they skip the build. - from vllm.models.deepseek_v41.amd import ops - from vllm.models.deepseek_v41.amd.ops import topk - - monkeypatch.setattr(topk, "build", lambda: True) - # create records --monokernel for the gfx942 helper ops. The test puts the - # recorded value back afterwards, so that no later test sees this one's. - monkeypatch.setattr(ops, "_monokernel", ops._monokernel) - 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_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 - ) - - -@pytest.mark.parametrize("on", [True, False]) -def test_create_records_the_flag_for_the_gfx942_helpers(monkeypatch, on): - """The gfx942 decode helpers (amd/ops) run in forward passes, where vLLM's - current config is not set, so create records --monokernel for them.""" - from vllm.models.deepseek_v41.amd import ops - - _deployment(monkeypatch, cdna=3, tp=4) - md.MonoDecodeLayer.create(_decoder_layer(), _config(monokernel=on)) - assert ops._monokernel is on - - -@pytest.mark.parametrize("cdna", [2, 5]) -def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): - _deployment(monkeypatch, cdna=cdna) - with pytest.raises(ValueError, match="needs CDNA3/4"): - md.MonoDecodeLayer.create(_decoder_layer(), _config()) - - -@pytest.mark.parametrize( - "cdna, tp, target", [(3, 4, "gfx950"), (4, 2, "gfx942")], ids=["gfx942", "gfx950"] -) -def test_opt_in_with_another_arch_for_flydsl_raises(monkeypatch, cdna, tp, target): - """An ARCH environment variable that names another architecture than the - GPU's would make FlyDSL build kernels that do not match the weight copies, - so the opt-in stops the server.""" - _deployment(monkeypatch, cdna=cdna, tp=tp, target=target) - with pytest.raises(ValueError, match=f"makes it compile for {target}"): - md.MonoDecodeLayer.create(_decoder_layer(), _config()) - - -def _other_backend(): - layer = _decoder_layer() - experts = layer.ffn.experts.routed_experts - experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON - return layer - - -@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", {}))) - - -@pytest.mark.parametrize("rows, runs", [(md.MAX_ROWS_GFX942, True), (48, False)]) -def test_gfx942_decode_steps_of_up_to_42_rows(run, rows, runs): - """On gfx942 the kernels take steps of up to MAX_ROWS_GFX942 rows. A step of - 48 rows, which gfx950 takes, runs vLLM's path there.""" - swa = _swa(num_decode_tokens=rows) - comp = SimpleNamespace(block_size=128, block_table=None) - out, calls = run( - metadata={"swa": swa, "comp": comp}, - max_rows=md.MAX_ROWS_GFX942, - **_inputs(rows), - ) - assert (out == ("mono outputs",)) == runs and len(calls) == int(runs) - - -def _kind(mono): - if mono is None: - return None - return f"ffn top {mono.topk}" if mono.ffn_only else "whole" - - -def _indexer_942(): - from vllm.models.deepseek_v41.attention import DeepseekV4Indexer - - return object.__new__(DeepseekV4Indexer) - - -@pytest.mark.parametrize( - "layer, kind", - [ - (_decoder_layer(), "whole"), - (_decoder_layer(compress_ratio=0), "whole"), - (_decoder_layer(compress_ratio=0, engram=object()), "whole"), - (_decoder_layer(engram=object()), "whole"), - (_decoder_layer(compressor=object(), indexer=_indexer_942()), "whole"), - (_decoder_layer(indexer=_indexer_942()), "whole"), - (_decoder_layer(compressor=object()), "ffn top 6"), - (_decoder_layer(indexer=object()), "ffn top 6"), - (_decoder_layer(kv_cache_dtype="auto"), "ffn top 6"), - (_decoder_layer(layer_id=40, routing=DRAFT_ROUTING), "ffn top 3"), - (_decoder_layer(layer_id=40), None), - ], - ids=[ - "standard", - "first", - "window-engram", - "engram", - "kv-source", - "index-source", - "compressor-only", - "other-indexer", - "kv", - "draft", - "draft-target-routing", - ], -) -def test_gfx942_layers(monkeypatch, layer, kind): - """On gfx942 every backbone layer runs whole, as on gfx950. Its index - layers have vLLM's FP8 indexer, because the MXFP4 indexer cache needs - CDNA4. The DSpark draft layers run the FFN launch with the draft's routing - of 128 experts and top 3.""" - _deployment(monkeypatch, cdna=3, tp=4) - assert _kind(md.MonoDecodeLayer.create(layer, _config())) == kind - - -def test_gfx942_tp2_runs_vllm_path(monkeypatch): - """The gfx942 kernels are built for TP4 only. At TP2 every layer runs - vLLM's path, and that is not an error.""" - _deployment(monkeypatch, cdna=3, tp=2) - assert md.MonoDecodeLayer.create(_decoder_layer(), _config()) is None - - -def test_gfx942_takes_any_moe_backend(monkeypatch): - """The gfx942 kernels read their own copies of the experts, made before - vLLM's MoE backend converts them, so a backend other than AITER_MXFP4_BF16 - is not an error there.""" - _deployment(monkeypatch, cdna=3, tp=4) - assert _kind(md.MonoDecodeLayer.create(_other_backend(), _config())) == "whole" - - -@pytest.mark.parametrize( - "make, copied", - [ - (_decoder_layer, ("gate_up_proj", "down_proj", *md.ATTN_LINEARS_942)), - ( - lambda: _decoder_layer(layer_id=40, routing=DRAFT_ROUTING), - ("gate_up_proj", "down_proj"), - ), - ], - ids=["standard", "draft"], -) -def test_gfx942_copies_dense_linears_at_load(monkeypatch, make, copied): - """On gfx942 the kernels' copies of the dense linears are made when vLLM - processes each linear's weights at load, so that vLLM counts them as model - memory when it sizes the KV cache. A draft layer runs vLLM's attention, so - its attention linears get no copies.""" - _deployment(monkeypatch, cdna=3, tp=4) - weights942 = ModuleType("vllm.models.deepseek_v41.amd.mono.weights942") - weights942.linear_copy = lambda w, s: ("copy", w, s) # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, weights942.__name__, weights942) - layer = make() - assert md.MonoDecodeLayer.create(layer, _config()) is not None - sh = layer.ffn.shared_experts - linears = {"gate_up_proj": sh.gate_up_proj, "down_proj": sh.down_proj} - linears |= {name: getattr(layer.attn, name) for name in md.ATTN_LINEARS_942} - for linear in linears.values(): - linear.quant_method.process_weights_after_loading(linear) - made = {name for name, linear in linears.items() if hasattr(linear, "mono942")} - assert made == set(copied) - for name in copied: - linear = linears[name] - assert linear.mono942 == ("copy", linear.weight, linear.weight_scale) - - -def test_gfx942_reload_writes_the_copies_in_place(monkeypatch): - """A reload of the weights (reload_weights, or a weight update in RL) runs - process_weights_after_loading again. The kernels and captured CUDA graphs - read the first copies by their addresses, so the copies of the reload are - written into the first copies instead of replacing them.""" - _deployment(monkeypatch, cdna=3, tp=4) - weights942 = ModuleType("vllm.models.deepseek_v41.amd.mono.weights942") - weights942.linear_copy = lambda w, s: (w.clone(), s.clone()) # type: ignore[attr-defined] - weights942.moe_copies = lambda w13, w13_s, w2, w2_s, inter: tuple( # type: ignore[attr-defined] - t.clone() for t in (w13, w13_s, w2, w2_s) - ) - monkeypatch.setitem(sys.modules, weights942.__name__, weights942) - layer = _decoder_layer() - assert md.MonoDecodeLayer.create(layer, _config()) is not None - experts, linear = layer.ffn.experts.routed_experts, layer.attn.wo_a - - def load(value): - u8 = dict(dtype=torch.uint8) - for name in ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"): - setattr(experts, name, torch.full((2, 4), value, **u8)) - linear.weight = torch.full((4, 8), value, **u8) - linear.weight_scale = torch.full((4, 1), value, **u8) - for module in (experts, linear): - module.quant_method.process_weights_after_loading(module) - - load(1) - first = (experts.mono942, linear.mono942) - load(2) - assert experts.mono942 is first[0] and linear.mono942 is first[1] - assert all(bool((t == 2).all()) for t in (*first[0], *first[1])) - - -# With --monokernel on gfx942 the indexer layer and the DSpark draft sampling -# run gfx942 launches from the model folder. vLLM's ROCm indexer op and its -# DSpark speculator loop stay as they are, and they run every step that the -# gfx942 code does not take. - - -@pytest.mark.parametrize( - "case, gfx942_path", - [ - ("decode", True), - ("profiling", False), - ("prefill", False), - ("mixed", False), - ("no rows", False), - ("writes its own K", False), - ("without the flag", False), - ], -) -def test_gfx942_indexer_runs_vllm_op_off_decode_steps(monkeypatch, case, gfx942_path): - """Only a step of only decode rows, whose K cache the model already wrote, - takes the gfx942 decode path. The profiling run, prefill and mixed steps, - a layer that writes its own K cache and a process without --monokernel on - gfx942 call vLLM's op with the same arguments.""" - from vllm.models.deepseek_v41.amd.ops import sparse_indexer as si - - prefills, decodes = {"prefill": (1, 0), "mixed": (1, 1), "no rows": (0, 0)}.get( - case, (0, 1) - ) - meta = SimpleNamespace(num_prefills=prefills, num_decodes=decodes) - attn_metadata = None if case == "profiling" else {"idx": meta} - monkeypatch.setattr(si, "DeepseekV32IndexerMetadata", SimpleNamespace) - monkeypatch.setattr( - si, "get_forward_context", lambda: SimpleNamespace(attn_metadata=attn_metadata) - ) - monkeypatch.setattr(si, "enabled", lambda: case != "without the flag") - calls = [] - - def vllm_op(*args, **kwargs): - calls.append(("vllm", args, kwargs)) - return "vllm" - - def decode_step(*args): - calls.append(("gfx942", args, {})) - return "gfx942" - - monkeypatch.setattr(si, "rocm_aiter_sparse_attn_indexer", vllm_op) - monkeypatch.setattr(si, "_decode_step", decode_step) - buffer = torch.zeros(6, 512, dtype=torch.int32) - args = (torch.zeros(6, 4), "idx", torch.zeros(1), torch.zeros(1), None) - args += (torch.zeros(1), 128, "ue8m0", 512, 128, 1024, 1024, buffer) - kwargs = dict( - skip_k_cache_insert=case != "writes its own K", - compress_ratio=4, - candidate_blocks=None, - candidate_block_size=0, - candidate_write=False, - ) - out = si.gfx942_sparse_attn_indexer(*args, **kwargs) - assert out == ("gfx942" if gfx942_path else "vllm") - assert len(calls) == 1 - if not gfx942_path: - assert calls[0][1] == args and calls[0][2] == kwargs - - -@pytest.mark.parametrize( - "monokernel, gfx942, takes_gfx942_indexer", - [(True, True, True), (True, False, False), (False, True, False)], -) -def test_rocm_attention_takes_the_gfx942_indexer_with_the_flag( - monkeypatch, monokernel, gfx942, takes_gfx942_indexer -): - """The attention builds its indexer before the mono decode layer records - the flag for ops.enabled(), so it reads --monokernel from the config of - the model that is being built. Without the flag, and on gfx950, the layer - keeps vLLM's indexer.""" - import vllm.platforms.rocm as rocm_platform - from vllm.models.deepseek_v41.amd import rocm - from vllm.models.deepseek_v41.amd.ops.sparse_indexer import ( - Gfx942SparseAttnIndexer, - ) - from vllm.models.deepseek_v41.attention import DeepseekV4Indexer - - config = SimpleNamespace(kernel_config=SimpleNamespace(monokernel=monokernel)) - monkeypatch.setattr(rocm, "get_current_vllm_config", lambda: config) - monkeypatch.setattr(rocm_platform, "on_gfx942", lambda: gfx942) - attn_cls = rocm.DeepseekV41ROCMAiterMLAAttention - indexer_cls = attn_cls._indexer_cls( - attn_cls.__new__(attn_cls), SimpleNamespace(rocm_mxfp4=False) - ) - expected = ( - rocm.DeepseekV41Gfx942Indexer if takes_gfx942_indexer else DeepseekV4Indexer - ) - assert indexer_cls is expected - assert rocm.DeepseekV41Gfx942Indexer.attn_cls is Gfx942SparseAttnIndexer - - -@pytest.mark.parametrize("monokernel", [True, False]) -@pytest.mark.parametrize("cdna", [3, 4]) -def test_swa_decode_meta_runs_on_cdna3_and_cdna4_with_the_flag( - monkeypatch, cdna, monokernel -): - """swa_decode_meta's kernel is plain Triton, so with --monokernel a decode - batch gets its SWA metadata from the one launch on CDNA4 (gfx950) as on - CDNA3 (gfx942). Without the flag, the builder keeps vLLM's build.""" - import vllm.platforms.rocm as rocm_platform - from vllm.models.deepseek_v41.amd import ops, rocm - from vllm.models.deepseek_v41.amd.ops import swa_decode_meta - - monkeypatch.setattr(ops, "current_platform", SimpleNamespace(is_rocm=lambda: True)) - monkeypatch.setattr(rocm_platform, "get_cdna_version", lambda: cdna) - monkeypatch.setattr(ops, "_monokernel", ops._monokernel) - ops.record_monokernel(_config(monokernel=monokernel)) - launches = [] - monkeypatch.setattr( - swa_decode_meta, "launch", lambda *args, noncausal: launches.append(noncausal) - ) - ints = lambda *shape: torch.zeros(shape, dtype=torch.int32) # noqa: E731 - builder = SimpleNamespace( - decode_threshold=M, - window_size=md.SWA_WIDTH, - block_size=128, - decode_swa_indices=ints(64, 1, md.SWA_WIDTH), - decode_swa_lens=ints(64), - is_valid_token=torch.zeros(64, dtype=torch.bool), - token_to_req_indices=ints(64), - no_replay_start=ints(2), - decode_swa_ragged_indices_buffer=ints(64 * md.SWA_WIDTH), - decode_swa_ragged_indptr_buffer=ints(65), - build_tile_scheduler=lambda rows: defaultdict(lambda: None), - _build_deepseek_v4_metadata=lambda *args: {}, - ) - query_start_loc = torch.tensor([0, M], dtype=torch.int32) - cam = SimpleNamespace( - causal=True, - max_query_len=M, - num_reqs=1, - num_actual_tokens=M, - query_start_loc=query_start_loc, - query_start_loc_cpu=query_start_loc, - seq_lens=torch.tensor([1000], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([1000], dtype=torch.int32), - slot_mapping=torch.arange(M), - block_table_tensor=ints(1, 8), - token_to_req_indices=lambda buffer: buffer[:M], - ) - out = swa_decode_meta.build( - builder, 0, cam, False, None, rocm.DeepseekV41ROCMAiterSparseSWAMetadata - ) - # No later test may see this test's cached gate. - ops.record_monokernel(_config(monokernel=False)) - assert launches == ([False] if monokernel else []) - assert (out is not None) is monokernel - - -def test_dspark_speculator_asks_the_draft_model_first(): - """A draft model that samples a step itself returns True, and vLLM's - loop then does not run. The stub speculator has none of the loop's - buffers, so the loop would fail here.""" - from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator - - calls = [] - - def sample_sequential(s, n, h): - calls.append((s, n, h)) - return True - - spec = DSparkSpeculator.__new__(DSparkSpeculator) - spec._draft_topk = None - spec.model = SimpleNamespace(sample_sequential=sample_sequential) - hidden = torch.zeros(4, 8) - DSparkSpeculator._sample_sequential(spec, 2, hidden) - assert calls == [(spec, 2, hidden)] - - -@pytest.mark.parametrize("taken", [True, False]) -def test_dspark_draft_model_samples_with_fused_markov(monkeypatch, taken): - """The DSv4.1 draft model answers the speculator with fused_markov, which - returns False for a step that the fused launches do not take.""" - from vllm.models.deepseek_v41.amd import dspark - from vllm.models.deepseek_v41.amd.ops import fused_markov - - calls = [] - - def sample_sequential(s, n, h): - calls.append((s, n, h)) - return taken - - monkeypatch.setattr(fused_markov, "sample_sequential", sample_sequential) - spec, hidden = object(), torch.zeros(4, 8) - model_cls = dspark.DSparkDeepseekV4ForCausalLM - assert model_cls.sample_sequential(None, spec, 2, hidden) is taken - assert calls == [(spec, 2, hidden)] diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index 1e428d83d302..3150e9391560 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -1,29 +1,9 @@ # 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, and at TP4 on gfx942, 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. - -On gfx942 the indexer is vLLM's FP8 indexer, because the MXFP4 indexer cache -needs CDNA4, and the kernels read the copies of the weights that the model -makes when it loads them. A decode step of 8 rows there takes two GEMV passes -of 6 rows, the second one padded, and a step of 42 rows, the most that the -gfx942 kernels take, takes seven. - -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 specialized decode against vLLM's path, through the real +model: at TP2 and TP4 on CDNA4, and at TP4 on gfx942, where the indexer is +vLLM's FP8 one (the MXFP4 indexer cache needs CDNA4) and the kernels read the +copies of the weights that the model makes when it loads them.""" import zlib @@ -169,19 +149,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 @@ -220,37 +203,61 @@ 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}") def mono_layers(self) -> dict: - from vllm.models.deepseek_v41.amd import mono_decode + from vllm.models.deepseek_v41.amd.specialized_decode.runtime.weights942 import ( + ATTN_LINEARS, + ) - def copies(m) -> bool: + model = self.model_runner.model # type: ignore[attr-defined] + decode = getattr(model, "specialized_decode", None) + plans = {} if decode is None else decode.kernel.plans + + def copies(plan) -> bool: """Whether the layer holds the gfx942 copies of its weights, which the model makes when it loads them (``mono942``).""" - ffn = m.ffn + ffn = plan.layer.ffn held = [ ffn.experts.routed_experts, ffn.shared_experts.gate_up_proj, ffn.shared_experts.down_proj, ] - if not m.mono.ffn_only: - held += [getattr(m.attn, n) for n in mono_decode.ATTN_LINEARS_942] + if not plan.ffn_only: + held += [getattr(plan.layer.attn, n) for n in ATTN_LINEARS] return all(hasattr(t, "mono942") for t in held) layers = { - i: None - if m.mono is None - else [ - "ffn" if m.mono.ffn_only else "whole", - m.mono._weights is not None, - copies(m), - ] - for i, m in self._decoder_layers().items() + i: ( + [ + "ffn" if plans[i].ffn_only else "whole", + plans[i].weights is not None, + copies(plans[i]), + ] + if i in plans + else None + ) + for i in self._decoder_layers() } - runner = mono_decode._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]), @@ -287,7 +294,7 @@ def _decode( hf_overrides=TRUNCATED, worker_extension_cls=f"{__name__}.MonoProbe", moe_backend="aiter", - monokernel=mono, + specialized_decode="monokernel" if mono else None, language_model_only=True, tokenizer_mode="deepseek_v41", **indexer, diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py index 5a620141bb8c..686770e42132 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) diff --git a/tests/models/test_dsv41_specialized_decode.py b/tests/models/test_dsv41_specialized_decode.py new file mode 100644 index 000000000000..bc94a4235025 --- /dev/null +++ b/tests/models/test_dsv41_specialized_decode.py @@ -0,0 +1,890 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1's DecodeSpec against the conformance kit, plus its own choices.""" + +import importlib +import sys +from collections import defaultdict +from types import ModuleType, SimpleNamespace +from typing import Any + +import pytest +import torch + +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.spec import DecodeUnsupported, Step +from vllm.models.deepseek_v41.amd.specialized_decode import spec as sd + +SPECS = (sd.DECODE_SPEC,) +LAYERS = 6 +HIDDEN = 16 +BLOCKS = 256 +BLOCK = 32 # KV records a cache block holds + + +class _StubRunner: + 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) + + 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, **kwargs): + return ("ffn", args, kwargs) + + +_MONO = "vllm.models.deepseek_v41.amd.mono" + + +def _stub_modules(monkeypatch) -> None: + """The FlyDSL runner and arch modules, as spec-shaped stubs.""" + 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(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) + + +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 _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, + kv_mxfp8=False, + kv_cache_dtype="fp8_ds_mla", + 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(), + 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(**(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)) + self.taken: list = [] + + 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() + + +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): + 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.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, replay=True + ): + return None + + def epilogue(self, hidden, *rest, replay=True): + return hidden + + +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( + 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, + 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=list(range(1, sd.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 _kernel_config(specialized_decode: str | None) -> Any: + return make_config( + kernel_config=SimpleNamespace(specialized_decode=specialized_decode) + ) + + +def _swa(**overrides) -> Any: + fields = dict( + num_prefills=0, + num_decodes=1, + num_decode_tokens=1, + 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), + block_size=BLOCK, + ) + return SimpleNamespace(**(fields | overrides)) + + +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=BLOCK, block_table=torch.zeros(8, 4, dtype=torch.int32) + ) + return meta + + +class TestDsv41Spec(DecodeConformance): + specs = SPECS + specs_module = "vllm.models.deepseek_v41.amd.specialized_decode.spec" + + @pytest.fixture(autouse=True) + def _stubs(self, monkeypatch): + _gpu(monkeypatch) + + def make_config(self, **overrides): + return make_config(**overrides) + + 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 bad_caches(self, config, model): + return {"swa": torch.zeros(4, dtype=torch.bfloat16)} + + +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" + ) + try: + importlib.import_module("vllm.platforms.rocm") + finally: + torch.cuda.get_device_properties = probe + return sys.modules["vllm.platforms.rocm"] + + +@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, 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): + _gpu(monkeypatch) + kit = TestDsv41Spec() + config = config or make_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 kit.make_decode(config, model), model + + +def _plan(decode, layer_id: int): + return decode.kernel.plans[layer_id] + + +def _indexer(): + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + return object.__new__(DeepseekV41RocmMxfp4Indexer) + + +@pytest.mark.parametrize( + "overrides, ffn_only", + [ + (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", + "kv-source", + "index-source", + "compressor-only", + "other-indexer", + "kv", + "mxfp8", + "ratio", + ], +) +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, from + wo_b's unreduced output.""" + decode, model = _build(monkeypatch, **overrides()) + 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) + + +@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) + decode, _ = _build(monkeypatch, config=config, model=model) + assert not decode.kernel.plans + + +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 + decode, _ = _build(monkeypatch, config=config, model=model) + assert not decode.kernel.plans + + +def _step(decode, metadata, width=1): + with forward(decode.vllm_config, CUDAGraphMode.FULL, metadata): + return sd._skip_step(decode.vllm_config, decode.kernel, width) + + +@pytest.mark.parametrize( + "case, why", + [ + (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=["prefill", "no-decode", "tokens-past-rows", "noncausal-window", "no-swa"], +) +def test_other_steps_take_vllm_path(monkeypatch, case, why): + decode, model = _build(monkeypatch) + assert why in _step(decode, _metadata(model, **case)) + + +def test_decode_step_takes_the_kernels(monkeypatch): + decode, model = _build(monkeypatch) + assert _step(decode, _metadata(model)) is None + + +def test_profile_run_takes_vllm_path(monkeypatch): + decode, model = _build(monkeypatch) + assert _step(decode, None) == "no attention metadata" + + +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) + + +@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, 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 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 sd._refuse_config(make_config())) + + +@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="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)) + + +def test_too_few_compute_units_refuses(monkeypatch): + config = make_config() + model = TestDsv41Spec().make_model(config) + _gpu(monkeypatch, cus=BLOCKS // 2) + with pytest.raises(DecodeUnsupported, match="compute units"): + SpecializedDecode(model, config, select_decode_spec(config, SPECS)) + + +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(DecodeUnsupported, match="AITER_MXFP4_BF16"): + 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() + config.kernel_config.specialized_decode = None + assert select_decode_spec(config, SPECS) is None + + +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)): + out = decode.kernel.launch_whole( + _plan(decode, 1), + 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:] == (sd.RECORD,)] + assert len(records) == 1 + 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).""" + decode, model = _build(monkeypatch) + positions = torch.zeros(1, dtype=torch.long) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + 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) + + +def test_ffn_launch_takes_wo_b_unreduced(monkeypatch): + decode, model = _build(monkeypatch, compressor=object()) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + 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), + torch.zeros(1, 4, 4), + torch.zeros(1, 4), + ) + 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 + decode, _ = _build(monkeypatch, config=config, model=model) + + 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, 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) + 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)], + ] + + +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.""" + decode, model = _build(monkeypatch) + positions = torch.zeros(1, dtype=torch.long) + with forward(decode.vllm_config, CUDAGraphMode.FULL, _metadata(model)): + out = decode.kernel.launch_whole( + _plan(decode, 1), + torch.zeros(1, 4, HIDDEN), + positions, + None, + None, + None, + None, + ) + assert out is None + + +@pytest.mark.parametrize("specialized_decode", ["monokernel", None]) +@pytest.mark.parametrize("cdna", [3, 4]) +def test_swa_decode_meta_runs_on_cdna3_and_cdna4_with_the_flag( + monkeypatch, cdna, specialized_decode +): + """swa_decode_meta's kernel is plain Triton, so with --monokernel a decode + batch gets its SWA metadata from the one launch on CDNA4 (gfx950) as on + CDNA3 (gfx942). Without the flag, the builder keeps vLLM's build.""" + import vllm.platforms.rocm as rocm_platform + from vllm.models.deepseek_v41.amd import ops, rocm + from vllm.models.deepseek_v41.amd.ops import swa_decode_meta + + rows = 6 + monkeypatch.setattr(ops, "current_platform", SimpleNamespace(is_rocm=lambda: True)) + monkeypatch.setattr(rocm_platform, "get_cdna_version", lambda: cdna) + monkeypatch.setattr(ops, "_specialized_decode", ops._specialized_decode) + ops.record_specialized_decode(_kernel_config(specialized_decode)) + launches = [] + monkeypatch.setattr( + swa_decode_meta, "launch", lambda *args, noncausal: launches.append(noncausal) + ) + ints = lambda *shape: torch.zeros(shape, dtype=torch.int32) # noqa: E731 + builder = SimpleNamespace( + decode_threshold=rows, + window_size=sd.SWA_WIDTH, + block_size=128, + decode_swa_indices=ints(64, 1, sd.SWA_WIDTH), + decode_swa_lens=ints(64), + is_valid_token=torch.zeros(64, dtype=torch.bool), + token_to_req_indices=ints(64), + no_replay_start=ints(2), + decode_swa_ragged_indices_buffer=ints(64 * sd.SWA_WIDTH), + decode_swa_ragged_indptr_buffer=ints(65), + build_tile_scheduler=lambda rows: defaultdict(lambda: None), + _build_deepseek_v4_metadata=lambda *args: {}, + ) + query_start_loc = torch.tensor([0, rows], dtype=torch.int32) + cam = SimpleNamespace( + causal=True, + max_query_len=rows, + num_reqs=1, + num_actual_tokens=rows, + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=torch.tensor([1000], dtype=torch.int32), + seq_lens_cpu_upper_bound=torch.tensor([1000], dtype=torch.int32), + slot_mapping=torch.arange(rows), + block_table_tensor=ints(1, 8), + token_to_req_indices=lambda buffer: buffer[:rows], + ) + out = swa_decode_meta.build( + builder, 0, cam, False, None, rocm.DeepseekV41ROCMAiterSparseSWAMetadata + ) + # No later test may see this test's cached gate. + ops.record_specialized_decode(_kernel_config(None)) + assert launches == ([False] if specialized_decode else []) + assert (out is not None) is bool(specialized_decode) diff --git a/tests/models/test_dsv4_mono_decode.py b/tests/models/test_dsv4_mono_decode.py index ee17ac9852f6..fc8bad604b9c 100644 --- a/tests/models/test_dsv4_mono_decode.py +++ b/tests/models/test_dsv4_mono_decode.py @@ -297,7 +297,9 @@ def _config(monokernel=True, max_model_len=8192): max_model_len=max_model_len, ), parallel_config=SimpleNamespace(**parallel), - kernel_config=SimpleNamespace(monokernel=monokernel), + kernel_config=SimpleNamespace( + specialized_decode="monokernel" if monokernel else 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 fd6333f267e3..54c7a789b61c 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -162,7 +162,7 @@ def m3_config(): parallel_config=ParallelConfig(), quant_config=None, cache_config=None, - kernel_config=SimpleNamespace(monokernel=False, moe_backend="auto"), + kernel_config=SimpleNamespace(specialized_decode=None, moe_backend="auto"), ) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 9e40c1528f59..ff743e1293f8 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"] + RDNAHybridW4A16PrefillDequantMode = Literal["enable", "disable"] @@ -299,11 +301,13 @@ 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, DeepSeek-V4-Pro at TP8 on - MI355X). 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 (DeepSeek-V4.1 on MI355X, + DeepSeek-V4-Pro at TP8 on MI355X); `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 316881e5d61b..9f85abe5968b 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 @@ -776,7 +777,8 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) - monokernel: bool = get_field(KernelConfig, "monokernel") + specialized_decode: SpecializedDecodeBackend | None = None + monokernel: bool | None = None worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1844,7 +1846,16 @@ 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="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) @@ -2801,8 +2812,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: + 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": 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 new file mode 100644 index 000000000000..0601c6298be1 --- /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: 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, tp_agree +from vllm.models.common.specialized_decode.spec import DecodeSpec, Step + +__all__ = [ + "DecodeSpec", + "SpecializedDecode", + "Step", + "select_decode_spec", + "tp_agree", +] diff --git a/vllm/models/common/specialized_decode/dispatch.py b/vllm/models/common/specialized_decode/dispatch.py new file mode 100644 index 000000000000..898cc5e7d831 --- /dev/null +++ b/vllm/models/common/specialized_decode/dispatch.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""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 +from typing import TYPE_CHECKING, Any, Generic + +import torch + +from vllm.config.compilation import CUDAGraphMode +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.specialized_decode.guards import ( + _refuse_model, + tp_agree, + tp_group, +) +from vllm.models.common.specialized_decode.spec import ( + DecodeSpec, + DecodeUnsupported, + DecodeWarmable, + KernelT, + Step, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +class DecodeWarmup(VllmJitKernel["DecodeWarmup.CompileKey"]): + """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, kernel: DecodeWarmable): + super().__init__() + self.spec = spec + self.kernel = kernel + + 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: "DecodeWarmup.CompileKey") -> None: + self.kernel.compile(compile_key.build) + self._compiled_cache[compile_key] = True + + +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 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. + """ + + def __init__(self, model, vllm_config: "VllmConfig", spec: DecodeSpec[KernelT]): + """Builds this rank's kernel. + + Args: + model: The loaded model, in ``process_weights_after_loading``. + vllm_config: The engine's config. + spec: The spec to build and run. + + """ + self.spec = spec + reasons = _refuse_model(spec, model) + agreed = tp_agree(not reasons, tp_group()[2]) + refused = f"Specialized decode {spec.name} refused for the loaded model" + if reasons: + raise DecodeUnsupported(f"{refused}: " + "; ".join(reasons)) + if not agreed: + raise DecodeUnsupported(f"{refused}: another TP rank refused") + self.vllm_config = 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 + 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: + """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} + reasons = self.spec.refuse_caches(self.vllm_config, self.kernel, caches) + if reasons: + raise DecodeUnsupported( + f"Specialized decode {self.spec.name} caches: " + "; ".join(reasons) + ) + 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. + + 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) + kernel_width = self._kernel_width(width) + why: str | None + if intermediate_tensors is not None: + why = "pipeline stage input" + elif kernel_width is None: + 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.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 + 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 + 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 new file mode 100644 index 000000000000..4d30aed75df4 --- /dev/null +++ b/vllm/models/common/specialized_decode/guards.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Common validation logic for specialized decode kernels.""" + +from collections.abc import 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.specialized_decode.spec import ( + DecodeSpec, + DecodeUnsupported, + Feature, + KernelT, +) + +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 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: + """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. + + """ + 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 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]: + """Features in use by this config.""" + 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.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, + 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: DecodeSpec[Any], vllm_config: "VllmConfig") -> list[str]: + """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) + ] + 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: DecodeSpec[Any], model) -> list[str]: + """The reasons the mode cannot run the specialized decode kernel, if any.""" + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + + reasons = [] + if isinstance(model, TorchCompileWithNoGuardsWrapper): + reasons.append( + "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[DecodeSpec[Any]], vllm_config: "VllmConfig" +) -> tuple[DecodeSpec[Any] | None, dict[str, list[str]]]: + """Get the first compatible spec.""" + 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 _set_capture_sizes(spec: DecodeSpec[Any], vllm_config: "VllmConfig") -> None: + """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 [] + new = sorted(set(widths) | {s for s in sizes if s > max(widths)}) + if new != sorted(sizes): + logger.info( + "Specialized decode %s: CUDA graph capture sizes %s", spec.name, new + ) + cc.cudagraph_capture_sizes = new + + +def select_decode_spec( + vllm_config: "VllmConfig", + specs: DecodeSpec[KernelT] | Sequence[DecodeSpec[KernelT]], +) -> DecodeSpec[KernelT] | None: + """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. + specs: The model's variants, in the order they are preferred. + + """ + if not _enabled(vllm_config): + return None + ordered = (specs,) if isinstance(specs, DecodeSpec) else tuple(specs) + spec, refusals = _select_spec(ordered, vllm_config) + if spec is None: + raise DecodeUnsupported( + "Specialized decode refused: " + + "; ".join(f"{n}: {', '.join(rs)}" for n, rs in refusals.items()) + ) + 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 " + "with FULL CUDA graphs", + spec.name, + ) + return spec diff --git a/vllm/models/common/specialized_decode/spec.py b/vllm/models/common/specialized_decode/spec.py new file mode 100644 index 000000000000..4359cc9fa147 --- /dev/null +++ b/vllm/models/common/specialized_decode/spec.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Interfaces for when and how to run a specialized decode kernel.""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from enum import Enum +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Protocol, + TypeVar, + runtime_checkable, +) + +if TYPE_CHECKING: + import torch + + from vllm.config import VllmConfig + + +class Feature(str, Enum): + """Features a specialized decode kernel may require.""" + + PIPELINE_PARALLEL = "pipeline_parallel" + DATA_PARALLEL = "data_parallel" + 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" + 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 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 ``SpecializedDecode.plan`` decided it.""" + + width: int + """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.""" + + +class DecodeImpl(Protocol): + """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. + """ + + cache_layers: tuple[str, ...] + """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.""" + ... + + 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.""" + ... + + def compile(self, key: Any) -> None: + """Compiles one of ``warmup_keys`` without launching it.""" + ... + + +KernelT = TypeVar("KernelT", bound=DecodeImpl) + + +@dataclass(frozen=True) +class DecodeSpec(Generic[KernelT]): + """Definition of when specialized decode can run, and how to build its + implementation. + + 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. + 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, + 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. 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. + + """ + + 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 diff --git a/vllm/models/deepseek_v4/amd/mono_decode.py b/vllm/models/deepseek_v4/amd/mono_decode.py index b9eb6d11fd5c..4260e8c62319 100644 --- a/vllm/models/deepseek_v4/amd/mono_decode.py +++ b/vllm/models/deepseek_v4/amd/mono_decode.py @@ -212,7 +212,7 @@ def __init__(self, hca_width: int) -> None: def create( cls, layer: "DeepseekV4DecoderLayer", vllm_config ) -> "MonoDecodeLayer | None": - if not vllm_config.kernel_config.monokernel: + if not vllm_config.kernel_config.specialized_decode: return None why = _deployment_error(layer, vllm_config) if why is not None: diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 62b663ff86b3..aa1059c76c74 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -43,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, @@ -55,8 +57,24 @@ if TYPE_CHECKING: from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator +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$") @@ -182,12 +200,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] @@ -201,7 +227,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( @@ -213,16 +270,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( @@ -305,6 +355,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 @@ -317,6 +372,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.model = DSparkDeepseekV4Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + from vllm.models.deepseek_v41.amd.specialized_decode.runtime.decode import ( + prepare_weights, + ) + + prepare_weights(vllm_config, self.model, draft=True) # Shared with the target (aliased by the speculator's load utility). self.lm_head = ParallelLMHead( self.config.vocab_size, @@ -355,6 +415,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: @@ -512,14 +576,17 @@ 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 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.py b/vllm/models/deepseek_v41/amd/model.py index 78a6b9f3c910..63876c482fcc 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -55,12 +55,13 @@ 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_v4.common.spec_decode import needs_mtp_hidden_buffer -from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention +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 @@ -78,6 +79,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__) @@ -239,6 +243,8 @@ def __init__( requires_grad=False, ) self.hc_attn_fn_broadcast: torch.Tensor | None = None + # set by the specialized decode when this layer runs whole on gfx942 + self.unfolded_entry = False self.hc_ffn_fn = nn.Parameter( torch.empty( (mix_hc, hc_dim), @@ -280,12 +286,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: @@ -303,21 +303,8 @@ def forward( residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, + decode_ffn: Callable | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - if self.mono is not None: - out = self.mono( - self, - 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 @@ -341,7 +328,7 @@ def forward( # layer input. The AITER op mhc_pre_delayed_aiter takes this # form. The folded form has no AITER kernel on gfx942 and runs # the eager reference instead, about 120 small kernels a step. - unfolded = self.mono is not None and self.mono.unfolded_entry + unfolded = self.unfolded_entry residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn if unfolded else self.hc_attn_fn_broadcast, @@ -426,10 +413,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 decode_ffn is not None: + 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) residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( @@ -513,6 +501,7 @@ 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 # Three aux streams: one per non-default input GEMM in # DeepseekV4Attention._run_parallel_input_projections @@ -691,6 +680,79 @@ def forward( inputs_embeds: torch.Tensor | None = None, lookback_token_ids: torch.Tensor | 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 @@ -775,37 +837,41 @@ 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, ) + + 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, + 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. + """ # 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, @@ -891,17 +957,16 @@ def _run_layers( engram_hashes, engram_mask, ) - # 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, @@ -914,6 +979,37 @@ 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], + 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, 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 ( + replay + and 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, @@ -1335,6 +1431,9 @@ 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_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( expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype) @@ -1355,6 +1454,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): "DeepSeek-V4.1 AMD fused shared experts are not supported yet; " "please open an issue." ) + from vllm.models.deepseek_v41.amd.specialized_decode.runtime.decode import ( + prepare_weights, + ) + + prepare_weights(vllm_config, self.model) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, @@ -1407,9 +1511,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: @@ -1430,6 +1534,15 @@ def forward( inputs_embeds: torch.Tensor | None = None, lookback_token_ids: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: + if self.specialized_decode is not None: + step = self.specialized_decode.plan(positions, intermediate_tensors) + if step is not None: + return self.specialized_decode.run( + step, + input_ids, + inputs_embeds, + lookback_token_ids=lookback_token_ids, + ) hidden_states = self.model( input_ids, positions, @@ -1448,12 +1561,19 @@ 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() + def process_weights_after_loading(self) -> None: + 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 + ) + 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..b7e412c0ef1e --- /dev/null +++ b/vllm/models/deepseek_v41/amd/model_state.py @@ -0,0 +1,24 @@ +# 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 +specialized decode'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 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/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py index 72a6ec2ab196..3c96927c8c02 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 ``..specialized_decode.runtime``.""" 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..07a3ba528637 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.specialized_decode.guards import tp_agree +from vllm.models.common.specialized_decode.spec import DecodeUnsupported -# ---------------------------------------------------------------- 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.""" @@ -51,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__( @@ -85,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_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py deleted file mode 100644 index 25e4fb88444e..000000000000 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ /dev/null @@ -1,766 +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 and CDNA3 (``--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 its SWA rows as the kernels read them (causal windows; -for the DSpark draft's layers, the block's non-causal rows). Anything else runs -the layer as before; the paths share every tensor at the layer boundary, so they -interleave freely. - -On CDNA3 (gfx942) the kernels compute the same layer on their own copies of -the weights (``mono/weights942.py``), at TP 4 and on steps of up to -``MAX_ROWS_GFX942`` rows. An index source's keys are written by -``mono/keylist.py`` between vLLM's indexer and K2 there, and the DSpark draft -layers also take the FFN launch, built for the draft's routing. - -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. -""" - -import functools -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 DeepseekV41ROCMAiterSparseSWAMetadata - 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) -# A gfx942 workgroup has 64 KB of LDS, not 160 KB as on gfx950. Its kernels -# take a step's tokens 6 at a time (``mono.stages.gemv.TILE``), and at TP 4 they -# take up to 42 rows, 7 requests with 5 DSpark drafts each. With 5 drafts vLLM -# captures FULL CUDA graphs of 6, 12, 18, 24, 36 and 42 rows, so a step of 5 -# requests runs padded to 36 rows. At 48 rows the shared expert has 8 token -# tiles of 36 tasks each, 288 units, and the MoE runs one unit a CTA on 256 -# CTAs. A larger step runs vLLM's path. -MAX_ROWS_GFX942 = 42 -# 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, or CDNA3 -# (gfx942), where the kernels compute the same layer with FNUZ FP8 MFMAs on -# their own copies of the weights (mono/weights942.py). -CDNA_VERSIONS = (3, 4) -# The routed experts' intermediate width a rank at TP 4 (2304 / 4). vLLM pads -# it to 640 on gfx942, and the gfx942 copies keep only the real columns. -INTER_TP4 = 576 - -_runner = None -_draft_runner = None - - -def _mono_runner(device: torch.device, draft: "MonoDecodeLayer | None" = None): - """The process's mono runner (one a TP rank, shared by every layer of the - target; the DSpark draft's layers share one of their own, as its routing and - SWA rows differ). Its peer-memory handshake is collective over the TP group: - every rank builds it at the same layer of the same step.""" - global _runner, _draft_runner - if draft is not None: - if _draft_runner is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "--monokernel's draft runner was first reached inside a CUDA " - "graph capture, before an eager draft step." - ) - from vllm.models.deepseek_v41.amd.mono.runner import DSV41MonoLayer - - _draft_runner = DSV41MonoLayer( - get_tensor_model_parallel_world_size(), - get_tensor_model_parallel_rank(), - get_tp_group().cpu_group, - device, - experts=draft.experts, - topk=draft.topk, - ug_groups=DRAFT_UG_GROUPS, - swa_width=draft.swa_width, - ) - return _draft_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 _on_gfx942() -> bool: - from vllm.platforms.rocm import get_cdna_version - - return get_cdna_version() == 3 - - -@functools.cache -def _max_rows() -> int: - return MAX_ROWS_GFX942 if _on_gfx942() else MAX_ROWS - - -def _capture_plain_experts(experts) -> None: - """gfx942: vLLM converts the routed experts to its Triton kernels' format - when the weights are loaded and frees the plain scales. Wrap that step so - the gfx942 copies (``weights942.moe_copies``) are made from the - plain tensors first.""" - qm = experts.quant_method - original = qm.process_weights_after_loading - - def process_weights_after_loading(layer): - if layer is experts: - from vllm.models.deepseek_v41.amd.mono.weights942 import moe_copies - - _keep_copies( - layer, - moe_copies( - layer.w13_weight.data, - _scale_bytes(layer.w13_weight_scale.data), - layer.w2_weight.data, - _scale_bytes(layer.w2_weight_scale.data), - INTER_TP4, - ), - ) - return original(layer) - - qm.process_weights_after_loading = process_weights_after_loading - - -def _keep_copies(layer, copies: tuple[torch.Tensor, ...]) -> None: - """Store a layer's gfx942 copies as ``layer.mono942``. A reload of the - weights (``reload_weights``, or a weight update in RL) runs - ``process_weights_after_loading`` again, and the kernels and captured - CUDA graphs read the first copies by their addresses. So the copies of a - reload are written into the first copies in place.""" - kept = getattr(layer, "mono942", None) - if kept is None: - layer.mono942 = copies - return - for old, new in zip(kept, copies, strict=True): - old.copy_(new) - - -# The attention linears that the gfx942 kernels read as ``linear_copy`` -# copies, as names of the layer's attention module. -ATTN_LINEARS_942 = ("fused_wqa_wkv", "wq_b", "wo_a", "wo_b") - - -def _copy_linear_at_load(linear) -> None: - """gfx942: make the dense MXFP8 linear's ``weights942.linear_copy`` right - after vLLM has processed its weights at load. By default - (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD) vLLM dequantizes them to bf16 - there, and the copy re-encodes them as the gfx942 kernels read them. - Weights that stay e4m3 are copied as they are.""" - qm = linear.quant_method - original = qm.process_weights_after_loading - - def process_weights_after_loading(layer): - out = original(layer) - if layer is linear: - from vllm.models.deepseek_v41.amd.mono.weights942 import linear_copy - - _keep_copies(layer, linear_copy(layer.weight, layer.weight_scale)) - return out - - qm.process_weights_after_loading = process_weights_after_loading - - -def _copy_942_at_load(layer: "DeepseekV4DecoderLayer", ffn_only: bool) -> None: - """gfx942: the kernels read their own copies of the layer's weights. Make - them when vLLM processes the weights at load, so that they count as model - memory when vLLM sizes the KV cache. Copies made at the first decode step - would count only when that step is also the CUDA graph memory profile, so - with that profile off or with ``enforce_eager`` they would not count. A - layer that takes only the FFN launch needs no attention copies.""" - ffn = layer.ffn - _capture_plain_experts(ffn.experts.routed_experts) - sh = ffn.shared_experts - # _create refuses layers without a shared expert. - assert sh is not None - linears = [sh.gate_up_proj, sh.down_proj] - if not ffn_only: - linears += [getattr(layer.attn, name) for name in ATTN_LINEARS_942] - for linear in linears: - _copy_linear_at_load(linear) - - -def _compute_units() -> int: - return current_platform.num_compute_units(torch.accelerator.current_device_index()) - - -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()) - return props.gcnArchName.split(":")[0].lower() - - -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 - if _on_gfx942(): - # The gfx942 kernels read their own copies of the experts. They are - # made from the loaded tensors before vLLM's MoE backend converts - # them (_capture_plain_experts), and moe_copies checks those tensors. - # So the backend matters on gfx950 only. - return None - 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, - draft: bool = False, - experts: int = 384, - topk: int = 6, - swa_width: int = SWA_WIDTH, - ) -> None: - self.ffn_only = ffn_only - # a layer of the DSpark draft: its own routing and (non-causal) SWA rows - self.draft, self.experts, self.topk, self.swa_width = ( - draft, - experts, - topk, - swa_width, - ) - self._weights: MonoLayerWeights | None = None - - @property - def unfolded_entry(self) -> bool: - """gfx942: when this whole first layer runs vLLM's path, its entry seam - projects the expanded residual with the unfolded hc_attn_fn, as K1 - does, and not the embedding with the folded copy (model.py). The - folded form has no AITER kernel on gfx942.""" - return not self.ffn_only and _on_gfx942() - - @staticmethod - def create( - layer: "DeepseekV4DecoderLayer", vllm_config - ) -> "MonoDecodeLayer | None": - from vllm.models.deepseek_v41.amd import ops - - # The gfx942 decode helpers run in forward passes, where vLLM's - # current config is not set, so they read the flag recorded here. - ops.record_monokernel(vllm_config) - mono = MonoDecodeLayer._create(layer, vllm_config) - if mono is not None and _on_gfx942(): - _copy_942_at_load(layer, mono.ffn_only) - return mono - - @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, 8): - why = "needs tensor parallel size 2, 4 or 8" - 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: - from vllm.models.deepseek_v41.amd.mono.common.arch import target_arch - - # every CTA of a launch stays resident: a partitioned GPU deadlocks - cus = _compute_units() - target, gpu = target_arch(), _gpu_arch() - if cus < BLOCKS: - why = f"needs {BLOCKS} compute units, the GPU has {cus}" - elif target != gpu: - # FlyDSL compiles the kernels for the ARCH environment - # variable when it is set, while the choice of the gfx942 - # weight copies follows the GPU. A stray ARCH would make - # the two disagree, so it stops the server here. - why = ( - f"needs FlyDSL to compile for the GPU's {gpu}, but the " - f"ARCH environment variable makes it compile for {target}" - ) - 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 - gfx942 = _on_gfx942() - if gfx942: - # The gfx942 indexer's top-k hooks need topk512_gfx942.cu. Build it - # when the model is created, so that a failed build shows at - # startup and the hooks leave the top-k to vLLM from the first - # step on. Later layers get the result of the first build. - from vllm.models.deepseek_v41.amd.ops import topk as topk_ops - - topk_ops.build() - draft = attn.layer_id >= config.num_hidden_layers - # a draft layer's routing is the DSpark expert set (128 / 3 on V4.1-Flash) - experts, topk = 384, 6 - if draft: - experts = getattr(config, "dspark_n_routed_experts", 0) or experts - topk = getattr(config, "dspark_num_experts_per_tok", 0) or topk - if layer.use_sequence_parallel: - why = "sequence-parallel layer" - elif not layer.fuse_seam_norm and not gfx942: - # gfx942 has no fused seam kernel. vLLM rounds the collapsed - # stream to bf16 before its RMSNorm there, which the kernels' - # fused seam does not, so the two paths differ by that rounding. - why = "needs aiter's fused mHC seam" - elif gfx942 and get_tensor_model_parallel_world_size() != 4: - why = "gfx942 builds for tensor parallel size 4 only" - elif ( - ffn.shared_experts is None - or ffn.gate.tid2eid is not None - 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 - ): - why = ( - "MoE routing outside the kernels' " - f"({experts} / {topk}, sqrtsoftplus, x1.5)" - ) - if why is not None: - logger.info_once("DSv4.1 mono decode off for some layers: %s", why) - return None - if draft and gfx942: - # On gfx942 the DSpark draft layers take the FFN launch, built for - # the draft's routing, after vLLM's attention. It runs on the - # target's runner. The whole draft layer below, on a runner of its - # own with the block's non-causal SWA rows, is checked on gfx950 - # only. vLLM's draft MoE multiplies bf16 activations with the FP4 - # experts, and the FFN launch quantizes them to MXFP8 as it does - # for the target. That changes only how many drafts the target - # accepts, not the target's output. - logger.info_once( - "DSv4.1 mono decode: the FFN launch for the DSpark draft layers " - "(%d / %d routing)", - experts, - topk, - ) - return MonoDecodeLayer(ffn_only=True, experts=experts, topk=topk) - # 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 - from vllm.models.deepseek_v41.attention import DeepseekV4Indexer - - # The indexer that _select drives. The MXFP4 indexer cache needs CDNA4, - # so gfx942 has vLLM's FP8 indexer, whose wq_b also reads K1's bf16 qr. - indexer_cls = DeepseekV4Indexer if gfx942 else 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, indexer_cls): - 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" - elif draft and (attn.compress_ratio or attn.indexer is not None): - why = "a draft layer with compressed keys" - if why is not None and draft: - logger.info_once("DSv4.1 mono decode off for the draft layers: %s", why) - return None - 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) - if draft: - from vllm.v1.attention.backends.mla.compressor_utils import ( - get_dspark_swa_index_width, - ) - - spec = vllm_config.speculative_config - logger.info_once("DSv4.1 mono decode: the DSpark draft layers") - return MonoDecodeLayer( - draft=True, - experts=experts, - topk=topk, - swa_width=get_dspark_swa_index_width( - config.sliding_window, spec.num_speculative_tokens - ), - ) - 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 - if _on_gfx942(): - self._weights = self._weights_942(layer) - return self._weights - 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(), self.experts) - return self._weights - - def _weights_942(self, layer: "DeepseekV4DecoderLayer"): - """gfx942: the layer's tensors as the gfx942 kernels read them. The - dense MXFP8 linears and the routed experts as the copies made at load - (``_copy_942_at_load``).""" - from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights - - f = layer.ffn - e, sh = f.experts.routed_experts, f.shared_experts - assert sh is not None, "DeepSeek-V4.1 layers have a shared expert" - w13, w13_s, w2, w2_s = e.mono942 - sgu, sgu_s = sh.gate_up_proj.mono942 - sw2, sw2_s = sh.down_proj.mono942 - weights = MonoLayerWeights( - attn=None if self.ffn_only else self._attn_weights_942(layer.attn), - 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=w13, - w13_s=w13_s, - w2=w2, - w2_s=w2_s, - sgu=sgu, - sgu_s=sgu_s, - sw2=sw2, - sw2_s=sw2_s, - ) - weights.check_942(f.n_routed_experts) - return weights - - @staticmethod - def _attn_weights_942(a): - from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims - from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights - - wqkv, wq_b, wo_a, wo_b = (getattr(a, name).mono942 for name in ATTN_LINEARS_942) - attn = AttnWeights( - layer_id=a.layer_id, - wqkv=wqkv[0], - wqkv_scale=wqkv[1], - q_norm=a.q_norm.weight, - kv_norm=a.kv_norm.weight, - wq_b=wq_b[0], - wq_b_scale=wq_b[1], - wo_a=wo_a[0], - wo_a_scale=wo_a[1], - wo_b=wo_b[0], - wo_b_scale=wo_b[1], - 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 _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( - "DeepseekV41ROCMAiterSparseSWAMetadata | 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 - # the draft's first layer takes the embedding already broadcast to the - # hc copies ([rows, hc, dim]), the target's the 2-D embedding - if first and x.dim() != (3 if self.draft else 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 != self.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, self if self.draft else None) - logger.info_once( - "DSv4.1 mono decode %slayer active (decode steps of <= %d rows)", - "draft " if self.draft else "", - _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: - 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 - 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, topk=self.topk) diff --git a/vllm/models/deepseek_v41/amd/ops/__init__.py b/vllm/models/deepseek_v41/amd/ops/__init__.py index 3be417427779..597c15b52c3c 100644 --- a/vllm/models/deepseek_v41/amd/ops/__init__.py +++ b/vllm/models/deepseek_v41/amd/ops/__init__.py @@ -17,18 +17,18 @@ from vllm.platforms import current_platform -# --monokernel (KernelConfig.monokernel) 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_monokernel). -_monokernel = False +# --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_monokernel(vllm_config) -> None: - """Record whether this process's model runs with --monokernel. The mono - decode layer calls this when the model is created, before any helper +def record_specialized_decode(vllm_config) -> None: + """Record whether this process's model runs with --monokernel. The + specialized decode records it when the model is created, before any helper runs.""" - global _monokernel - _monokernel = bool(vllm_config.kernel_config.monokernel) + global _specialized_decode + _specialized_decode = bool(vllm_config.kernel_config.specialized_decode) enabled.cache_clear() swa_decode_meta_enabled.cache_clear() @@ -39,7 +39,7 @@ def enabled() -> bool: 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 keeps vLLM's ops here, except for swa_decode_meta (swa_decode_meta_enabled).""" - if not _monokernel or not current_platform.is_rocm(): + if not _specialized_decode or not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_gfx942 @@ -51,7 +51,7 @@ def swa_decode_meta_enabled() -> bool: """True with --monokernel on a CDNA3 (gfx942) or CDNA4 (gfx950) GPU. swa_decode_meta is plain Triton and writes the buffers that the attention of both reads.""" - if not _monokernel or not current_platform.is_rocm(): + if not _specialized_decode or not current_platform.is_rocm(): return False from vllm.platforms.rocm import get_cdna_version diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 3d694b2b9106..8fa52373638f 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -514,13 +514,15 @@ class DeepseekV41Gfx942Indexer(DeepseekV4Indexer): attn_cls = Gfx942SparseAttnIndexer -def _gfx942_monokernel() -> bool: +def _gfx942_specialized_decode() -> bool: """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 - return bool(get_current_vllm_config().kernel_config.monokernel) and on_gfx942() + return ( + bool(get_current_vllm_config().kernel_config.specialized_decode) and on_gfx942() + ) class DeepseekV41ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): @@ -554,7 +556,7 @@ def _indexer_cls( ) -> type[DeepseekV4Indexer]: if k_cache is not None and k_cache.rocm_mxfp4: return DeepseekV41RocmMxfp4Indexer - if _gfx942_monokernel(): + if _gfx942_specialized_decode(): return DeepseekV41Gfx942Indexer return super()._indexer_cls(k_cache) diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py b/vllm/models/deepseek_v41/amd/specialized_decode/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v41/amd/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/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..b2fb5a89b576 --- /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 ``--monokernel``. 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..54d01c610fc1 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/decode.py @@ -0,0 +1,633 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""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 +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 ( + DRAFT_UG_GROUPS, + EXPERTS, + MAX_ROWS, + RECORD, + SWA_WIDTH, + TOPK, + UG_GROUPS, + draft_routing, + draft_swa_width, + on_gfx942, +) +from .weights942 import scale_bytes as _scale_bytes + +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). + + 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)) + + +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 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.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.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. + 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] = {} + 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) + if plan is None: + continue + self.plans[layer.attn.layer_id] = plan + if plan.ffn_only: + continue + 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 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: + """Info needed to run the specialized decode kernel for this layer.""" + attn = layer.attn + why = _skip_layer(layer, self.experts, self.topk, tp_size) + 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 + if self.draft and on_gfx942(): + # On gfx942 the draft layers take the FFN launch after vLLM's + # attention. vLLM's draft MoE multiplies bf16 activations with the + # FP4 experts, and the FFN launch quantizes them to MXFP8 as it does + # for the target. That changes only how many drafts the target + # accepts, not the target's output. + logger.info_once( + "DSv4.1 specialized decode: the FFN launch for the draft layers " + "(%d / %d routing)", + self.experts, + self.topk, + ) + return self._ffn_plan(layer, tp_size) + why = _ffn_only(layer) + 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 + if why is not None: + logger.info_once( + "DSv4.1 specialized decode: the FFN launch for layers with %s", why + ) + return self._ffn_plan(layer, tp_size) + weights = _weights(layer, False, tp_size, self.experts) + # gfx942: the entry seam of a whole first layer projects the expanded + # residual with the unfolded hc_attn_fn, as K1 does, and not the + # embedding with the folded copy, which has no AITER kernel there. + layer.unfolded_entry = on_gfx942() + 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 = first 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 _ffn_plan(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 + return LayerPlan( + layer, _weights(layer, True, tp_size, self.experts), ffn_only=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", self.experts, self.topk] = 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], key[2], key[3]) + 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 + # 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: + 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 + 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, topk=self.topk) + + +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=UG_GROUPS, + swa_width=SWA_WIDTH, + ) + + 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, + ) + + +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, + # the draft's unions are small on gfx950, where it runs whole; on + # gfx942 it takes the FFN launch, validated at the target's groups + ug_groups=UG_GROUPS if on_gfx942() else 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) + 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, 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, + decode_ffn=ffn, + ) + hidden, residual, post_mix, res_mix, pre_mix = out + return model.epilogue( + hidden, residual, post_mix, res_mix, pre_mix, full_num_tokens + ) + + +def prepare_weights(vllm_config, model, draft: bool = False) -> None: + """Make ready, at model creation, whatever the launches need from the + weights before vLLM loads them: the flag the gfx942 decode helpers read, the + top-k hooks they call, and the gfx942 kernels' own copies of the weights, + which vLLM then makes as it processes each layer's at load. + + Args: + vllm_config: the config of the model being created. + model: the backbone, or the DSpark draft head's. + draft: whether ``model`` is the draft head's backbone. + + """ + from vllm.models.deepseek_v41.amd import ops + + # The gfx942 decode helpers run in forward passes, where vLLM's current + # config is not set, so they read the flag recorded here. + ops.record_specialized_decode(vllm_config) + if not vllm_config.kernel_config.specialized_decode or not on_gfx942(): + return + from vllm.models.deepseek_v41.amd.ops import topk as topk_ops + + from .weights942 import copy_at_load + + # A failed build shows at startup, and the hooks then leave the top-k to + # vLLM from the first step on. + topk_ops.build() + experts, topk = draft_routing(vllm_config) if draft else (EXPERTS, TOPK) + tp_size = vllm_config.parallel_config.tensor_parallel_size + for layer in model.layers: + if _skip_layer(layer, experts, topk, tp_size) is None: + copy_at_load(layer, ffn_only=draft or _ffn_only(layer) is not 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 + from vllm.models.deepseek_v41.attention import DeepseekV4Indexer + + # The MXFP4 indexer cache needs CDNA4, so gfx942 drives vLLM's FP8 indexer, + # whose wq_b also reads K1's bf16 qr. + indexer_cls = DeepseekV4Indexer if on_gfx942() else 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, indexer_cls): + 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, experts: int, topk: int, tp_size: 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 + gfx942 = on_gfx942() + if layer.use_sequence_parallel: + return "sequence-parallel layer" + if not layer.fuse_seam_norm and not gfx942: + # gfx942 has no fused seam kernel. vLLM rounds the collapsed stream to + # bf16 before its RMSNorm there, which the kernels' fused seam does not, + # so the two paths differ by that rounding. + return "needs aiter's fused mHC seam" + if gfx942 and tp_size != 4: + return "gfx942 builds for tensor parallel size 4 only" + if ( + ffn.shared_experts is None + or ffn.gate.tid2eid is not None + 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 ( + f"MoE routing outside the kernels' ({experts} / {topk}, sqrtsoftplus, x1.5)" + ) + return None + + +def _weights(layer, ffn_only: bool, world: int, experts: int): + from vllm.models.deepseek_v41.amd.mono.runner import ( + MonoLayerWeights, + ) + + if on_gfx942(): + from .weights942 import layer_weights + + return layer_weights(layer, ffn_only, world) + 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, experts) + 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 diff --git a/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py new file mode 100644 index 000000000000..05b81c2b5ae4 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/runtime/weights942.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The gfx942 kernels' own copies of a decoder layer's weights. + +On CDNA3 the kernels compute the same layer with FNUZ FP8 MFMAs, from tensors +re-encoded as they read them (``mono/weights942.py``). The copies are made as +vLLM processes the loaded weights, so that they count as model memory when vLLM +sizes the KV cache: copies made at the first decode step would count only when +that step is also the CUDA graph memory profile. +""" + +from typing import TYPE_CHECKING + +import torch + +from ..spec import INTER_TP4 + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + +# the attention linears the kernels read as ``linear_copy`` copies, as names of +# the layer's attention module +ATTN_LINEARS = ("fused_wqa_wkv", "wq_b", "wo_a", "wo_b") + + +def scale_bytes(scale: torch.Tensor) -> torch.Tensor: + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _keep(layer, copies: tuple[torch.Tensor, ...]) -> None: + """Store a layer's copies as ``layer.mono942``. A reload of the weights + runs ``process_weights_after_loading`` again, and the kernels and captured + CUDA graphs read the first copies by their addresses, so a reload's copies + are written into the first ones in place.""" + kept = getattr(layer, "mono942", None) + if kept is None: + layer.mono942 = copies + return + for old, new in zip(kept, copies, strict=True): + old.copy_(new) + + +def _capture_plain_experts(experts) -> None: + """Take the copies before vLLM converts the routed experts to its Triton kernels' + format when the weights are loaded, which frees the plain scales.""" + qm = experts.quant_method + original = qm.process_weights_after_loading + + def process_weights_after_loading(layer): + if layer is experts: + from ...mono.weights942 import moe_copies + + _keep( + layer, + moe_copies( + layer.w13_weight.data, + scale_bytes(layer.w13_weight_scale.data), + layer.w2_weight.data, + scale_bytes(layer.w2_weight_scale.data), + INTER_TP4, + ), + ) + return original(layer) + + qm.process_weights_after_loading = process_weights_after_loading + + +def _copy_linear(linear) -> None: + """By default (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD) vLLM dequantizes a + dense MXFP8 linear's weights to bf16, and the copy re-encodes them as the + kernels read them. Weights that stay e4m3 are copied as they are.""" + qm = linear.quant_method + original = qm.process_weights_after_loading + + def process_weights_after_loading(layer): + out = original(layer) + if layer is linear: + from ...mono.weights942 import linear_copy + + _keep(layer, linear_copy(layer.weight, layer.weight_scale)) + return out + + qm.process_weights_after_loading = process_weights_after_loading + + +def copy_at_load(layer: "DeepseekV4DecoderLayer", ffn_only: bool) -> None: + """Make this layer's copies when vLLM processes its weights at load. A + layer that takes only the FFN launch needs no attention copies.""" + ffn = layer.ffn + _capture_plain_experts(ffn.experts.routed_experts) + sh = ffn.shared_experts + assert sh is not None + linears = [sh.gate_up_proj, sh.down_proj] + if not ffn_only: + linears += [getattr(layer.attn, name) for name in ATTN_LINEARS] + for linear in linears: + _copy_linear(linear) + + +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 + + f = layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + assert sh is not None + w13, w13_s, w2, w2_s = e.mono942 + sgu, sgu_s = sh.gate_up_proj.mono942 + sw2, sw2_s = sh.down_proj.mono942 + weights = MonoLayerWeights( + attn=None if ffn_only else _attn_weights(layer.attn, 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=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + ) + weights.check_942(f.n_routed_experts) + return weights + + +def _attn_weights(a, world: int): + 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( + layer_id=a.layer_id, + wqkv=wqkv[0], + wqkv_scale=wqkv[1], + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=wq_b[0], + wq_b_scale=wq_b[1], + wo_a=wo_a[0], + wo_a_scale=wo_a[1], + wo_b=wo_b[0], + wo_b_scale=wo_b[1], + 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..6dc5ae39e553 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/specialized_decode/spec.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1's specialized decode.""" + +import functools +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 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) +# A gfx942 workgroup has 64 KB of LDS, not 160 KB as on gfx950. Its kernels take +# a step's tokens 6 at a time (mono.stages.gemv.TILE), and at TP 4 they take up +# to 7 requests with 5 DSpark drafts each. A larger step runs vLLM's path. +MAX_ROWS_GFX942 = 42 +EXPERTS = 384 # the backbone's routed experts, top-TOPK of them a token +TOPK = 6 +# FP4 groups of an ug task +UG_GROUPS = 2 +# the draft's 128-expert unions are small +DRAFT_UG_GROUPS = 1 +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats, or CDNA3 +# (gfx942), where the kernels compute the same layer with FNUZ FP8 MFMAs on +# their own copies of the weights (runtime/weights942.py) +CDNA_VERSIONS = (3, 4) +# the routed experts' intermediate width a rank at TP 4 (2304 / 4): vLLM pads it +# to 640 on gfx942, and the gfx942 copies keep only the real columns +INTER_TP4 = 576 + + +def on_gfx942() -> bool: + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 3 + + +@functools.cache +def max_rows() -> int: + """The widest step the kernels take on this GPU.""" + return MAX_ROWS_GFX942 if on_gfx942() else MAX_ROWS + + +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 + none.""" + rows = max_rows() + sizes = vllm_config.compilation_config.cudagraph_capture_sizes or [] + served = tuple(s for s in sorted(sizes) if 1 <= s <= rows) + 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()) + return props.gcnArchName.split(":")[0].lower() + + +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, 8): + reasons.append("needs tensor parallel size 2, 4 or 8") + 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.common.arch import ( + target_arch, + ) + 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 = current_platform.num_compute_units(torch.accelerator.current_device_index()) + if cus < BLOCKS: + reasons.append(f"needs {BLOCKS} compute units, the GPU has {cus}") + target, gpu = target_arch(), _gpu_arch() + if target != gpu: + # FlyDSL compiles the kernels for the ARCH environment variable when it + # is set, while the choice of the gfx942 weight copies follows the GPU. + reasons.append( + f"needs FlyDSL to compile for the GPU's {gpu}, but the ARCH " + f"environment variable makes it compile for {target}" + ) + if on_gfx942(): + # The gfx942 kernels read their own copies of the experts, made from the + # loaded tensors before vLLM's MoE backend converts them, so the backend + # matters on gfx950 only. + return reasons + 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: "LayerLaunches", caches) -> list[str]: + """Every reason the kernels cannot read these caches: they index each as + [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): + 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]") + 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 + + +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: + 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 != decode.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 _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=_SUPPORTS, + widths=_widths, + refuse_config=_refuse_config, + refuse_model=_refuse_model, + refuse_caches=_refuse_caches, + skip_step=_skip_step, + build=_build, +) + +DRAFT_DECODE_SPEC = DecodeSpec["DSV41DraftKernel"]( + name="dsv41-dspark-cdna4", + supports=_SUPPORTS, + widths=_draft_widths, + set_capture_sizes=False, + refuse_config=_refuse_config, + refuse_model=_refuse_model, + refuse_caches=_refuse_caches, + skip_step=_skip_draft_step, + build=_build_draft, +) diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 60e4e55cde26..2bbe3dc9f19f 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -293,9 +293,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: @@ -340,15 +344,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # one contiguous group. mapped = stream_language_model_first(self.hf_to_vllm_mapper.apply(weights)) 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() 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: 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 aec4f52579d5..37bff0e39171 100644 --- a/vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py +++ b/vllm/models/minimax_m3/amd/specialized_decode/runtime/spec.py @@ -100,7 +100,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)