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

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

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():