diff --git a/tests/test_patches.py b/tests/test_patches.py index d84a077e43c1..3d58decb92b7 100644 --- a/tests/test_patches.py +++ b/tests/test_patches.py @@ -479,6 +479,227 @@ def test_tp2_keeps_musa_cudagraph_capture(self): assert vllm_config.compilation_config.max_cudagraph_capture_size == 512 assert vllm_config.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8] + def _make_hybrid_spec_config( + self, + *, + method="qwen3_5_mtp", + num_speculative_tokens=2, + is_hybrid=True, + cudagraph_mode=None, + use_dflash=None, + compilation_mode=None, + enable_sp=False, + fuse_gemm_comms=False, + fuse_attn_quant=False, + ): + from types import SimpleNamespace + + from vllm.config import CompilationMode + + vllm_config = self._make_vllm_config(cudagraph_mode=cudagraph_mode) + vllm_config.model_config.is_hybrid = is_hybrid + # The real SpeculativeConfig.use_dflash is a method, so mock it as a + # callable (a plain value would not reproduce the always-truthy pitfall). + _dflash_flag = use_dflash if use_dflash is not None else (method == "dflash") + vllm_config.speculative_config = SimpleNamespace( + method=method, + num_speculative_tokens=num_speculative_tokens, + use_dflash=lambda flag=_dflash_flag: flag, + ) + # The guard reads compilation_config.mode (piecewise needs VLLM_COMPILE) + # and clears the full-graph-only fusion flags on pass_config. + vllm_config.compilation_config.mode = ( + CompilationMode.VLLM_COMPILE + if compilation_mode is None + else compilation_mode + ) + vllm_config.compilation_config.pass_config = SimpleNamespace( + enable_sp=enable_sp, + fuse_gemm_comms=fuse_gemm_comms, + fuse_attn_quant=fuse_attn_quant, + ) + return vllm_config + + def test_hybrid_mtp_full_cudagraph_coerced_to_piecewise(self): + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + vllm_config = self._make_hybrid_spec_config( + method="qwen3_5_mtp", + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE + + def test_hybrid_spec_coercion_is_method_agnostic(self): + # An unknown/renamed draft method must still be routed off the corrupt + # full-capture verify path. + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + vllm_config = self._make_hybrid_spec_config( + method="some_unrecognized_future_method", + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE + + def test_hybrid_dflash_target_verify_coerced_to_piecewise(self): + # dFlash's target verify hits the same corruption and must be coerced too + # (the drafter graph is configured separately, earlier in the hook). + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + os.environ.pop("VLLM_MUSA_DFLASH_FULL_WRAP", None) + vllm_config = self._make_hybrid_spec_config( + method="dflash", + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE + + def test_hybrid_spec_disables_full_graph_only_fusions(self): + # SP / fused-GEMM-comms / attn-quant fusion require full-graph compilation + # and set_splitting_ops_for_v1 would re-promote PIECEWISE to FULL; the guard + # clears them. + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + vllm_config = self._make_hybrid_spec_config( + method="qwen3_5_mtp", + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + enable_sp=True, + fuse_gemm_comms=True, + fuse_attn_quant=True, + ) + + MUSAPlatformBase.check_and_update_config(vllm_config) + + pc = vllm_config.compilation_config.pass_config + assert vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE + assert pc.enable_sp is False + assert pc.fuse_gemm_comms is False + assert pc.fuse_attn_quant is False + + def test_hybrid_spec_compilation_disabled_coerced_to_none(self): + # PIECEWISE requires VLLM_COMPILE; with compilation disabled the guard must + # fall back to NONE (PIECEWISE would trip the runtime dispatcher assertion). + from vllm.config import CompilationMode, CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + vllm_config = self._make_hybrid_spec_config( + method="qwen3_5_mtp", + cudagraph_mode=CUDAGraphMode.FULL, + compilation_mode=CompilationMode.NONE, + ) + + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE + + def test_non_hybrid_spec_keeps_full_cudagraph(self): + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + vllm_config = self._make_hybrid_spec_config( + method="qwen3_5_mtp", + is_hybrid=False, + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert ( + vllm_config.compilation_config.cudagraph_mode + == CUDAGraphMode.FULL_AND_PIECEWISE + ) + + def test_non_spec_hybrid_keeps_full_cudagraph(self): + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + os.environ.pop("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", None) + vllm_config = self._make_vllm_config( + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + vllm_config.model_config.is_hybrid = True + + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert ( + vllm_config.compilation_config.cudagraph_mode + == CUDAGraphMode.FULL_AND_PIECEWISE + ) + + def test_hybrid_spec_full_cudagraph_env_override_keeps_full(self): + from vllm.config import CUDAGraphMode + + from vllm_musa.platform import MUSAPlatformBase + + vllm_config = self._make_hybrid_spec_config( + method="qwen3_5_mtp", + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + + with patch.dict(os.environ, {"VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL": "1"}): + MUSAPlatformBase.check_and_update_config(vllm_config) + + assert ( + vllm_config.compilation_config.cudagraph_mode + == CUDAGraphMode.FULL_AND_PIECEWISE + ) + + def test_set_splitting_ops_promotes_piecewise_with_sp(self): + # Control (finding 2): confirms the promotion the guard defends against is + # real — SP forces a piecewise mode back to FULL in set_splitting_ops_for_v1. + from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode + + cc = CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + ) + cc.pass_config.enable_sp = True + + cc.set_splitting_ops_for_v1(all2all_backend="allgather_reducescatter") + + assert cc.cudagraph_mode == CUDAGraphMode.FULL + + def test_set_splitting_ops_preserves_piecewise_without_fusions(self): + # finding 2: with the full-graph-only fusions cleared (what the guard does), + # set_splitting_ops_for_v1 must NOT re-promote PIECEWISE to FULL. + from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode + + cc = CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + cudagraph_mode=CUDAGraphMode.PIECEWISE, + ) + cc.pass_config.enable_sp = False + cc.pass_config.fuse_gemm_comms = False + cc.pass_config.fuse_attn_quant = False + + cc.set_splitting_ops_for_v1(all2all_backend="allgather_reducescatter") + + assert cc.cudagraph_mode == CUDAGraphMode.PIECEWISE + def test_deepseek_v4_defaults_moe_gemv_block32x8(self): from vllm_musa.platform import MUSAPlatformBase diff --git a/vllm_musa/platform.py b/vllm_musa/platform.py index 4efd84e5f2fa..97e05e31d0e6 100644 --- a/vllm_musa/platform.py +++ b/vllm_musa/platform.py @@ -502,6 +502,59 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: _block, ) + # MUSA: a full decode CUDAGraph corrupts the multi-query speculative verify + # of a hybrid (GatedDeltaNet / Mamba) target. The paged multi-query verify + # attention miscomputes when captured inside the full model graph at + # query_len = num_speculative_tokens + 1; query_len=1 decode in the same + # graph is correct. The defect is method-agnostic, so gate on any + # hybrid-model speculative decode, including dFlash's target verify — an + # unknown/renamed draft method must not slip back into the corrupt path. + # + # The verify must run piecewise (eager attention). Sequence parallelism, + # fused GEMM comms, and attn-quant fusion require full-graph compilation, so + # set_splitting_ops_for_v1() would otherwise promote PIECEWISE back to FULL; + # they are disabled here for this case. When compilation is not enabled, + # piecewise cudagraphs are unavailable, so fall back to NONE. This avoids the + # full-graph corruption and preserves the piecewise result (it does not claim + # bit-exactness versus non-spec decode). + import os as _spec_os + + if ( + spec_config is not None + and getattr(spec_config, "num_speculative_tokens", 0) + and getattr(getattr(vllm_config, "model_config", None), "is_hybrid", False) + and _spec_os.environ.get("VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL", "0") != "1" + ): + _comp = getattr(vllm_config, "compilation_config", None) + if _comp is not None: + from vllm.config import CompilationMode as _CM + from vllm.config import CUDAGraphMode as _CGM + + _cg = getattr(_comp, "cudagraph_mode", None) + # None == "not yet defaulted"; the v1 default is FULL_AND_PIECEWISE. + if _cg is None or _cg.has_full_cudagraphs(): + if getattr(_comp, "mode", None) == _CM.VLLM_COMPILE: + _comp.cudagraph_mode = _CGM.PIECEWISE + _pc = getattr(_comp, "pass_config", None) + if _pc is not None: + for _flag in ( + "enable_sp", + "fuse_gemm_comms", + "fuse_attn_quant", + ): + if hasattr(_pc, _flag): + setattr(_pc, _flag, False) + else: + _comp.cudagraph_mode = _CGM.NONE + logger.warning( + "MUSA: hybrid-model speculative decode corrupts the " + "multi-query verify under a full decode CUDAGraph; setting " + "cudagraph_mode=%s (and disabling SP / fused-GEMM-comms / " + "attn-quant fusion, which would re-promote to FULL). Set " + "VLLM_MUSA_HYBRID_SPEC_ALLOW_FULL=1 to override.", + _comp.cudagraph_mode.name, + ) + # FP8 correctness under torch.compile + PIECEWISE CUDAGraph capture: the # MUSA RMSNorm / SiluAndMul / per-token-group FP8-quant custom-op kernels # (forward_oot) are opaque to Inductor's quant fusion and leave the quant's