Fix eagle3 cuda graph (#8163)

2025-07-20 00:28:06 -07:00
parent bbcfbc1a02
commit 99aefa037e
1 changed files with 9 additions and 1 deletions
--- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
@@ -84,7 +84,15 @@ class EAGLEDraftExtendCudaGraphRunner:
                self.hidden_states = torch.zeros(
                    (
                        self.max_num_token,
-                        self.model_runner.model_config.hidden_size * 3,
+                        (
+                            self.model_runner.model_config.hf_config.target_hidden_size
+                            * 3
+                            if hasattr(
+                                self.model_runner.model_config.hf_config,
+                                "target_hidden_size",
+                            )
+                            else self.model_runner.model_config.hidden_size * 3
+                        ),
                    ),
                    dtype=self.model_runner.dtype,
                )