feat(eagle3):support qwen3 dense model

Signed-off-by: xq25478 <[email protected]>

Author: xq25478

tensorrt_llm/_torch/models/modeling_qwen3.py

@@ -16,8 +16,9 @@
 from ..modules.linear import TensorParallelMode
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
-from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
-                             register_auto_model)
+from ..speculative import SpecMetadata
+from .modeling_speculative import SpecDecOneEngineForCausalLM
+from .modeling_utils import DecoderModel, EagerFusionConfig, register_auto_model
 
 
 class Qwen3Attention(Attention):
@@ -141,13 +142,21 @@ def __init__(
                                                 eps=config.rms_norm_eps,
                                                 dtype=config.torch_dtype)
 
+        self.fusion_config = EagerFusionConfig()
+        # self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp(
+        # )
+        # TODO: re-enable these fusions
+        self.fusion_config.PRE_MOE_FUSION = False
+        self.fusion_config.POST_MLP_FUSION = False
+
     def forward(
         self,
         position_ids: torch.IntTensor,
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
         residual: Optional[torch.Tensor],
         mrope_config: Optional[Tuple[torch.Tensor, int]] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         if residual is None:
@@ -171,6 +180,10 @@ def forward(
             hidden_states, residual)
         hidden_states = self.mlp(hidden_states)
 
+        if spec_metadata is not None:
+            spec_metadata.maybe_capture_hidden_states(self.layer_idx,
+                                                      hidden_states, residual)
+
         return hidden_states, residual
 
 
@@ -207,6 +220,7 @@ def forward(
         position_ids: Optional[torch.IntTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         mrope_config: Optional[Tuple[torch.Tensor, int]] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         if (input_ids is None) ^ (inputs_embeds is not None):
@@ -227,24 +241,23 @@ def forward(
                 attn_metadata=attn_metadata,
                 residual=residual,
                 mrope_config=mrope_config,
+                spec_metadata=spec_metadata,
             )
 
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
 
 @register_auto_model("Qwen3ForCausalLM")
-class Qwen3ForCausalLM(DecoderModelForCausalLM[Qwen3Model, Qwen3Config]):
+class Qwen3ForCausalLM(SpecDecOneEngineForCausalLM[Qwen3Model, Qwen3Config]):
 
     def __init__(
         self,
         model_config: ModelConfig[Qwen3Config],
     ):
         super().__init__(
             Qwen3Model(model_config),
-            config=model_config,
-            hidden_size=model_config.pretrained_config.hidden_size,
-            vocab_size=model_config.pretrained_config.vocab_size,
+            model_config,
         )
 
     # NOTE: Qwen2-VL needs special mrope_config so adding separate forward() function to accept 'mrope_config'.
@@ -256,6 +269,7 @@ def forward(
         inputs_embeds: Optional[torch.FloatTensor] = None,
         return_context_logits: bool = False,
         mrope_config: Optional[dict] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         output = self.model(
@@ -264,6 +278,7 @@ def forward(
             position_ids=position_ids,
             inputs_embeds=inputs_embeds,
             mrope_config=mrope_config,
+            spec_metadata=spec_metadata,
         )
 
         return self.logits_processor.forward(

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -1650,6 +1650,31 @@ def test_bf16(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph,
             task = MMLU(self.MODEL_NAME)
             task.evaluate(llm)
 
+    def test_eagle3(self):
+        pytorch_config = dict(
+            disable_overlap_scheduler=True,
+            cuda_graph_config=CudaGraphConfig(batch_sizes=[1]),
+        )
+        kv_cache_config = KvCacheConfig(enable_block_reuse=False)
+
+        eagle_model_dir = f"{llm_models_root()}/qwen3_8b_eagle3"
+        target_model_dir = f"{llm_models_root()}/Qwen3/Qwen3-8B"
+
+        draft_len = 4
+        spec_config = EagleDecodingConfig(max_draft_len=draft_len,
+                                          pytorch_weights_path=eagle_model_dir,
+                                          eagle3_one_model=False)
+
+        llm = LLM(model=target_model_dir,
+                  **pytorch_config,
+                  kv_cache_config=kv_cache_config,
+                  speculative_config=spec_config,
+                  build_config=None)
+
+        with llm:
+            task = MMLU(self.MODEL_NAME)
+            task.evaluate(llm)
+
 
 class TestQwen3_30B_A3B(LlmapiAccuracyTestHarness):
     MODEL_NAME = "Qwen3/Qwen3-30B-A3B"