[Spec Decode][ROCm][Perf] Default EAGLE3 Llama drafter to AITER unified attention

tests/test_config.py

@@ -28,13 +28,15 @@
 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,
     OptimizationLevel,
 )
 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,

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(