[1/N][Mix] Mixed Chunk Prefill Base (#36288)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/model_executor/runner/__init__.py

@@ -12,7 +12,8 @@ Public API:
   - BaseCudaGraphRunner — abstract cuda-graph base; bucket padding +
     capture-loop scaffolding on top of BaseRunner.
   - DecodeCudaGraphRunner — concrete decode-phase runner.
-  - PrefillCudaGraphRunner — concrete prefill-phase runner.
+  - PrefillCudaGraphRunner — concrete prefill-phase runner (extend family;
+    MIXED batches replay the EXTEND-captured graphs).
   - EagerRunner — no-cuda-graph runner; runs model.forward live (the
     eager dual of the cuda-graph runners), mode-dispatched over decode +
     extend + idle.

python/sglang/srt/model_executor/runner/eager_runner.py

@@ -27,6 +27,7 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.cp.utils import (
     cp_gather_after_forward,
     cp_shard_model_inputs,
+    get_cp_strategy,
     is_cp_v2_active,
     prepare_cp_forward,
 )
@@ -37,7 +38,11 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import (
 from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
     create_chunked_prefix_cache_kv_indices,
 )
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    ForwardMode,
+    PPProxyTensors,
+)
 from sglang.srt.model_executor.forward_context import (
     ForwardContext,
     forward_context,
@@ -59,7 +64,7 @@ from sglang.srt.runtime_context import (
     max_prefill_buffer_tokens,
     max_speculative_num_draft_tokens,
 )
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_hip, is_npu
 from sglang.srt.utils.common import (
     ceil_align,
     get_eager_max_batch_size,
@@ -208,6 +213,12 @@ class EagerRunner(BaseRunner):
         self, forward_batch: ForwardBatch, pp_proxy_tensors=None, **kwargs
     ) -> Any:
         mode = forward_batch.forward_mode
+        if mode.is_mixed() and not is_npu() and get_cp_strategy() is None:
+            # A mixed batch is extend-shaped (decode tails are 1-token
+            # extends); run it as EXTEND. NPU keeps MIXED for its dedicated
+            # kernel; CP keeps it to skip the zigzag split.
+            forward_batch.forward_mode = ForwardMode.EXTEND
+            mode = ForwardMode.EXTEND
         if mode.is_decode():
             return self._execute_decode(forward_batch, pp_proxy_tensors)
         if mode.is_idle():

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -262,6 +262,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     """
 
     def __init__(self, model_runner: ModelRunner):
+        if get_schedule().enable_mixed_chunk:
+            backend = get_exec().graph.cuda_graph_config.prefill.backend
+            assert backend == Backend.BREAKABLE, (
+                "Mixed chunk prefill requires the breakable prefill CUDA "
+                f"graph backend; got '{backend}'."
+            )
         super().__init__(model_runner)
         # --- model flags ----------------------------------------------
         self.quant_config = getattr(model_runner.model, "quant_config", None)
@@ -1527,8 +1533,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             else forward_batch.num_token_non_padded
         )
 
-        # Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1)
-        # doesn't fail on MIXED=3.
+        # MIXED replays the EXTEND-captured graphs.
         pcg_forward_mode = (
             ForwardMode.EXTEND
             if forward_batch.forward_mode == ForwardMode.MIXED