[Spec Decode] Add Qwen3 architecture support for EAGLE3

[benchislett]

Purpose

As proven by DFlash, the Qwen3 base architecture is more stable for EAGLE3 training due to the extra norms in the attention module. We can improve performance of EAGLE3 and P-EAGLE by using Qwen3 as a base architecture and leveraging the higher learning rates now possible in training. I have manually ablated this and confirmed the results.

This PR is experimental as no public checkpoints currently use the Qwen3 base arch for EAGLE3 or P-EAGLE.

If this becomes likely, I will update and merge the PR accordingly. Until then, it will remain a draft.

[gemini-code-assist]

Code Review

This pull request adds support for the Eagle3 speculative decoding architecture for Qwen3 models by implementing the Qwen3Eagle3DecoderLayer, Qwen3Eagle3Model, and Eagle3Qwen3ForCausalLM classes. These changes include logic for auxiliary hidden states, specialized weight loading, and parallel drafting support. The model is also registered and integrated into the speculative decoding proposer. I have no feedback to provide.

[orestis-z]

Hey @benchislett — I tested this PR with speculators-format Qwen3 Eagle3 checkpoints (trained with --draft-arch qwen3 in speculators) and hit a bug: loading fails with KeyError: 'layers.0.self_attn.k_norm.weight'.

Root cause: The speculators config loader in vllm/transformers_utils/configs/speculators/algos.py hardcodes Eagle3LlamaForCausalLM as the architecture for all Eagle3 models (line 39), ignoring transformer_layer_config.model_type. So Qwen3-based drafters get loaded with the Llama model class, which doesn't have q_norm/k_norm.

Fix: Resolve the architecture from model_type:

diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py
--- a/vllm/transformers_utils/configs/speculators/algos.py
+++ b/vllm/transformers_utils/configs/speculators/algos.py
@@ -36,7 +36,13 @@
         "norm_before_residual", True
     )
     pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
-    pre_trained_config["architectures"] = ["Eagle3LlamaForCausalLM"]
+    eagle3_arch_map = {
+        "qwen3": "Eagle3Qwen3ForCausalLM",
+    }
+    model_type = pre_trained_config.get("model_type", "llama")
+    pre_trained_config["architectures"] = [
+        eagle3_arch_map.get(model_type, "Eagle3LlamaForCausalLM")
+    ]
     if config_dict.get("eagle_aux_hidden_state_layer_ids"):

I verified this fixes the issue — Qwen3 Eagle3 drafters load and run correctly with speculative decoding after the change. Happy to open a PR against your branch if that helps.

[orestis-z]

Hey @benchislett, I have a qwen3-arch based checkpoint here: https://huggingface.co/inference-optimization/Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1.
Do you need anything else to get this in?

[orestis-z]

Same issue applies to P-EAGLE with Qwen3 draft architecture — three additional changes are needed on top of the Eagle3 fix above:

1. Architecture resolution for P-EAGLE (algos.py) — update_peagle() hardcodes PeagleLlamaForCausalLM, same pattern as Eagle3:

diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py
--- a/vllm/transformers_utils/configs/speculators/algos.py
+++ b/vllm/transformers_utils/configs/speculators/algos.py
@@ -53,7 +59,13 @@
     PEagle specific fields:
     ...
     """
-    pre_trained_config["architectures"] = ["PeagleLlamaForCausalLM"]
+    peagle_arch_map = {
+        "qwen3": "PeagleQwen3ForCausalLM",
+    }
+    model_type = pre_trained_config.get("model_type", "llama")
+    pre_trained_config["architectures"] = [
+        peagle_arch_map.get(model_type, "PeagleLlamaForCausalLM")
+    ]

2. Model registry (registry.py) — add the Qwen3 P-EAGLE entry:

diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
     "Eagle3Qwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"),
+    "PeagleQwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"),

3. Parallel drafting for P-EAGLE (speculative.py) — peagle method should enable parallel_drafting (only needed when method is explicitly passed, since base.py already handles auto-detection):

diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
-        if self.method in ("dflash",):
+        if self.method in ("dflash", "peagle"):
             self.parallel_drafting = True

Verified with a Qwen3-arch P-EAGLE checkpoint — all three changes are required for correct loading and serving.

[benchislett]

@orestis-z I added support for speculators models as you suggested. I did not add the "method": "peagle" check since this is not a supported method, we should always be using either "eagle" (with or without "parallel_drafting": True) or "dflash".

I added some tests to ensure the correct behaviour here. They're passing locally. Let me know if you want me to remove them until official nm-testing checkpoints are released, or merge as-is and we can follow-up with LTS checkpoints when they are available

[orestis-z]

@benchislett Thanks for adding the tests! I'd say merge as-is. We can swap in official nm-testing checkpoints in a follow-up once they're published.

[shanjiaz]

@benchislett Just curious, are quantized eagle3 drafters supported now? Are we planning on supporting quantized dflash drafters as well? (Just saw in qwen3-dflash seems like dflash also supports quantized draft now) Very cool.

[benchislett]

I see no blockers for quantized drafters. It's planned to be supported but we don't really have any FP8 drafter checkpoints that we can use to pipeclean.

[mgoin]

LGTM, just the seemingly bad call on get_cache_scale