[Bugfix] Fix full prefill CUDA graph padding and EAGLE capture (#35588)

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -1403,13 +1403,12 @@ class TestPrefillNumTokenNonPaddedPostFill(unittest.TestCase):
         # pads. local = clamp(1018 - 512, 0, 512).
         self.assertEqual(self._fill(attn_tp_rank=1, attn_tp_size=2), 506)
 
-    def test_non_gathered_keeps_plain_fb_copy(self):
-        # Without a gathered buffer there is no attn-TP scatter; the plain FB
-        # copy must be preserved (post_fill no-op), mirroring the decode
-        # registry's contract.
+    def test_non_gathered_uses_raw_token_count(self):
+        # Full prefill graphs need the live raw boundary even without a
+        # gathered buffer so model layers can discard the padded bucket tail.
         self.assertEqual(
             self._fill(attn_tp_rank=0, attn_tp_size=2, require_gathered_buffer=False),
-            509,
+            1018,
         )
 
 

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -152,28 +152,25 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
 
     def test_eagle_target_tc_piecewise_skips_last_mode_capture(self):
         eager_runner = object()
-        # The server-side hidden-state ceiling is a bag leaf.
+        # The server-side hidden-state ceiling and graph config are bag leaves.
         override = get_context().override_server_args(
             enable_return_hidden_states=True,
             return_hidden_states_mode="last",
+            cuda_graph_config=SimpleNamespace(
+                prefill=SimpleNamespace(backend=Backend.TC_PIECEWISE)
+            ),
         )
         override.install()
         self.addCleanup(override.restore)
         model_runner = SimpleNamespace(
             is_draft_worker=False,
             spec_algorithm=SimpleNamespace(is_eagle=lambda: True),
-            server_args=SimpleNamespace(),
         )
 
-        with patch.object(
-            graph_setup,
-            "check_cuda_graph_backend",
-            return_value=False,
-        ):
-            capture = capture_prefill_graph(
-                model_runner=model_runner,
-                eager_runner=eager_runner,
-            )
+        capture = capture_prefill_graph(
+            model_runner=model_runner,
+            eager_runner=eager_runner,
+        )
 
         self.assertIs(capture.runner, eager_runner)
 
@@ -234,6 +231,36 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
 
         self.assertIs(static_batch.mm_input_embeds, mm_input_embeds)
 
+    def test_eagle_target_full_reaches_graph_construction(self):
+        override = get_context().override_server_args(
+            enable_return_hidden_states=True,
+            return_hidden_states_mode="last",
+            cuda_graph_config=SimpleNamespace(
+                prefill=SimpleNamespace(backend=Backend.FULL)
+            ),
+        )
+        override.install()
+        self.addCleanup(override.restore)
+        model_runner = SimpleNamespace(
+            is_draft_worker=False,
+            lora_manager=None,
+            model=object(),
+            spec_algorithm=SimpleNamespace(is_eagle=lambda: True),
+        )
+
+        with (
+            patch.object(
+                graph_setup,
+                "resolve_language_model",
+                side_effect=RuntimeError("reached graph construction"),
+            ),
+            self.assertRaisesRegex(RuntimeError, "reached graph construction"),
+        ):
+            capture_prefill_graph(
+                model_runner=model_runner,
+                eager_runner=object(),
+            )
+
     def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
         graph_config = SimpleNamespace(
             prefill=SimpleNamespace(full_prefill_prefix_chunk_tokens=None, max_bs=8)
@@ -243,7 +270,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         override = get_context().override_server_args(
             chunked_prefill_size=16, cuda_graph_config=graph_config
         )
-        published = override.install()
+        override.install()
         self.addCleanup(override.restore)
         model_runner = SimpleNamespace(
             server_args=SimpleNamespace(),