[AMD] Support two batch overlap with MTP on DeepSeekV4 (#30238)

python/sglang/jit_kernel/dsv4/compress.py

@@ -10,6 +10,7 @@ from sglang.jit_kernel.utils import (
     load_jit,
     make_cpp_args,
 )
+from sglang.srt.utils import is_hip
 
 from .utils import make_name
 
@@ -213,6 +214,20 @@ class CompressorPrefillPlan(NamedTuple):
             dtype=torch.uint8,
             pin_memory=not is_gpu_input,
         )
+        # DP-safe empty-batch guard: a TBO ubatch (or tail batch) can have 0
+        # query tokens (num_q_tokens==0) on THIS rank while other DP ranks are
+        # non-empty. The global TBO decision must stay uniform across ranks, so
+        # return an empty plan here (downstream compressor then processes 0
+        # tokens = no-op) instead of skipping TBO per-rank. Avoids the
+        # c_plan.cuh RuntimeCheck(batch_size <= num_q_tokens) failure at B>=1.
+        if int(num_q_tokens) == 0 and is_hip():
+            _dev = req_to_token.device
+            return CompressorPrefillPlan(
+                compress_ratio,
+                torch.empty((0, 16), dtype=torch.uint8, device=_dev),
+                torch.empty((0, 8), dtype=torch.uint8, device=_dev),
+                pin_buffer,
+            )
         module = _jit_compress_plan_module()
         plan_c, plan_w = module.plan_prefill(
             req_pool_indices,

python/sglang/srt/models/deepseek_v4.py

@@ -2148,7 +2148,9 @@ class DeepseekV4Model(nn.Module):
             and forward_batch.can_run_tbo
             and forward_batch.tbo_children is not None
             and forward_batch.global_forward_mode is not None
-            and forward_batch.global_forward_mode.is_extend()
+            # MTP target-verify also reports is_extend(); only real prefill
+            # should enter the prefill TBO strategy.
+            and forward_batch.global_forward_mode.is_extend_without_speculative()
             and not dsa_use_prefill_cp(forward_batch)
             and self.pp_group.world_size == 1
         )