From 42c2a93c48b48c8ebc9c0a3bc3225f5fec7e0db8 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 3 Sep 2026 07:51:53 +0000 Subject: [PATCH 1/4] [Bugfix] Fall back to full decode graphs for noncompiled models Resolve model paths without a piecewise provider to safe full-decode or eager modes when breakable graphs are unavailable, preserving platform compile paths and rejecting late-invalid combinations. Co-authored-by: OpenAI Codex Signed-off-by: Andreas Karatzas --- tests/test_config.py | 416 ++++++++++++++++++++++++++++++++++++++++--- vllm/config/vllm.py | 184 +++++++++++++++++-- 2 files changed, 563 insertions(+), 37 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 5254327f71bd..50c4fb5fa279 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -239,8 +239,7 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): - """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled - MRV1 paths and off breakable cudagraphs by default.""" + """ROCm keeps the DSA models on MRV1 and off breakable graphs by default.""" from vllm.config.vllm import ( ROCM_DEFAULT_MRV1_ARCHITECTURES, default_breakable_cudagraph_architectures, @@ -329,6 +328,63 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph( default_breakable_cudagraph_architectures.cache_clear() +class _CudagraphConfigStub(SimpleNamespace): + _uses_breakable_cudagraph_by_default = ( + VllmConfig._uses_breakable_cudagraph_by_default + ) + _uses_noncompiled_cudagraph_path = VllmConfig._uses_noncompiled_cudagraph_path + _piecewise_cudagraph_provider_available = ( + VllmConfig._piecewise_cudagraph_provider_available + ) + + +@pytest.fixture +def make_cudagraph_config(monkeypatch): + from vllm.config.vllm import default_breakable_cudagraph_architectures + + def make( + architecture="DeepseekV4ForConditionalGeneration", + *, + rocm=True, + gfx950=False, + breakable=False, + optimization_level=OptimizationLevel.O2, + compilation_mode=None, + cudagraph_mode=None, + use_v2=True, + compilation_config=None, + ): + if breakable is None: + monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False) + else: + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", str(int(breakable))) + monkeypatch.setattr(current_platform, "is_cuda", lambda: not rocm) + monkeypatch.setattr(current_platform, "is_rocm", lambda: rocm) + monkeypatch.setattr( + current_platform, + "is_device_capability", + lambda value: gfx950 and value == 95, + ) + default_breakable_cudagraph_architectures.cache_clear() + resolved_compilation_config = ( + compilation_config + if compilation_config is not None + else CompilationConfig(mode=compilation_mode, cudagraph_mode=cudagraph_mode) + ) + return _CudagraphConfigStub( + model_config=SimpleNamespace( + architecture=architecture, architectures=[architecture] + ), + compilation_config=resolved_compilation_config, + optimization_level=optimization_level, + use_v2_model_runner=use_v2, + ) + + yield make + os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None) + default_breakable_cudagraph_architectures.cache_clear() + + @pytest.mark.parametrize( ("architecture", "is_rocm", "expected"), [ @@ -336,34 +392,354 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph( ("DeepseekV32ForCausalLM", True, False), ("DeepseekV32MTPModel", False, True), ("DeepseekV32MTPModel", True, False), + ("DeepseekV4ForCausalLM", False, True), + ("DeepseekV4ForCausalLM", True, False), ("GlmMoeDsaForCausalLM", False, True), ("GlmMoeDsaForCausalLM", True, False), ], ) def test_dsa_breakable_cudagraph_platform_default( - monkeypatch, architecture, is_rocm, expected + make_cudagraph_config, architecture, is_rocm, expected ): - from vllm.config.vllm import default_breakable_cudagraph_architectures - from vllm.platforms import current_platform + config = make_cudagraph_config( + architecture, + rocm=is_rocm, + breakable=None, + use_v2=not is_rocm, + ) - monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False) - monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm) - default_breakable_cudagraph_architectures.cache_clear() - config = SimpleNamespace( - model_config=SimpleNamespace(architectures=[architecture]), - compilation_config=CompilationConfig(), + assert VllmConfig._maybe_enable_breakable_cudagraph(config) is expected + if expected: + assert config.compilation_config.mode == CompilationMode.NONE + else: + assert config.compilation_config.mode is None + assert config.compilation_config.cudagraph_mode is None + + +@pytest.mark.parametrize( + "architecture", + sorted(vllm_config_module.NON_COMPILED_CUDAGRAPH_FALLBACK_ARCHITECTURES), +) +def test_noncompiled_architectures_fall_back_when_breakable_disabled( + make_cudagraph_config, architecture +): + config = make_cudagraph_config(architecture, rocm=False) + + assert not VllmConfig._maybe_enable_breakable_cudagraph(config) + assert config.compilation_config.mode == CompilationMode.NONE + assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY + + +@pytest.mark.parametrize( + ("optimization_level", "expected_cudagraph_mode"), + [ + (OptimizationLevel.O0, CUDAGraphMode.NONE), + (OptimizationLevel.O1, CUDAGraphMode.NONE), + (OptimizationLevel.O2, CUDAGraphMode.FULL_DECODE_ONLY), + (OptimizationLevel.O3, CUDAGraphMode.FULL_DECODE_ONLY), + ], +) +def test_noncompiled_cudagraph_fallback_respects_optimization_level( + make_cudagraph_config, optimization_level, expected_cudagraph_mode +): + config = make_cudagraph_config( + optimization_level=optimization_level, ) - config._uses_breakable_cudagraph_by_default = lambda: ( - VllmConfig._uses_breakable_cudagraph_by_default(config) + + assert not VllmConfig._maybe_enable_breakable_cudagraph(config) + assert config.compilation_config.mode == CompilationMode.NONE + assert config.compilation_config.cudagraph_mode == expected_cudagraph_mode + + +@pytest.mark.parametrize( + ("architecture", "expected_cudagraph_mode"), + [ + ("DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY), + ("DeepseekV32MTPModel", None), + ("DeepseekV32ForCausalLM", None), + ("GlmMoeDsaForCausalLM", None), + ], +) +def test_rocm_forced_v2_only_falls_back_for_noncompiled_dsa_model( + make_cudagraph_config, architecture, expected_cudagraph_mode +): + config = make_cudagraph_config(architecture) + + assert not VllmConfig._maybe_enable_breakable_cudagraph(config) + expected_compilation_mode = ( + CompilationMode.NONE if expected_cudagraph_mode is not None else None ) + assert config.compilation_config.mode == expected_compilation_mode + assert config.compilation_config.cudagraph_mode == expected_cudagraph_mode - try: - assert VllmConfig._maybe_enable_breakable_cudagraph(config) is expected - if expected: - assert config.compilation_config.mode == CompilationMode.NONE - finally: - os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None) - default_breakable_cudagraph_architectures.cache_clear() + +@pytest.mark.parametrize( + ("architecture", "use_v2", "input_mode", "expected_compile", "expected_graph"), + [ + ("DeepseekV4ForCausalLM", True, None, CompilationMode.NONE, CUDAGraphMode.NONE), + ( + "DeepseekV4ForConditionalGeneration", + True, + None, + CompilationMode.NONE, + CUDAGraphMode.NONE, + ), + ( + "DeepseekV4ForConditionalGeneration", + True, + CUDAGraphMode.FULL, + CompilationMode.NONE, + CUDAGraphMode.FULL, + ), + ( + "DeepseekV4ForConditionalGeneration", + True, + CUDAGraphMode.FULL_DECODE_ONLY, + CompilationMode.NONE, + CUDAGraphMode.FULL_DECODE_ONLY, + ), + ( + "DeepseekV4ForConditionalGeneration", + True, + CUDAGraphMode.FULL_AND_PIECEWISE, + CompilationMode.NONE, + CUDAGraphMode.FULL_DECODE_ONLY, + ), + ( + "KimiK3ForConditionalGeneration", + True, + None, + CompilationMode.NONE, + CUDAGraphMode.FULL_DECODE_ONLY, + ), + ], +) +def test_rocm_gfx950_noncompiled_cudagraph_policy( + make_cudagraph_config, + architecture, + use_v2, + input_mode, + expected_compile, + expected_graph, +): + config = make_cudagraph_config( + architecture, gfx950=True, cudagraph_mode=input_mode, use_v2=use_v2 + ) + + breakable_enabled = VllmConfig._maybe_enable_breakable_cudagraph(config) + VllmConfig._normalize_unavailable_piecewise_cudagraphs( + config, breakable_cudagraph_enabled=breakable_enabled + ) + + assert not breakable_enabled + assert config.compilation_config.mode == expected_compile + assert config.compilation_config.cudagraph_mode == expected_graph + + +@pytest.mark.parametrize( + ("compile_mode", "graph_mode", "expected_graph", "should_raise"), + [ + (None, CUDAGraphMode.NONE, CUDAGraphMode.NONE, False), + (None, CUDAGraphMode.PIECEWISE, CUDAGraphMode.NONE, False), + (None, CUDAGraphMode.FULL, CUDAGraphMode.FULL, False), + (None, CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_DECODE_ONLY, False), + ( + None, + CUDAGraphMode.FULL_AND_PIECEWISE, + CUDAGraphMode.FULL_DECODE_ONLY, + False, + ), + ( + CompilationMode.VLLM_COMPILE, + None, + None, + True, + ), + ( + CompilationMode.VLLM_COMPILE, + CUDAGraphMode.NONE, + None, + True, + ), + ( + CompilationMode.VLLM_COMPILE, + CUDAGraphMode.FULL_DECODE_ONLY, + None, + True, + ), + (CompilationMode.VLLM_COMPILE, CUDAGraphMode.PIECEWISE, None, True), + (CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL, None, True), + ( + CompilationMode.VLLM_COMPILE, + CUDAGraphMode.FULL_AND_PIECEWISE, + None, + True, + ), + ], +) +def test_noncompiled_cudagraph_fallback_validates_explicit_modes( + make_cudagraph_config, + compile_mode, + graph_mode, + expected_graph, + should_raise, +): + config = make_cudagraph_config( + compilation_mode=compile_mode, cudagraph_mode=graph_mode + ) + if should_raise: + with pytest.raises(ValueError, match="unavailable piecewise capture"): + VllmConfig._maybe_enable_breakable_cudagraph(config) + return + + breakable_enabled = VllmConfig._maybe_enable_breakable_cudagraph(config) + VllmConfig._normalize_unavailable_piecewise_cudagraphs( + config, breakable_cudagraph_enabled=breakable_enabled + ) + + assert not breakable_enabled + assert config.compilation_config.mode == (compile_mode or CompilationMode.NONE) + assert config.compilation_config.cudagraph_mode == expected_graph + + +@pytest.mark.parametrize( + ("architecture", "breakable", "expected_enabled", "expected_compile"), + [ + ("LlamaForCausalLM", False, False, None), + ( + "DeepseekV4ForConditionalGeneration", + True, + True, + CompilationMode.NONE, + ), + ], +) +def test_noncompiled_cudagraph_fallback_controls( + make_cudagraph_config, + architecture, + breakable, + expected_enabled, + expected_compile, +): + config = make_cudagraph_config(architecture, breakable=breakable) + + assert VllmConfig._maybe_enable_breakable_cudagraph(config) is expected_enabled + assert config.compilation_config.mode == expected_compile + assert config.compilation_config.cudagraph_mode is None + + +@pytest.mark.parametrize( + ( + "architecture", + "compile_mode", + "graph_mode", + "breakable", + "expected_graph", + "expected_sizes", + ), + [ + ( + "DeepseekV4ForConditionalGeneration", + CompilationMode.NONE, + CUDAGraphMode.NONE, + False, + CUDAGraphMode.NONE, + [], + ), + ( + "DeepseekV4ForConditionalGeneration", + CompilationMode.NONE, + CUDAGraphMode.PIECEWISE, + False, + CUDAGraphMode.NONE, + [], + ), + ( + "DeepseekV4ForConditionalGeneration", + CompilationMode.NONE, + CUDAGraphMode.FULL_AND_PIECEWISE, + False, + CUDAGraphMode.FULL_DECODE_ONLY, + [1, 2, 4], + ), + ( + "DeepseekV4ForConditionalGeneration", + CompilationMode.NONE, + CUDAGraphMode.PIECEWISE, + True, + CUDAGraphMode.PIECEWISE, + [1, 2, 4], + ), + ( + "LlamaForCausalLM", + CompilationMode.VLLM_COMPILE, + CUDAGraphMode.PIECEWISE, + False, + CUDAGraphMode.PIECEWISE, + [1, 2, 4], + ), + ], +) +def test_late_piecewise_override_is_normalized_by_available_provider( + make_cudagraph_config, + architecture, + compile_mode, + graph_mode, + breakable, + expected_graph, + expected_sizes, +): + compilation_config = CompilationConfig( + mode=compile_mode, + cudagraph_mode=graph_mode, + cudagraph_capture_sizes=[1, 2, 4], + max_cudagraph_capture_size=4, + ) + config = make_cudagraph_config( + architecture, + breakable=breakable, + compilation_config=compilation_config, + ) + + VllmConfig._normalize_unavailable_piecewise_cudagraphs( + config, breakable_cudagraph_enabled=breakable + ) + + assert compilation_config.cudagraph_mode == expected_graph + assert compilation_config.cudagraph_capture_sizes == expected_sizes + expected_max_size = 0 if expected_graph == CUDAGraphMode.NONE else 4 + assert compilation_config.max_cudagraph_capture_size == expected_max_size + + +@pytest.mark.parametrize( + ("architecture", "mode", "breakable_enabled", "should_raise"), + [ + ("DeepseekV4ForConditionalGeneration", CompilationMode.NONE, False, True), + ("DeepseekV4ForConditionalGeneration", CompilationMode.NONE, True, False), + ("LlamaForCausalLM", CompilationMode.VLLM_COMPILE, False, False), + ], +) +def test_adaptive_verification_requires_piecewise_cudagraph_provider( + make_cudagraph_config, architecture, mode, breakable_enabled, should_raise +): + config = make_cudagraph_config( + architecture, + breakable=breakable_enabled, + compilation_mode=mode, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + ) + config.speculative_config = SimpleNamespace(enable_adaptive_verification=True) + config.lora_config = None + config.parallel_config = SimpleNamespace(pipeline_parallel_size=1) + + validate = lambda: VllmConfig._validate_adaptive_verification( + config, breakable_cudagraph_enabled=breakable_enabled + ) + if should_raise: + with pytest.raises(ValueError, match="requires piecewise CUDA graphs"): + validate() + else: + validate() @pytest.mark.parametrize( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index e4c223daf69b..ecb98b462789 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -72,7 +72,11 @@ {"DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", "GlmMoeDsaForCausalLM"} ) -DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( +# At least one platform implementation of each architecture below has no +# active torch.compile boundary and uses breakable graphs for PIECEWISE +# capture. Platform-specific compiled implementations are carved out by +# _uses_noncompiled_cudagraph_path(), including compatibility-policy exceptions. +NON_COMPILED_CUDAGRAPH_FALLBACK_ARCHITECTURES = frozenset( { "DeepseekV32MTPModel", "DeepseekV32ForCausalLM", @@ -94,6 +98,22 @@ } ) +# These implementations also define the current default policy. Keep the +# capability set separately named because policy can change independently. +DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = ( + NON_COMPILED_CUDAGRAPH_FALLBACK_ARCHITECTURES +) + +# DeepSeek V3.2 and its MTP model use their noncompiled implementations on +# CUDA, but the shared implementations selected elsewhere are compiled. +NON_CUDA_COMPILED_CUDAGRAPH_ARCHITECTURES = frozenset( + { + "DeepseekV32MTPModel", + "DeepseekV32ForCausalLM", + "GlmMoeDsaForCausalLM", + } +) + @lru_cache def default_breakable_cudagraph_architectures() -> frozenset[str]: @@ -712,6 +732,33 @@ def _uses_breakable_cudagraph_by_default(self) -> bool: architectures = set(model_config.architectures) return bool(architectures & default_breakable_cudagraph_architectures()) + def _uses_noncompiled_cudagraph_path(self) -> bool: + """Whether the selected path needs the noncompiled graph fallback.""" + model_config = self.model_config + if model_config is None: + return False + + architecture = model_config.architecture + if architecture not in NON_COMPILED_CUDAGRAPH_FALLBACK_ARCHITECTURES: + return False + + from vllm.platforms import current_platform + + if ( + not current_platform.is_cuda() + and architecture in NON_CUDA_COMPILED_CUDAGRAPH_ARCHITECTURES + ): + return False + + # Compatibility policy: keep DeepSeek V4's existing ROCm MRV1 graph + # defaults unchanged. MRV2 adds a runtime provider guard and therefore + # needs the fallback when explicitly selected. + return not ( + current_platform.is_rocm() + and architecture == "DeepseekV4ForCausalLM" + and not self.use_v2_model_runner + ) + def _maybe_enable_breakable_cudagraph(self) -> bool: if ( "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ @@ -730,6 +777,53 @@ def _maybe_enable_breakable_cudagraph(self) -> bool: enabled = is_breakable_cudagraph_enabled() if enabled: self.compilation_config.mode = CompilationMode.NONE + elif self._uses_noncompiled_cudagraph_path(): + # These architectures do not have an active torch.compile wrapper, + # so PIECEWISE capture is unavailable without breakable graphs. Keep + # full graphs for uniform decode at O2/O3 and run other batches + # eagerly. This also handles platforms such as ROCm that deliberately + # leave breakable graphs disabled by default for performance. + compilation_config = self.compilation_config + if compilation_config.mode == CompilationMode.VLLM_COMPILE: + raise ValueError( + f"{self.model_config.architecture} does not expose an active " + "torch.compile boundary, so compilation mode VLLM_COMPILE " + "cannot be honored and later CUDA graph resolution can require " + "unavailable piecewise capture. Use compilation mode NONE with " + "cudagraph mode NONE/FULL_DECODE_ONLY, or enable breakable CUDA " + "graphs." + ) + if compilation_config.cudagraph_mode is None: + from vllm.platforms import current_platform + + architecture = self.model_config.architecture + unsafe_full_graph = ( + current_platform.is_rocm() + and current_platform.is_device_capability(95) + and self.use_v2_model_runner + and architecture + in { + "DeepseekV4ForCausalLM", + "DeepseekV4ForConditionalGeneration", + } + ) + compilation_config.cudagraph_mode = ( + CUDAGraphMode.FULL_DECODE_ONLY + if self.optimization_level >= OptimizationLevel.O2 + and not unsafe_full_graph + else CUDAGraphMode.NONE + ) + logger.info_once( + "Breakable CUDA graphs are disabled for a model without " + "torch.compile support; defaulting cudagraph mode to %s.", + compilation_config.cudagraph_mode.name, + ) + + # None is the unset value, not an explicit compile request. These + # wrappers cannot provide piecewise compilation, so resolve it now + # and let the compatibility pass below remove impossible modes. + if compilation_config.mode is None: + compilation_config.mode = CompilationMode.NONE return enabled @property @@ -966,6 +1060,51 @@ def _maybe_disable_dynamic_sd_for_data_parallel(self) -> None: ) speculative_config.num_speculative_tokens_per_batch_size = None + def _normalize_unavailable_piecewise_cudagraphs( + self, *, breakable_cudagraph_enabled: bool + ) -> None: + """Remove piecewise graphs when no capture implementation is active.""" + compilation_config = self.compilation_config + if compilation_config.cudagraph_mode == CUDAGraphMode.NONE: + compilation_config.max_cudagraph_capture_size = 0 + compilation_config.cudagraph_capture_sizes = [] + return + + piecewise_capture_available = ( + VllmConfig._piecewise_cudagraph_provider_available( + self, + breakable_cudagraph_enabled=breakable_cudagraph_enabled, + ) + ) + if ( + compilation_config.cudagraph_mode.requires_piecewise_compilation() + and not piecewise_capture_available + ): + fallback_mode = ( + CUDAGraphMode.FULL_DECODE_ONLY + if compilation_config.cudagraph_mode.has_full_cudagraphs() + else CUDAGraphMode.NONE + ) + logger.info_once( + "Cudagraph mode %s is not compatible with compilation mode %s. " + "Overriding to %s.", + compilation_config.cudagraph_mode, + compilation_config.mode, + fallback_mode, + ) + compilation_config.cudagraph_mode = fallback_mode + if fallback_mode == CUDAGraphMode.NONE: + compilation_config.max_cudagraph_capture_size = 0 + compilation_config.cudagraph_capture_sizes = [] + + def _piecewise_cudagraph_provider_available( + self, *, breakable_cudagraph_enabled: bool + ) -> bool: + return breakable_cudagraph_enabled or ( + self.compilation_config.mode == CompilationMode.VLLM_COMPILE + and not self._uses_noncompiled_cudagraph_path() + ) + def _post_init_kv_transfer_config(self) -> None: """Update KVTransferConfig based on top-level configs in VllmConfig. @@ -1409,7 +1548,7 @@ def __post_init__(self): ) ): logger.warning_once( - "Inductor compilation was disabled by user settings, " + "Inductor compilation is disabled by configuration, " "optimizations settings that are only active during " "inductor compilation will be ignored." ) @@ -1471,19 +1610,6 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() - if ( - self.compilation_config.cudagraph_mode.requires_piecewise_compilation() - and self.compilation_config.mode != CompilationMode.VLLM_COMPILE - and not envs.VLLM_USE_BREAKABLE_CUDAGRAPH - ): - logger.info_once( - "Cudagraph mode %s is not compatible with compilation mode %s." - "Overriding to NONE.", - self.compilation_config.cudagraph_mode, - self.compilation_config.mode, - ) - self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - # async tp is built on top of sequence parallelism and requires it. pass_config = self.compilation_config.pass_config if pass_config.fuse_gemm_comms: @@ -1595,6 +1721,9 @@ def has_blocked_weights(): else: self.compilation_config.cudagraph_num_of_warmups = 1 + self._normalize_unavailable_piecewise_cudagraphs( + breakable_cudagraph_enabled=breakable_cudagraph_enabled + ) self._set_cudagraph_sizes() else: @@ -1651,6 +1780,13 @@ def has_blocked_weights(): ) current_platform.check_and_update_config(self) + # Platform and connector compatibility checks above can replace a full + # graph mode with PIECEWISE. Re-normalize after those late updates and + # before validating consumers of the resolved mode. + self._normalize_unavailable_piecewise_cudagraphs( + breakable_cudagraph_enabled=breakable_cudagraph_enabled + ) + self._resolve_allow_missing_mm_embeddings() self._resolve_mm_processor_device() self._validate_mm_processor_device() @@ -1661,7 +1797,9 @@ def has_blocked_weights(): self._validate_v1_model_runner() self._validate_batch_sharded_sampling() - self._validate_adaptive_verification() + self._validate_adaptive_verification( + breakable_cudagraph_enabled=breakable_cudagraph_enabled + ) # Re-compute compile ranges after platform-specific config updates # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled) @@ -2648,7 +2786,9 @@ def _get_v1_model_runner_unsupported_features(self) -> list[str]: return unsupported - def _validate_adaptive_verification(self) -> None: + def _validate_adaptive_verification( + self, *, breakable_cudagraph_enabled: bool + ) -> None: spec_config = self.speculative_config if not spec_config or not spec_config.enable_adaptive_verification: return @@ -2660,6 +2800,16 @@ def _validate_adaptive_verification(self) -> None: "Adaptive verification is not currently compatible with LoRA" ) + if not VllmConfig._piecewise_cudagraph_provider_available( + self, + breakable_cudagraph_enabled=breakable_cudagraph_enabled, + ): + raise ValueError( + "Adaptive verification requires piecewise CUDA graphs, but no " + "torch.compile or breakable CUDA graph provider is active. Enable " + "breakable CUDA graphs or disable adaptive verification." + ) + if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: # The draft budget divides by step costs profiled from captured # cudagraphs; eager execution captures none. From 9a5cec83fee92177b13d7f31c7e266f5bef63e06 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 3 Sep 2026 07:55:07 +0000 Subject: [PATCH 2/4] [Model][ROCm] Enable DeepSeek V4 Vision Enable the shared DeepSeek V4 vision wrapper on ROCm, preserve platform-specific text behavior, and add focused multimodal and attention coverage. Co-authored-by: OpenAI Codex Signed-off-by: Andreas Karatzas --- .../attention/test_rocm_triton_attn_dsv4.py | 69 ++++ .../processing/test_tensor_schema.py | 4 +- tests/models/test_deepseek_v4_vl_rocm.py | 228 ++++++++++++ tests/models/test_initialization.py | 12 +- tests/models/test_registry.py | 4 +- vllm/models/deepseek_v4/__init__.py | 8 +- vllm/models/deepseek_v4/amd/model.py | 24 +- vllm/models/deepseek_v4/amd/rocm.py | 56 ++- vllm/models/deepseek_v4/common/vl_model.py | 331 +++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 1 + vllm/models/deepseek_v4/nvidia/vl_model.py | 332 +----------------- vllm/models/deepseek_v4/vl_stub.py | 4 +- 12 files changed, 732 insertions(+), 341 deletions(-) create mode 100644 tests/models/test_deepseek_v4_vl_rocm.py create mode 100644 vllm/models/deepseek_v4/common/vl_model.py diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 2f1e8d279356..b2effcb29cdd 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -398,6 +398,75 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None: torch.testing.assert_close(actual_lens, expected_lens) +@torch.inference_mode() +def test_combine_topk_swa_indices_adds_image_visibility() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + num_tokens = 8 + topk_indices = torch.full((num_tokens, 1), -1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gather_lens = seq_lens.clone() + left_visible = torch.tensor( + [0, 0, 0, 1, 2, 3, 4, 0], dtype=torch.int32, device=device + ) + right_visible = torch.tensor( + [0, 0, 4, 3, 2, 1, 0, 0], dtype=torch.int32, device=device + ) + + indices, lens = combine_topk_swa_indices( + topk_indices, + query_start_loc, + seq_lens, + gather_lens, + window_size=4, + compress_ratio=1, + topk=0, + M=16, + N=0, + max_image_tokens=5, + left_visible=left_visible, + right_visible=right_visible, + ) + + expected_rows = [ + [0], + [0, 1], + [0, 1, 2, 3, 4, 5, 6], + [0, 1, 2, 3, 4, 5, 6], + [1, 2, 3, 4, 5, 6], + [2, 3, 4, 5, 6], + [2, 3, 4, 5, 6], + [4, 5, 6, 7], + ] + for token_idx, expected in enumerate(expected_rows): + actual = indices[token_idx, : lens[token_idx]].cpu().tolist() + assert actual == expected + + +@torch.inference_mode() +def test_combine_topk_swa_indices_keeps_vision_row_width_without_images() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + indices, lens = combine_topk_swa_indices( + torch.full((1, 1), -1, dtype=torch.int32, device=device), + torch.tensor([0, 1], dtype=torch.int32, device=device), + torch.tensor([1], dtype=torch.int32, device=device), + torch.tensor([1], dtype=torch.int32, device=device), + window_size=120, + compress_ratio=1, + topk=0, + M=256, + N=0, + max_image_tokens=16, + ) + + assert indices.shape == (1, 256) + assert lens.item() == 1 + + def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None: from vllm.models.deepseek_v4.amd import rocm as mod diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 8dc59fd4f5dc..5e507f595706 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -165,9 +165,9 @@ def test_model_tensor_schema(model_id: str): ) if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not ( - current_platform.is_cuda() + current_platform.is_cuda() or current_platform.is_rocm() ): - pytest.skip("Deepseek V4 is only supported on CUDA") + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) model_info.check_available_online(on_fail="skip") diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py new file mode 100644 index 000000000000..618cd5fd05f1 --- /dev/null +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.utils import WeightsMapper + +pytestmark = pytest.mark.cpu_test + + +def test_vl_mapper_preserves_rocm_weight_mapping() -> None: + from vllm.models.deepseek_v4.amd.model import _make_deepseek_v4_weights_mapper + from vllm.models.deepseek_v4.common.vl_model import ( + _make_deepseek_v4_vl_weights_mapper, + ) + + text_mapper = _make_deepseek_v4_weights_mapper("fp4", fuse_shared_experts=True) + mapper = _make_deepseek_v4_vl_weights_mapper(text_mapper, image_enabled=True) + + assert mapper._map_name("layers.3.attn.wq_a.input_scale") == ( + "language_model.model.layers.3.attn.wq_a.input_scale_2" + ) + assert mapper._map_name("layers.3.ffn.shared_experts.w2.weight") == ( + "language_model.model.layers.3.ffn.shared_experts.w2.weight" + ) + assert mapper._map_name("head.weight") == "language_model.lm_head.weight" + + +def test_rocm_moe_wires_vision_routing_on_hash_and_regular_layers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.deepseek_v4.amd import model as rocm_model + from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + + captured: list[dict] = [] + + class FakeGate(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() + + def fake_factory(**kwargs): + captured.append(kwargs) + return nn.Identity() + + monkeypatch.setattr(rocm_model, "GateLinear", FakeGate) + monkeypatch.setattr(rocm_model, "FusedMoEFactory", fake_factory) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0) + + config = SimpleNamespace( + hidden_size=16, + n_routed_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + swiglu_limit=None, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + num_hash_layers=1, + vocab_size=32, + topk_method="noaux_tc", + vision_n_layers=1, + n_shared_experts=None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), quant_config=None + ) + + hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn") + regular_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.1.ffn") + + assert hash_moe.gate.tid2eid is not None + assert regular_moe.gate.tid2eid is None + for moe, factory_kwargs in zip((hash_moe, regular_moe), captured, strict=True): + assert moe.gate.e_score_correction_bias is not None + assert moe.gate.bias_vl is not None + assert factory_kwargs["bias_vl"] is moe.gate.bias_vl + assert factory_kwargs["image_sentinel_lo"] == IMAGE_SENTINEL_BASE_ID + + +def test_rocm_compute_logits_local_skips_gather() -> None: + from vllm.models.deepseek_v4.amd.model import DeepseekV4ForCausalLM + + calls: list[tuple[nn.Module, torch.Tensor, bool]] = [] + + def logits_processor( + lm_head: nn.Module, hidden_states: torch.Tensor, *, skip_gather: bool = False + ) -> torch.Tensor: + calls.append((lm_head, hidden_states, skip_gather)) + return hidden_states + 1 + + model = object.__new__(DeepseekV4ForCausalLM) + nn.Module.__init__(model) + model.lm_head = nn.Identity() + model.logits_processor = logits_processor + hidden_states = torch.tensor([4.0]) + + result = model.compute_logits_local(hidden_states) + + assert torch.equal(result, torch.tensor([5.0])) + assert calls == [(model.lm_head, hidden_states, True)] + + +class _FakeLanguageModel(nn.Module): + finalizes_weights_during_load = False + + def __init__(self) -> None: + super().__init__() + self.tensor_a = nn.Parameter(torch.zeros(1)) + self.tensor_c = nn.Parameter(torch.zeros(1)) + self.finalized_values: list[tuple[float, float]] = [] + + def process_weights_after_loading(self) -> None: + self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item())) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + 1 + + +def test_vl_wrapper_streams_then_delegates_finalization() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeLanguageModel() + model.vision = nn.Module() + model.vision.tensor_b = nn.Parameter(torch.zeros(1)) + model.hf_to_vllm_mapper = WeightsMapper() + + def interleaved_weights(): + yield "language_model.tensor_a", torch.tensor([1.0]) + assert model.language_model.tensor_a.item() == 1.0 + yield "vision.tensor_b", torch.tensor([2.0]) + assert model.vision.tensor_b.item() == 2.0 + yield "language_model.tensor_c", torch.tensor([3.0]) + + loaded = model.load_weights(interleaved_weights()) + + assert loaded == { + "language_model.tensor_a", + "vision.tensor_b", + "language_model.tensor_c", + } + assert model.language_model.finalized_values == [] + + model.process_weights_after_loading() + + assert model.language_model.finalized_values == [(1.0, 3.0)] + assert torch.equal( + model.compute_logits_local(torch.tensor([4.0])), torch.tensor([5.0]) + ) + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(1.0, 3.0)] + + +class _FakeFinalizingLanguageModel(_FakeLanguageModel): + finalizes_weights_during_load = True + + def __init__(self) -> None: + super().__init__() + self.load_calls = 0 + + def load_weights(self, weights) -> set[str]: + self.load_calls += 1 + loaded = set() + for name, value in weights: + getattr(self, name).data.copy_(value) + loaded.add(name) + self.process_weights_after_loading() + return loaded + + +def test_vl_wrapper_groups_child_that_finalizes_during_load() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeFinalizingLanguageModel() + model.vision = nn.Module() + model.vision.tensor_b = nn.Parameter(torch.zeros(1)) + model.hf_to_vllm_mapper = WeightsMapper() + + loaded = model.load_weights( + iter( + ( + ("language_model.tensor_a", torch.tensor([1.0])), + ("vision.tensor_b", torch.tensor([2.0])), + ("language_model.tensor_c", torch.tensor([3.0])), + ) + ) + ) + + assert loaded == { + "language_model.tensor_a", + "vision.tensor_b", + "language_model.tensor_c", + } + assert model.language_model.load_calls == 1 + assert model.language_model.finalized_values == [(1.0, 3.0)] + + # The framework's later model-level hook must not double-finalize a child + # which already completed this work in load_weights. + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(1.0, 3.0)] + + +def test_vl_wrapper_dummy_load_delegates_finalization() -> None: + from vllm.models.deepseek_v4.common.vl_model import ( + DeepseekV4ForConditionalGeneration, + ) + + model = object.__new__(DeepseekV4ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = _FakeFinalizingLanguageModel() + + # DummyModelLoader bypasses model.load_weights(), so no finalized marker + # exists and the framework-level hook must still delegate to the child. + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(0.0, 0.0)] + model.process_weights_after_loading() + assert model.language_model.finalized_values == [(0.0, 0.0)] diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 3566eb534133..cdf2ea098961 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from functools import partial +from typing import Any from unittest.mock import patch import pytest @@ -135,8 +136,8 @@ def _initialize_kv_caches_v1(self, vllm_config): if model_arch == "DeepseekV4ForConditionalGeneration": from vllm.platforms import current_platform - if not current_platform.is_cuda(): - pytest.skip("Deepseek V4 is only supported on CUDA") + if not (current_platform.is_cuda() or current_platform.is_rocm()): + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") with ( patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1), @@ -160,7 +161,12 @@ def _initialize_kv_caches_v1(self, vllm_config): if model_arch == "WhisperForConditionalGeneration": m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - kwargs = {} + kwargs: dict[str, Any] = {} + if ( + model_arch == "DeepseekV4ForConditionalGeneration" + and current_platform.is_rocm() + ): + kwargs["kv_cache_dtype"] = "fp8" if not model_info.enable_prefix_caching: kwargs["enable_prefix_caching"] = False diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index d271cb677cab..10bf92b1700d 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -65,9 +65,9 @@ def test_registry_imports(model_arch): pytest.skip("HY V4 is only supported on CUDA") if model_arch == "DeepseekV4ForConditionalGeneration" and not ( - current_platform.is_cuda() + current_platform.is_cuda() or current_platform.is_rocm() ): - pytest.skip("Deepseek V4 is only supported on CUDA") + pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") # Ensure all model classes can be imported successfully model_cls = ModelRegistry._try_load_model_cls(model_arch) diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index d71bddae3328..386ed3c3fa2a 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -20,7 +20,7 @@ ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP - from .vl_stub import ( # type: ignore[assignment] + from .common.vl_model import ( # type: ignore[assignment] DeepseekV4ForConditionalGeneration, ) elif current_platform.is_xpu(): @@ -31,14 +31,14 @@ from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] else: + from .common.vl_model import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, ) from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment] - from .nvidia.vl_model import ( # type: ignore[assignment] - DeepseekV4ForConditionalGeneration, - ) __all__ = [ "DSparkDeepseekV4ForCausalLM", diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 7a67c8f420e6..1f6f832a06de 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -67,6 +67,8 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + logger = init_logger(__name__) @@ -204,6 +206,10 @@ def __init__( self.gate.e_score_correction_bias = None self.gate.tid2eid = None + self.gate.bias_vl = None + self.image_sentinel_lo = ( + IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0 + ) is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers self.hash_indices_dtype = torch.int32 if is_hash_moe: @@ -219,12 +225,20 @@ def __init__( ), requires_grad=False, ) - elif getattr(config, "topk_method", None) == "noaux_tc": + if getattr(config, "topk_method", None) == "noaux_tc" and ( + not is_hash_moe or getattr(config, "vision_n_layers", 0) > 0 + ): self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32), requires_grad=False, ) + if getattr(config, "vision_n_layers", 0) > 0: + self.gate.bias_vl = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + self.n_shared_experts = config.n_shared_experts # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` @@ -285,6 +299,8 @@ def __init__( routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, hash_indices_table=self.gate.tid2eid, + bias_vl=self.gate.bias_vl, + image_sentinel_lo=self.image_sentinel_lo, swiglu_limit=self.swiglu_limit, router_logits_dtype=torch.float32, ) @@ -294,6 +310,8 @@ def forward( ) -> torch.Tensor: if self.gate.tid2eid is not None and input_ids is None: raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + if self.gate.bias_vl is not None and input_ids is None: + raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") org_shape = hidden_states.shape final_hidden_states = self.experts( @@ -974,6 +992,7 @@ def _make_deepseek_v4_weights_mapper( class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): model_cls = DeepseekV4Model + finalizes_weights_during_load = False # Default mapper assumes the original FP4-expert checkpoint layout. # Overridden per-instance in __init__ when expert_dtype != "fp4". @@ -1020,6 +1039,9 @@ def compute_logits( logits = self.logits_processor(self.lm_head, hidden_states) return logits + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True) + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3c8d23c4b4fc..a2841b4f6795 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -106,13 +106,18 @@ def _combine_topk_swa_indices_kernel( query_start_loc_ptr, seq_lens_ptr, gather_lens_ptr, + left_visible_ptr, + right_visible_ptr, M, N, TOP_K: tl.constexpr, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, + SWA_WIDTH: tl.constexpr, TOPK_WIDTH: tl.constexpr, PADDED_TOP_K: tl.constexpr, + PADDED_SWA_WIDTH: tl.constexpr, + HAS_IMAGE: tl.constexpr, ): batch_idx = tl.program_id(0) worker_id = tl.program_id(1) @@ -131,7 +136,16 @@ def _combine_topk_swa_indices_kernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + if HAS_IMAGE: + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) + swa_end = pos + right + 1 + swa_len = swa_end - swa_start topk_offset = tl.arange(0, PADDED_TOP_K) topk_mask = topk_offset < topk_len @@ -149,14 +163,14 @@ def _combine_topk_swa_indices_kernel( mask=topk_mask, ) - swa_offset = tl.arange(0, WINDOW_SIZE) + swa_offset = tl.arange(0, PADDED_SWA_WIDTH) tl.store( combined_indices_ptr + token_idx * combined_indices_stride + topk_len + swa_offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_offset < swa_len, + M * batch_idx + N + swa_offset + swa_start - gather_start, + mask=(swa_offset < swa_len) & (swa_offset < SWA_WIDTH), ) tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) @@ -172,12 +186,22 @@ def combine_topk_swa_indices( topk: int, M: int, N: int, + max_image_tokens: int = 0, + left_visible: torch.Tensor | None = None, + right_visible: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + if (left_visible is None) != (right_visible is None): + raise ValueError("left_visible and right_visible must be provided together") topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() num_tokens = topk_indices.shape[0] num_reqs = seq_lens.shape[0] + has_image = left_visible is not None + # Keep the row shape fixed for a vision model even when a particular batch + # has no image. The gathered KV workspace needs no matching expansion: its + # query portion already contains each atomically-prefilled image span. + swa_width = window_size + max_image_tokens combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + (topk + swa_width + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) @@ -201,13 +225,18 @@ def combine_topk_swa_indices( query_start_loc, seq_lens, gather_lens, + left_visible if left_visible is not None else topk_indices, + right_visible if right_visible is not None else topk_indices, M, N, TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, + SWA_WIDTH=swa_width, TOPK_WIDTH=topk_indices.shape[-1], PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), + PADDED_SWA_WIDTH=triton.next_power_of_2(swa_width), + HAS_IMAGE=has_image, ) return combined_indices, combined_lens @@ -909,6 +938,12 @@ def _forward_prefill( assert query_start_loc_cpu is not None assert query_start_loc is not None prefill_token_base = query_start_loc_cpu[num_decodes] + left_visible = swa_metadata.prefill_left_visible + right_visible = swa_metadata.prefill_right_visible + if left_visible is not None: + left_visible = left_visible[num_decode_tokens:] + assert right_visible is not None + right_visible = right_visible[num_decode_tokens:] if not swa_only: if self.compress_ratio == 4: @@ -987,6 +1022,17 @@ def _forward_prefill( top_k, M, N, + max_image_tokens=self.max_image_tokens, + left_visible=( + left_visible[query_start:query_end] + if left_visible is not None + else None + ), + right_visible=( + right_visible[query_start:query_end] + if right_visible is not None + else None + ), ) rocm_sparse_attn_prefill( q=q[query_start:query_end], diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py new file mode 100644 index 000000000000..3470845300d0 --- /dev/null +++ b/vllm/models/deepseek_v4/common/vl_model.py @@ -0,0 +1,331 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). + +Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: + +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE + sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / + ``image_newline`` / ``image_end``) fill the remaining sentinel positions. +- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor + in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved + in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` + (see ``common/vision.py`` for the tower itself). +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before + its hyper-connection stream expansion. Raw ``input_ids`` still flow into + the model so the MoE router can apply ``bias_vl`` to image tokens + (``requires_raw_input_tokens``). +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .mm_preprocess import ( + IMAGE_PLACEHOLDER, + IMAGE_SENTINEL_BASE_ID, + DeepseekV4VLDummyInputsBuilder, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessingInfo, + image_sentinel_mask, +) +from .vision import DeepseekV4Aligner, DeepseekV4ViT + + +def _make_deepseek_v4_vl_weights_mapper( + text_mapper: WeightsMapper, image_enabled: bool +) -> WeightsMapper: + """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" + orig_to_new_prefix = { + src: None if dst is None else f"language_model.{dst}" + for src, dst in text_mapper.orig_to_new_prefix.items() + } + if not image_enabled: + orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) + return WeightsMapper( + orig_to_new_renaming=text_mapper.orig_to_new_renaming, + orig_to_new_prefix=orig_to_new_prefix, + orig_to_new_regex=text_mapper.orig_to_new_regex, + orig_to_new_stacked=text_mapper.orig_to_new_stacked, + orig_to_new_suffix={ + **text_mapper.orig_to_new_suffix, + "head.weight": "language_model.lm_head.weight", + }, + orig_to_new_substr={ + **text_mapper.orig_to_new_substr, + # The MTP/DSpark draft heads are not supported for the vision + # variant; drop their weights. + "mtp.": None, + }, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + DeepseekV4VLMultiModalProcessor, + info=DeepseekV4VLProcessingInfo, + dummy_inputs=DeepseekV4VLDummyInputsBuilder, +) +class DeepseekV4ForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. + + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) + delegates through ``language_model`` via the protocol defaults. + """ + + # The MoE router needs raw token ids to detect image sentinel tokens + # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. + requires_raw_input_tokens = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return IMAGE_PLACEHOLDER + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config, prefix: str = "") -> None: + super().__init__() + model_config = vllm_config.model_config + config = model_config.hf_config + self.config = config + self.multimodal_config = model_config.multimodal_config + assert self.multimodal_config is not None + + image_enabled = ( + config.vision_n_layers > 0 + and self.multimodal_config.get_limit_per_prompt("image") > 0 + ) + with self._mark_tower_model(vllm_config, {"image"}): + self.vision: DeepseekV4ViT | None = None + self.aligner: DeepseekV4Aligner | None = None + self.image_start: nn.Parameter | None = None + self.image_end: nn.Parameter | None = None + self.image_newline: nn.Parameter | None = None + self.image_pad: nn.Parameter | None = None + if image_enabled: + self.vision = DeepseekV4ViT(config) + self.aligner = DeepseekV4Aligner(config) + for name in ( + "image_start", + "image_end", + "image_newline", + "image_pad", + ): + setattr( + self, + name, + nn.Parameter( + torch.empty(config.hidden_size, dtype=torch.float32) + ), + ) + self.vision.to(dtype=model_config.dtype) + self.aligner.to(dtype=model_config.dtype) + + with self._mark_language_model(vllm_config): + # The arch convertor routes any config with a vision tower to + # this wrapper class; mark the copy handed to the text backbone + # so it resolves to DeepseekV4ForCausalLM instead of recursing + # (with_hf_config deepcopies the config, the marker survives). + config._dsv4_vl_inner = True # type: ignore[attr-defined] + try: + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["DeepseekV4ForCausalLM"], + ) + finally: + del config._dsv4_vl_inner # type: ignore[attr-defined] + # The outer mapper (see load_weights) fully resolves HF names into + # this wrapper's namespace before AutoWeightsLoader strips the + # "language_model." prefix and delegates to the child's load_weights, + # so the child's own mapper must be a no-op. Its suffix rules are not + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would + # re-fire "head.weight" -> "lm_head.weight"). + text_mapper = self.language_model.hf_to_vllm_mapper + self.language_model.hf_to_vllm_mapper = WeightsMapper() + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( + text_mapper, image_enabled + ) + self._weights_finalized = False + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + patches = kwargs.pop("patches", None) + if patches is None: + return None + vit_grid = kwargs.pop("vit_grid", None) + llm_grid = kwargs.pop("llm_grid", None) + perm = kwargs.pop("perm", None) + assert vit_grid is not None and llm_grid is not None and perm is not None + return { + "patches": patches, + "vit_grid": vit_grid, + "llm_grid": llm_grid, + "perm": perm, + } + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + perm: torch.Tensor, + ) -> torch.Tensor: + assert self.vision is not None and self.aligner is not None + image_embeds = self.aligner( + self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w + ) + # Reorder into the N-layout block order used in the prompt. + return image_embeds[perm.to(image_embeds.device)] + + def _process_image_input( + self, + patches: torch.Tensor, + vit_grid: torch.Tensor, + llm_grid: torch.Tensor, + perm: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + assert self.vision is not None and self.aligner is not None + patches = patches.to(self.aligner.w1.weight.dtype) + + embeds: list[torch.Tensor] = [] + vit_offset = 0 + llm_offset = 0 + for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( + vit_grid.tolist(), llm_grid.tolist(), strict=True + ): + n_vit = n_vit_h * n_vit_w + n_llm = n_llm_h * n_llm_w + embeds.append( + self._encode_image( + patches[vit_offset : vit_offset + n_vit], + n_vit_h, + n_vit_w, + perm[llm_offset : llm_offset + n_llm], + ) + ) + vit_offset += n_vit + llm_offset += n_llm + return tuple(embeds) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None or self.vision is None: + return [] + return self._process_image_input( + image_input["patches"], + image_input["vit_grid"], + image_input["llm_grid"], + image_input["perm"], + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # All ids are in-vocab here: image-block sentinels are borrowed + # reserved tokens (their embedding rows are always overwritten below). + inputs_embeds = self.language_model.embed_input_ids(input_ids) + + if self.image_start is not None: + # Branch-free sentinel overwrite: safe inside compiled/captured + # regions (no data-dependent control flow). + sentinel_mask = image_sentinel_mask(input_ids) + if is_multimodal is not None: + # IMAGE positions get vision embeddings via the merge below. + sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) + table = torch.stack( + [ + self.image_start, + self.image_pad, + self.image_pad, + self.image_newline, + self.image_end, + ] + ).to(inputs_embeds.dtype) + idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) + inputs_embeds = torch.where( + sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits_local(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" + return self.language_model.get_mtp_target_hidden_states() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + child_finalizes = getattr( + self.language_model, "finalizes_weights_during_load", True + ) + mapped = self.hf_to_vllm_mapper.apply(weights) + if child_finalizes: + # A child which finalizes inside load_weights must see all of its + # weights in one contiguous delegation from AutoWeightsLoader. + mapped = iter(sorted(mapped, key=lambda x: x[0])) + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(mapped) + self._weights_finalized = child_finalizes + return loaded_params + + def process_weights_after_loading(self) -> None: + # Backbones such as the ROCm implementation require this to run only + # after the loader's generic per-layer quantization finalization. + if getattr(self, "_weights_finalized", False): + return + self.language_model.process_weights_after_loading() + self._weights_finalized = True diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 3247877db071..0608d0c45c6a 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1744,6 +1744,7 @@ class DeepseekV4ForCausalLM( SupportsLoRA, DeepseekV4MixtureOfExperts, ): + finalizes_weights_during_load = True model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/models/deepseek_v4/nvidia/vl_model.py b/vllm/models/deepseek_v4/nvidia/vl_model.py index bd40b7bce9e7..a66217f2fe65 100644 --- a/vllm/models/deepseek_v4/nvidia/vl_model.py +++ b/vllm/models/deepseek_v4/nvidia/vl_model.py @@ -1,333 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). +"""Compatibility imports for the platform-neutral DeepSeek-V4 vision model.""" -Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: - -- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE - sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / - ``image_newline`` / ``image_end``) fill the remaining sentinel positions. -- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor - in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved - in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` - (see ``common/vision.py`` for the tower itself). -- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before - its hyper-connection stream expansion. Raw ``input_ids`` still flow into - the model so the MoE router can apply ``bias_vl`` to image tokens - (``requires_raw_input_tokens``). -""" - -from collections.abc import Iterable - -import torch -from torch import nn - -from vllm.model_executor.models.interfaces import ( - MultiModalEmbeddings, - SupportsEagle3, - SupportsMultiModal, - SupportsPP, +from ..common.vl_model import ( + DeepseekV4ForConditionalGeneration, ) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - WeightsMapper, - init_vllm_registered_model, - maybe_prefix, +from ..common.vl_model import ( + _make_deepseek_v4_vl_weights_mapper as _make_common_vl_weights_mapper, ) -from vllm.multimodal import MULTIMODAL_REGISTRY - -from ..common.mm_preprocess import ( - IMAGE_PLACEHOLDER, - IMAGE_SENTINEL_BASE_ID, - DeepseekV4VLDummyInputsBuilder, - DeepseekV4VLMultiModalProcessor, - DeepseekV4VLProcessingInfo, - image_sentinel_mask, -) -from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT from .model import _make_deepseek_v4_weights_mapper -def _make_deepseek_v4_vl_weights_mapper( - expert_dtype: str, image_enabled: bool -) -> WeightsMapper: - """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" - base = _make_deepseek_v4_weights_mapper(expert_dtype) - orig_to_new_prefix: dict[str, str | None] = { - "layers.": "language_model.model.layers.", - "embed.": "language_model.model.embed.", - "norm.": "language_model.model.norm.", - "hc_head": "language_model.model.hc_head", - "mtp.": "language_model.model.mtp.", - } - if not image_enabled: - orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) - return WeightsMapper( - orig_to_new_prefix=orig_to_new_prefix, - orig_to_new_regex=base.orig_to_new_regex, - orig_to_new_suffix={ - "head.weight": "language_model.lm_head.weight", - "embed.weight": "embed_tokens.weight", - ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", - }, - orig_to_new_substr={ - ".shared_experts.w2": ".shared_experts.down_proj", - # The MTP/DSpark draft heads are not supported for the vision - # variant; drop their weights. - "mtp.": None, - }, +def _make_deepseek_v4_vl_weights_mapper(expert_dtype: str, image_enabled: bool): + """Retain the original NVIDIA helper signature for downstream imports.""" + return _make_common_vl_weights_mapper( + _make_deepseek_v4_weights_mapper(expert_dtype), image_enabled ) -@MULTIMODAL_REGISTRY.register_processor( - DeepseekV4VLMultiModalProcessor, - info=DeepseekV4VLProcessingInfo, - dummy_inputs=DeepseekV4VLDummyInputsBuilder, -) -class DeepseekV4ForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 -): - """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. - - ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) - delegates through ``language_model`` via the protocol defaults. - """ - - # The MoE router needs raw token ids to detect image sentinel tokens - # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. - requires_raw_input_tokens = True - - @classmethod - def get_placeholder_str(cls, modality: str, i: int) -> str | None: - if modality == "image": - return IMAGE_PLACEHOLDER - raise ValueError(f"Unsupported modality: {modality!r}") - - def __init__(self, *, vllm_config, prefix: str = "") -> None: - super().__init__() - model_config = vllm_config.model_config - config = model_config.hf_config - self.config = config - self.multimodal_config = model_config.multimodal_config - assert self.multimodal_config is not None - - image_enabled = ( - config.vision_n_layers > 0 - and self.multimodal_config.get_limit_per_prompt("image") > 0 - ) - with self._mark_tower_model(vllm_config, {"image"}): - self.vision: DeepseekV4ViT | None = None - self.aligner: DeepseekV4Aligner | None = None - self.image_start: nn.Parameter | None = None - self.image_end: nn.Parameter | None = None - self.image_newline: nn.Parameter | None = None - self.image_pad: nn.Parameter | None = None - if image_enabled: - self.vision = DeepseekV4ViT(config) - self.aligner = DeepseekV4Aligner(config) - for name in ( - "image_start", - "image_end", - "image_newline", - "image_pad", - ): - setattr( - self, - name, - nn.Parameter( - torch.empty(config.hidden_size, dtype=torch.float32) - ), - ) - self.vision.to(dtype=model_config.dtype) - self.aligner.to(dtype=model_config.dtype) - - with self._mark_language_model(vllm_config): - # The arch convertor routes any config with a vision tower to - # this wrapper class; mark the copy handed to the text backbone - # so it resolves to DeepseekV4ForCausalLM instead of recursing - # (with_hf_config deepcopies the config, the marker survives). - config._dsv4_vl_inner = True # type: ignore[attr-defined] - try: - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - hf_config=config, - prefix=maybe_prefix(prefix, "language_model"), - architectures=["DeepseekV4ForCausalLM"], - ) - finally: - del config._dsv4_vl_inner # type: ignore[attr-defined] - # The outer mapper (see load_weights) fully resolves HF names into - # this wrapper's namespace before AutoWeightsLoader strips the - # "language_model." prefix and delegates to the child's load_weights, - # so the child's own mapper must be a no-op. Its suffix rules are not - # idempotent (e.g. "lm_head.weight".endswith("head.weight") would - # re-fire "head.weight" -> "lm_head.weight"). - self.language_model.hf_to_vllm_mapper = WeightsMapper() - self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] - self.language_model.make_empty_intermediate_tensors - ) - - expert_dtype = getattr(config, "expert_dtype", "fp4") - self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( - expert_dtype, image_enabled - ) - - def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: - patches = kwargs.pop("patches", None) - if patches is None: - return None - vit_grid = kwargs.pop("vit_grid", None) - llm_grid = kwargs.pop("llm_grid", None) - perm = kwargs.pop("perm", None) - assert vit_grid is not None and llm_grid is not None and perm is not None - return { - "patches": patches, - "vit_grid": vit_grid, - "llm_grid": llm_grid, - "perm": perm, - } - - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - perm: torch.Tensor, - ) -> torch.Tensor: - assert self.vision is not None and self.aligner is not None - image_embeds = self.aligner( - self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w - ) - # Reorder into the N-layout block order used in the prompt. - return image_embeds[perm.to(image_embeds.device)] - - def _process_image_input( - self, - patches: torch.Tensor, - vit_grid: torch.Tensor, - llm_grid: torch.Tensor, - perm: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - assert self.vision is not None and self.aligner is not None - patches = patches.to(self.aligner.w1.weight.dtype) - - embeds: list[torch.Tensor] = [] - vit_offset = 0 - llm_offset = 0 - for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( - vit_grid.tolist(), llm_grid.tolist(), strict=True - ): - n_vit = n_vit_h * n_vit_w - n_llm = n_llm_h * n_llm_w - embeds.append( - self._encode_image( - patches[vit_offset : vit_offset + n_vit], - n_vit_h, - n_vit_w, - perm[llm_offset : llm_offset + n_llm], - ) - ) - vit_offset += n_vit - llm_offset += n_llm - return tuple(embeds) - - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - image_input = self._parse_and_validate_image_input(**kwargs) - if image_input is None or self.vision is None: - return [] - return self._process_image_input( - image_input["patches"], - image_input["vit_grid"], - image_input["llm_grid"], - image_input["perm"], - ) - - def embed_input_ids( - self, - input_ids: torch.Tensor, - multimodal_embeddings: MultiModalEmbeddings | None = None, - *, - is_multimodal: torch.Tensor | None = None, - ) -> torch.Tensor: - from vllm.model_executor.models.utils import _merge_multimodal_embeddings - - # All ids are in-vocab here: image-block sentinels are borrowed - # reserved tokens (their embedding rows are always overwritten below). - inputs_embeds = self.language_model.embed_input_ids(input_ids) - - if self.image_start is not None: - # Branch-free sentinel overwrite: safe inside compiled/captured - # regions (no data-dependent control flow). - sentinel_mask = image_sentinel_mask(input_ids) - if is_multimodal is not None: - # IMAGE positions get vision embeddings via the merge below. - sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) - table = torch.stack( - [ - self.image_start, - self.image_pad, - self.image_pad, - self.image_newline, - self.image_end, - ] - ).to(inputs_embeds.dtype) - idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) - inputs_embeds = torch.where( - sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds - ) - - if multimodal_embeddings is None or len(multimodal_embeddings) == 0: - return inputs_embeds - - assert is_multimodal is not None - return _merge_multimodal_embeddings( - inputs_embeds=inputs_embeds, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors=None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - return self.language_model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: - return self.language_model.compute_logits(hidden_states) - - def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.language_model.compute_logits_local(hidden_states) - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - return self.language_model.get_expert_mapping() - - def get_mtp_target_hidden_states(self) -> torch.Tensor | None: - """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" - return self.language_model.get_mtp_target_hidden_states() - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Map HF names into this wrapper's namespace up front and sort, so - # the "language_model." group reaches the child loader as one - # contiguous block (AutoWeightsLoader delegates per contiguous group, - # and the child's load_weights finalizes fused expert weights, which - # must not run on a partially loaded model). - mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0]) - loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(mapped) - # The child's load_weights already ran its post-load finalization. - self._weights_finalized = True - return loaded_params - - 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 - self.language_model.process_weights_after_loading() +__all__ = ["DeepseekV4ForConditionalGeneration"] diff --git a/vllm/models/deepseek_v4/vl_stub.py b/vllm/models/deepseek_v4/vl_stub.py index 0f15a82f0190..4ade7b17d387 100644 --- a/vllm/models/deepseek_v4/vl_stub.py +++ b/vllm/models/deepseek_v4/vl_stub.py @@ -9,6 +9,6 @@ class DeepseekV4ForConditionalGeneration(nn.Module): def __init__(self, *, vllm_config, prefix: str = ""): super().__init__() raise NotImplementedError( - "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only " - "supported on NVIDIA GPUs for now." + "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is not " + "supported on the current platform; use CUDA or ROCm." ) From b751ff84e3c109df9596c4b4a7d08fd315c0a343 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 4 Sep 2026 04:19:59 +0000 Subject: [PATCH 3/4] [Bugfix] Fix noncompiled cudagraph fallback edge cases Preserve default compilation dispatch, accept graph-safe auxiliary configs, and keep adaptive verification on a full decode graph when piecewise capture is unavailable. Signed-off-by: Andreas Karatzas --- tests/test_config.py | 64 +++++++++---------- .../spec_decode/test_adaptive_verification.py | 22 +++++++ vllm/config/vllm.py | 38 ++++------- vllm/v1/worker/gpu/model_runner.py | 10 ++- .../gpu/spec_decode/adaptive_verification.py | 11 ++++ 5 files changed, 85 insertions(+), 60 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 50c4fb5fa279..77bc6749a6d9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -426,7 +426,7 @@ def test_noncompiled_architectures_fall_back_when_breakable_disabled( config = make_cudagraph_config(architecture, rocm=False) assert not VllmConfig._maybe_enable_breakable_cudagraph(config) - assert config.compilation_config.mode == CompilationMode.NONE + assert config.compilation_config.mode is None assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY @@ -447,7 +447,7 @@ def test_noncompiled_cudagraph_fallback_respects_optimization_level( ) assert not VllmConfig._maybe_enable_breakable_cudagraph(config) - assert config.compilation_config.mode == CompilationMode.NONE + assert config.compilation_config.mode is None assert config.compilation_config.cudagraph_mode == expected_cudagraph_mode @@ -466,50 +466,47 @@ def test_rocm_forced_v2_only_falls_back_for_noncompiled_dsa_model( config = make_cudagraph_config(architecture) assert not VllmConfig._maybe_enable_breakable_cudagraph(config) - expected_compilation_mode = ( - CompilationMode.NONE if expected_cudagraph_mode is not None else None - ) - assert config.compilation_config.mode == expected_compilation_mode + assert config.compilation_config.mode is None assert config.compilation_config.cudagraph_mode == expected_cudagraph_mode @pytest.mark.parametrize( ("architecture", "use_v2", "input_mode", "expected_compile", "expected_graph"), [ - ("DeepseekV4ForCausalLM", True, None, CompilationMode.NONE, CUDAGraphMode.NONE), + ("DeepseekV4ForCausalLM", True, None, None, CUDAGraphMode.NONE), ( "DeepseekV4ForConditionalGeneration", True, None, - CompilationMode.NONE, + None, CUDAGraphMode.NONE, ), ( "DeepseekV4ForConditionalGeneration", True, CUDAGraphMode.FULL, - CompilationMode.NONE, + None, CUDAGraphMode.FULL, ), ( "DeepseekV4ForConditionalGeneration", True, CUDAGraphMode.FULL_DECODE_ONLY, - CompilationMode.NONE, + None, CUDAGraphMode.FULL_DECODE_ONLY, ), ( "DeepseekV4ForConditionalGeneration", True, CUDAGraphMode.FULL_AND_PIECEWISE, - CompilationMode.NONE, + None, CUDAGraphMode.FULL_DECODE_ONLY, ), ( "KimiK3ForConditionalGeneration", True, None, - CompilationMode.NONE, + None, CUDAGraphMode.FULL_DECODE_ONLY, ), ], @@ -555,20 +552,26 @@ def test_rocm_gfx950_noncompiled_cudagraph_policy( None, True, ), - ( + pytest.param( CompilationMode.VLLM_COMPILE, CUDAGraphMode.NONE, - None, - True, + CUDAGraphMode.NONE, + False, + id="ngram-gpu-auxiliary-config", ), ( CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL_DECODE_ONLY, - None, - True, + CUDAGraphMode.FULL_DECODE_ONLY, + False, ), (CompilationMode.VLLM_COMPILE, CUDAGraphMode.PIECEWISE, None, True), - (CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL, None, True), + ( + CompilationMode.VLLM_COMPILE, + CUDAGraphMode.FULL, + CUDAGraphMode.FULL, + False, + ), ( CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL_AND_PIECEWISE, @@ -598,7 +601,7 @@ def test_noncompiled_cudagraph_fallback_validates_explicit_modes( ) assert not breakable_enabled - assert config.compilation_config.mode == (compile_mode or CompilationMode.NONE) + assert config.compilation_config.mode == compile_mode assert config.compilation_config.cudagraph_mode == expected_graph @@ -712,31 +715,28 @@ def test_late_piecewise_override_is_normalized_by_available_provider( @pytest.mark.parametrize( - ("architecture", "mode", "breakable_enabled", "should_raise"), + ("mode", "graph_mode", "should_raise"), [ - ("DeepseekV4ForConditionalGeneration", CompilationMode.NONE, False, True), - ("DeepseekV4ForConditionalGeneration", CompilationMode.NONE, True, False), - ("LlamaForCausalLM", CompilationMode.VLLM_COMPILE, False, False), + (CompilationMode.NONE, CUDAGraphMode.FULL_DECODE_ONLY, False), + (CompilationMode.VLLM_COMPILE, CUDAGraphMode.FULL_AND_PIECEWISE, False), + (CompilationMode.NONE, CUDAGraphMode.NONE, True), + (CompilationMode.VLLM_COMPILE, CUDAGraphMode.PIECEWISE, True), ], ) -def test_adaptive_verification_requires_piecewise_cudagraph_provider( - make_cudagraph_config, architecture, mode, breakable_enabled, should_raise +def test_adaptive_verification_requires_full_cudagraphs( + make_cudagraph_config, mode, graph_mode, should_raise ): config = make_cudagraph_config( - architecture, - breakable=breakable_enabled, compilation_mode=mode, - cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + cudagraph_mode=graph_mode, ) config.speculative_config = SimpleNamespace(enable_adaptive_verification=True) config.lora_config = None config.parallel_config = SimpleNamespace(pipeline_parallel_size=1) - validate = lambda: VllmConfig._validate_adaptive_verification( - config, breakable_cudagraph_enabled=breakable_enabled - ) + validate = lambda: VllmConfig._validate_adaptive_verification(config) if should_raise: - with pytest.raises(ValueError, match="requires piecewise CUDA graphs"): + with pytest.raises(ValueError, match="requires full CUDA graphs"): validate() else: validate() diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index d92f78a71535..ff9db905ef49 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -4,7 +4,9 @@ from types import SimpleNamespace import numpy as np +import pytest +from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.worker.gpu.async_utils import StepTimingSample from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo @@ -12,6 +14,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, maybe_create_adaptive_verification_manager, + resolve_adaptive_cudagraph_mode, ) from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping @@ -35,6 +38,25 @@ def make_manager( return manager +@pytest.mark.parametrize( + ("mode", "piecewise_capture_available", "expected"), + [ + (CUDAGraphMode.FULL_DECODE_ONLY, True, CUDAGraphMode.FULL_DECODE_ONLY), + (CUDAGraphMode.FULL_DECODE_ONLY, False, CUDAGraphMode.FULL_DECODE_ONLY), + (CUDAGraphMode.FULL, True, CUDAGraphMode.FULL_AND_PIECEWISE), + (CUDAGraphMode.FULL, False, CUDAGraphMode.FULL_DECODE_ONLY), + (CUDAGraphMode.FULL_AND_PIECEWISE, False, CUDAGraphMode.FULL_DECODE_ONLY), + ], +) +def test_resolve_adaptive_cudagraph_mode(mode, piecewise_capture_available, expected): + assert ( + resolve_adaptive_cudagraph_mode( + mode, piecewise_capture_available=piecewise_capture_available + ) + == expected + ) + + def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatch): class Backend: @classmethod diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ecb98b462789..73022ce085d2 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -784,7 +784,11 @@ def _maybe_enable_breakable_cudagraph(self) -> bool: # eagerly. This also handles platforms such as ROCm that deliberately # leave breakable graphs disabled by default for performance. compilation_config = self.compilation_config - if compilation_config.mode == CompilationMode.VLLM_COMPILE: + cudagraph_mode = compilation_config.cudagraph_mode + if compilation_config.mode == CompilationMode.VLLM_COMPILE and ( + cudagraph_mode is None + or cudagraph_mode.requires_piecewise_compilation() + ): raise ValueError( f"{self.model_config.architecture} does not expose an active " "torch.compile boundary, so compilation mode VLLM_COMPILE " @@ -819,11 +823,6 @@ def _maybe_enable_breakable_cudagraph(self) -> bool: compilation_config.cudagraph_mode.name, ) - # None is the unset value, not an explicit compile request. These - # wrappers cannot provide piecewise compilation, so resolve it now - # and let the compatibility pass below remove impossible modes. - if compilation_config.mode is None: - compilation_config.mode = CompilationMode.NONE return enabled @property @@ -1797,9 +1796,7 @@ def has_blocked_weights(): self._validate_v1_model_runner() self._validate_batch_sharded_sampling() - self._validate_adaptive_verification( - breakable_cudagraph_enabled=breakable_cudagraph_enabled - ) + self._validate_adaptive_verification() # Re-compute compile ranges after platform-specific config updates # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled) @@ -2786,9 +2783,7 @@ def _get_v1_model_runner_unsupported_features(self) -> list[str]: return unsupported - def _validate_adaptive_verification( - self, *, breakable_cudagraph_enabled: bool - ) -> None: + def _validate_adaptive_verification(self) -> None: spec_config = self.speculative_config if not spec_config or not spec_config.enable_adaptive_verification: return @@ -2800,22 +2795,11 @@ def _validate_adaptive_verification( "Adaptive verification is not currently compatible with LoRA" ) - if not VllmConfig._piecewise_cudagraph_provider_available( - self, - breakable_cudagraph_enabled=breakable_cudagraph_enabled, - ): - raise ValueError( - "Adaptive verification requires piecewise CUDA graphs, but no " - "torch.compile or breakable CUDA graph provider is active. Enable " - "breakable CUDA graphs or disable adaptive verification." - ) - - if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: - # The draft budget divides by step costs profiled from captured - # cudagraphs; eager execution captures none. + if not self.compilation_config.cudagraph_mode.has_full_cudagraphs(): raise ValueError( - "Adaptive verification is not currently compatible with " - "enforce_eager/cudagraph_mode=none" + "Adaptive verification requires full CUDA graphs. Use cudagraph " + "mode FULL, FULL_DECODE_ONLY, or FULL_AND_PIECEWISE, or disable " + "adaptive verification." ) if self.parallel_config.pipeline_parallel_size > 1: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 24e38d47a74b..0e89d1dd2351 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -100,6 +100,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, ModelCudaGraphManager, + has_compiled_submodule, ) from vllm.v1.worker.gpu.cudagraph_utils import ( profile_cudagraph_memory as _profile_cudagraph_memory, @@ -147,6 +148,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, maybe_create_adaptive_verification_manager, + resolve_adaptive_cudagraph_mode, ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, @@ -638,7 +640,13 @@ def initialize_kv_cache( use_replayssm=self.vllm_config.cache_config.use_replayssm, ) if self.adaptive_verification is not None: - self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode( + self.compilation_config.cudagraph_mode, + piecewise_capture_available=( + envs.VLLM_USE_BREAKABLE_CUDAGRAPH + or has_compiled_submodule(self.model) + ), + ) cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index d3b3743d67fd..af58b8993bfa 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -10,6 +10,7 @@ import torch import vllm.envs as envs +from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -33,6 +34,16 @@ from vllm.v1.worker.utils import AttentionGroup +def resolve_adaptive_cudagraph_mode( + mode: CUDAGraphMode, *, piecewise_capture_available: bool +) -> CUDAGraphMode: + """Select a separate decode route for variable-length verification.""" + assert mode.has_full_cudagraphs() + if mode == CUDAGraphMode.FULL_DECODE_ONLY or not piecewise_capture_available: + return CUDAGraphMode.FULL_DECODE_ONLY + return CUDAGraphMode.FULL_AND_PIECEWISE + + def _assign_draft_token_budget( confidence_probs: torch.Tensor, idx_mapping: torch.Tensor, From dba58f44dc04a2a2b7722b7308a297e0f99797cf Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 4 Sep 2026 07:25:58 +0000 Subject: [PATCH 4/4] Address DeepSeek V4 ROCm review feedback Signed-off-by: Andreas Karatzas --- .buildkite/test_areas/models_basic.yaml | 3 +- .../attention/test_rocm_triton_attn_dsv4.py | 26 +++++ .../processing/test_tensor_schema.py | 5 +- tests/models/test_deepseek_v4_vl_rocm.py | 105 +++++++++++++++++- tests/models/test_initialization.py | 8 +- tests/models/test_registry.py | 5 +- tests/test_config.py | 39 +++++++ vllm/models/deepseek_v4/amd/mtp.py | 2 +- vllm/models/deepseek_v4/amd/rocm.py | 13 ++- vllm/models/deepseek_v4/attention.py | 13 ++- vllm/models/deepseek_v4/common/vl_model.py | 4 +- vllm/platforms/rocm.py | 15 +++ 12 files changed, 213 insertions(+), 25 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 6d4b948e8708..abdfd7a60ed5 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -48,8 +48,9 @@ steps: - tests/models/test_terratorch.py - tests/models/transformers/test_backend.py - tests/models/test_registry.py + - tests/models/test_deepseek_v4_vl_rocm.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py mirror: amd: label: ":amd: (MI300) Basic Models (Other)" diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index b2effcb29cdd..f66dbdfd1ce6 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -445,6 +445,32 @@ def test_combine_topk_swa_indices_adds_image_visibility() -> None: assert actual == expected +@torch.inference_mode() +def test_combine_topk_swa_indices_apc_hit_inside_image() -> None: + from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices + + device = torch.device("cuda") + indices, lens = combine_topk_swa_indices( + torch.full((2, 1), -1, dtype=torch.int32, device=device), + torch.tensor([0, 2], dtype=torch.int32, device=device), + torch.tensor([10], dtype=torch.int32, device=device), + # Only positions [5, 10) exist in the gathered SWA workspace. + torch.tensor([5], dtype=torch.int32, device=device), + window_size=4, + compress_ratio=1, + topk=0, + M=10, + N=0, + max_image_tokens=10, + left_visible=torch.tensor([8, 9], dtype=torch.int32, device=device), + # Deliberately extends beyond seq_len to exercise the upper clamp too. + right_visible=torch.tensor([5, 5], dtype=torch.int32, device=device), + ) + + assert lens.cpu().tolist() == [5, 5] + assert indices[:, :5].cpu().tolist() == [list(range(5)), list(range(5))] + + @torch.inference_mode() def test_combine_topk_swa_indices_keeps_vision_row_width_without_images() -> None: from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 5e507f595706..98c5b1dd30ab 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -164,8 +164,9 @@ def test_model_tensor_schema(model_id: str): "Kimi-K2.5's offline inference has issues about vision chunks. Fix later." ) - if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not ( - current_platform.is_cuda() or current_platform.is_rocm() + if ( + model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" + and not current_platform.is_cuda_alike() ): pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 618cd5fd05f1..b9ea654e5b31 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -8,8 +8,42 @@ from torch import nn from vllm.model_executor.models.utils import WeightsMapper +from vllm.platforms import current_platform -pytestmark = pytest.mark.cpu_test +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + + +def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None: + from vllm.config import CacheConfig + from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype + + cache_config = CacheConfig() + + resolved_dtype, torch_dtype = _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout=True, + kv_cache_dtype=cache_config.cache_dtype, + cache_config=cache_config, + ) + + assert resolved_dtype == "fp8_ds_mla" + assert torch_dtype is torch.uint8 + assert cache_config.cache_dtype == "fp8_ds_mla" + + +def test_rocm_packed_kv_cache_rejects_unquantized_dtype() -> None: + from vllm.config import CacheConfig + from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype + + cache_config = CacheConfig(cache_dtype="bfloat16") + + with pytest.raises(ValueError, match="only supports fp8 kv-cache"): + _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout=True, + kv_cache_dtype=cache_config.cache_dtype, + cache_config=cache_config, + ) def test_vl_mapper_preserves_rocm_weight_mapping() -> None: @@ -81,6 +115,75 @@ def fake_factory(**kwargs): assert factory_kwargs["image_sentinel_lo"] == IMAGE_SENTINEL_BASE_ID +def test_rocm_mtp_forwards_input_ids_for_vision_routing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.deepseek_v4.amd import mtp as rocm_mtp + + hidden_size = 4 + hc_mult = 2 + + class FakeNorm(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = 1e-6 + + class FakeMTPBlock(nn.Module): + use_fused_mhc = False + + def __init__(self) -> None: + super().__init__() + self.seen_input_ids: torch.Tensor | None = None + + def forward( + self, + *, + positions: torch.Tensor, + x: torch.Tensor, + input_ids: torch.Tensor | None, + ): + self.seen_input_ids = input_ids + return x, None, None, None + + def passthrough_mtp_input( + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + *args, + ) -> tuple[torch.Tensor, torch.Tensor]: + return inputs_embeds, previous_hidden_states + + monkeypatch.setattr(rocm_mtp, "fused_mtp_input_rmsnorm", passthrough_mtp_input) + + layer = object.__new__(rocm_mtp.DeepSeekV4MultiTokenPredictorLayer) + nn.Module.__init__(layer) + layer.config = SimpleNamespace(hidden_size=hidden_size) + layer.hc_mult = hc_mult + layer.enorm = FakeNorm() + layer.hnorm = FakeNorm() + layer.e_proj = nn.Identity() + layer.h_proj = nn.Identity() + layer.mtp_block = FakeMTPBlock() + + input_ids = torch.tensor([11, 12]) + positions = torch.tensor([3, 4]) + inputs_embeds = torch.arange(8, dtype=torch.float32).view(2, hidden_size) + previous_hidden_states = torch.arange(16, dtype=torch.float32).view(2, -1) + + output = layer( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + ) + + assert layer.mtp_block.seen_input_ids is input_ids + expected = previous_hidden_states.view(2, hc_mult, hidden_size) + expected = expected + inputs_embeds.unsqueeze(-2) + torch.testing.assert_close(output, expected.flatten(1)) + + def test_rocm_compute_logits_local_skips_gather() -> None: from vllm.models.deepseek_v4.amd.model import DeepseekV4ForCausalLM diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index cdf2ea098961..d7fe12d21d01 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from functools import partial -from typing import Any from unittest.mock import patch import pytest @@ -161,12 +160,7 @@ def _initialize_kv_caches_v1(self, vllm_config): if model_arch == "WhisperForConditionalGeneration": m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - kwargs: dict[str, Any] = {} - if ( - model_arch == "DeepseekV4ForConditionalGeneration" - and current_platform.is_rocm() - ): - kwargs["kv_cache_dtype"] = "fp8" + kwargs = {} if not model_info.enable_prefix_caching: kwargs["enable_prefix_caching"] = False diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 10bf92b1700d..7d57918a8e26 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -64,8 +64,9 @@ def test_registry_imports(model_arch): ): pytest.skip("HY V4 is only supported on CUDA") - if model_arch == "DeepseekV4ForConditionalGeneration" and not ( - current_platform.is_cuda() or current_platform.is_rocm() + if ( + model_arch == "DeepseekV4ForConditionalGeneration" + and not current_platform.is_cuda_alike() ): pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") diff --git a/tests/test_config.py b/tests/test_config.py index 50c4fb5fa279..7f53d7dd2379 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -55,6 +55,45 @@ def _write_json(path: Path, value: object) -> None: path.write_text(json.dumps(value), encoding="utf-8") +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") +@pytest.mark.parametrize( + ("is_mm_prefix_lm", "is_multimodal_model", "expected"), + [ + pytest.param(True, True, True, id="multimodal-prefix-lm"), + pytest.param(False, True, False, id="multimodal-causal"), + pytest.param(True, False, False, id="text-prefix-lm"), + pytest.param(None, True, False, id="missing-model-config"), + ], +) +def test_rocm_mm_prefix_lm_disables_chunked_mm_input( + is_mm_prefix_lm: bool | None, + is_multimodal_model: bool, + expected: bool, +) -> None: + from vllm.platforms.rocm import RocmPlatform + + config = SimpleNamespace( + compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + worker_cls="test-worker", + ), + model_config=( + None + if is_mm_prefix_lm is None + else SimpleNamespace(is_mm_prefix_lm=is_mm_prefix_lm) + ), + scheduler_config=SimpleNamespace( + is_multimodal_model=is_multimodal_model, + disable_chunked_mm_input=False, + ), + ) + + RocmPlatform.check_and_update_config(config) + + assert config.scheduler_config.disable_chunked_mm_input is expected + + def test_kda_recoverssm_derivation_is_revalidated(): config = SimpleNamespace( cache_config=SimpleNamespace( diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 51479af238b5..9fe106cf9f3c 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -156,7 +156,7 @@ def forward( inputs_embeds ).unsqueeze(-2) hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=None + positions=positions, x=hidden_states, input_ids=input_ids ) if self.mtp_block.use_fused_mhc: hidden_states = self.mtp_block.hc_post( diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index a2841b4f6795..00f500944296 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -143,9 +143,13 @@ def _combine_topk_swa_indices_kernel( left = 0 right = 0 left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) - swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) - swa_end = pos + right + 1 - swa_len = swa_end - swa_start + # Prefix caching can resume inside an image span. Do not generate + # indices outside the SWA rows present in the gathered workspace. + swa_start = tl.maximum( + tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0), gather_start + ) + swa_end = tl.minimum(pos + right + 1, seq_len) + swa_len = tl.maximum(swa_end - swa_start, 0) topk_offset = tl.arange(0, PADDED_TOP_K) topk_mask = topk_offset < topk_len @@ -197,8 +201,7 @@ def combine_topk_swa_indices( num_reqs = seq_lens.shape[0] has_image = left_visible is not None # Keep the row shape fixed for a vision model even when a particular batch - # has no image. The gathered KV workspace needs no matching expansion: its - # query portion already contains each atomically-prefilled image span. + # has no image. swa_width = window_size + max_image_tokens combined_topk = ( (topk + swa_width + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index ce9fe06a2fb8..d62e4d5bf5f0 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -102,10 +102,15 @@ def _resolve_dsv4_kv_cache_dtype( """ if use_fp8_ds_mla_layout: # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. - assert kv_cache_dtype.startswith("fp8"), ( - f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, " - f"got {kv_cache_dtype}" - ) + if kv_cache_dtype == "auto": + kv_cache_dtype = "fp8" + if not kv_cache_dtype.startswith("fp8"): + raise ValueError( + "DeepseekV4 fp8_ds_mla layout only supports fp8 " + f"kv-cache, got {kv_cache_dtype}. Please set " + "`--kv-cache-dtype fp8` or select a backend that supports " + "bfloat16 KV cache." + ) if kv_cache_dtype != "fp8_ds_mla": if cache_config is not None: cache_config.cache_dtype = "fp8_ds_mla" diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py index 3470845300d0..4cccb85a8302 100644 --- a/vllm/models/deepseek_v4/common/vl_model.py +++ b/vllm/models/deepseek_v4/common/vl_model.py @@ -68,8 +68,8 @@ def _make_deepseek_v4_vl_weights_mapper( }, orig_to_new_substr={ **text_mapper.orig_to_new_substr, - # The MTP/DSpark draft heads are not supported for the vision - # variant; drop their weights. + # Draft models load the checkpoint's MTP weights separately from + # the target model. "mtp.": None, }, ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 280dec2d3181..7803ef55b22c 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -915,6 +915,21 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: if parallel_config.worker_cls == "auto": parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker" + model_config = vllm_config.model_config + scheduler_config = vllm_config.scheduler_config + # Note: model_config may be None during testing + if ( + model_config is not None + and model_config.is_mm_prefix_lm + and scheduler_config.is_multimodal_model + and not scheduler_config.disable_chunked_mm_input + ): + logger.warning_once( + "Forcing --disable_chunked_mm_input for models " + "with multimodal-bidirectional attention." + ) + scheduler_config.disable_chunked_mm_input = True + @classmethod def verify_model_arch(cls, model_arch: str) -> None: if model_arch in _ROCM_UNSUPPORTED_MODELS: