support mtp with deepseek r1 nvfp4 model (#13115)

Co-authored-by: Trevor Morris <tmorris@nvidia.com>
2025-12-06 00:45:54 -08:00
parent 42fcf5438f
commit ea177372bd
11 changed files with 123 additions and 59 deletions
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -764,7 +764,12 @@ class ForwardBatch:

        bs = self.batch_size

-        if self.forward_mode.is_decode():
+        if (
+            self.forward_mode.is_decode()
+            or self.forward_mode.is_target_verify()
+            or self.forward_mode.is_draft_extend(include_v2=True)
+            or self.forward_mode.is_idle()
+        ):
            if self.is_extend_in_batch and dp_padding_mode.is_max_len():
                setattr(self, "_original_forward_mode", self.forward_mode)
                self.forward_mode = ForwardMode.EXTEND