Fix eagle3 cuda graph (#8163)

99aefa03 · Jay Zhou · GitHub · bbcfbc1a · 99aefa03
Unverified Commit 99aefa03 authored Jul 20, 2025 by Jay Zhou Committed by GitHub Jul 20, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 9 additions and 1 deletion

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py ...g/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +9 -1

No files found.
--- 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,
                )