diff --git a/tests/test_config.py b/tests/test_config.py index 3837057658b5..f5ac597efafe 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -28,6 +28,7 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.config.kernel import IrOpPriorityConfig from vllm.config.load import LoadConfig +from vllm.config.speculative import _should_default_rocm_eagle3_attention_backend from vllm.config.utils import get_field from vllm.config.vllm import ( OPTIMIZATION_LEVEL_TO_CONFIG, @@ -35,6 +36,7 @@ ) from vllm.platforms import current_platform from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.registry import AttentionBackendEnum DEVICE_TYPE = current_platform.device_type @@ -271,6 +273,87 @@ def test_is_default_v2_model_runner_model(model_config, expected): assert VllmConfig._is_default_v2_model_runner_model(config) is expected +@pytest.mark.parametrize( + ( + "method", + "attention_backend", + "draft_model_config", + "is_rocm", + "expected", + ), + [ + ( + "eagle3", + None, + SimpleNamespace(architectures=["LlamaForCausalLMEagle3"]), + True, + True, + ), + ( + "eagle3", + AttentionBackendEnum.ROCM_ATTN, + SimpleNamespace(architectures=["LlamaForCausalLMEagle3"]), + True, + False, + ), + ( + "eagle3", + None, + SimpleNamespace(architectures=["LlamaForCausalLMEagle3"]), + False, + False, + ), + ( + "eagle3", + None, + SimpleNamespace(architectures=["Eagle3Qwen3ForCausalLM"]), + True, + False, + ), + ( + "eagle3", + None, + SimpleNamespace( + architectures=[], + hf_config=SimpleNamespace(architectures=["Eagle3LlamaForCausalLM"]), + ), + True, + True, + ), + ], +) +def test_should_default_rocm_eagle3_attention_backend( + method, attention_backend, draft_model_config, is_rocm, expected +): + assert ( + _should_default_rocm_eagle3_attention_backend( + method, + attention_backend, + draft_model_config, + is_rocm=is_rocm, + ) + is expected + ) + + +def test_rocm_eagle3_lite_llama_defaults_to_aiter_unified_attention( + monkeypatch, +): + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + speculative_config = SimpleNamespace( + method="eagle3", + attention_backend=None, + draft_model_config=SimpleNamespace(architectures=["LlamaForCausalLMEagle3"]), + ) + + SpeculativeConfig._maybe_default_rocm_eagle3_attention_backend(speculative_config) + + assert ( + speculative_config.attention_backend + is AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN + ) + + @pytest.mark.skip_global_cleanup def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping( monkeypatch, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 7fc4700fa2b5..aee7a8e9884d 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -74,6 +74,41 @@ DraftSampleMethod = Literal["greedy", "probabilistic"] +_LLAMA_EAGLE3_DRAFT_ARCHITECTURES = frozenset( + ("Eagle3LlamaForCausalLM", "LlamaForCausalLMEagle3") +) + + +def _has_llama_eagle3_draft_architecture( + draft_model_config: ModelConfig | None, +) -> bool: + if draft_model_config is None: + return False + + architectures = getattr(draft_model_config, "architectures", None) or () + if _LLAMA_EAGLE3_DRAFT_ARCHITECTURES.intersection(architectures): + return True + + hf_config = getattr(draft_model_config, "hf_config", None) + hf_architectures = getattr(hf_config, "architectures", None) or () + return bool(_LLAMA_EAGLE3_DRAFT_ARCHITECTURES.intersection(hf_architectures)) + + +def _should_default_rocm_eagle3_attention_backend( + method: SpeculativeMethod | None, + attention_backend: AttentionBackendEnum | None, + draft_model_config: ModelConfig | None, + *, + is_rocm: bool, +) -> bool: + return ( + is_rocm + and attention_backend is None + and method == "eagle3" + and _has_llama_eagle3_draft_architecture(draft_model_config) + ) + + @config class SpeculativeConfig: """Configuration for speculative decoding.""" @@ -869,6 +904,8 @@ def __post_init__(self): if self.method in ("dflash", "dspark"): self.parallel_drafting = True + self._maybe_default_rocm_eagle3_attention_backend() + if self.num_speculative_tokens is not None and hasattr( self.draft_model_config.hf_config, "num_lookahead_tokens" ): @@ -922,6 +959,23 @@ def __post_init__(self): ) return self + def _maybe_default_rocm_eagle3_attention_backend(self) -> None: + from vllm.platforms import current_platform + + if not _should_default_rocm_eagle3_attention_backend( + self.method, + self.attention_backend, + self.draft_model_config, + is_rocm=current_platform.is_rocm(), + ): + return + + self.attention_backend = AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN + logger.info_once( + "Using ROCM_AITER_UNIFIED_ATTN for ROCm EAGLE3 Llama draft model " + "attention. Set speculative_config.attention_backend to override." + ) + def _validate_suffix_decoding(self): if not has_arctic_inference(): raise ImportError(