update dp ep mpt

Author: grimoire

lmdeploy/pytorch/spec_decode/spec_agent.py

@@ -205,7 +205,10 @@ def _prepare_inputs_from_main(self, model_inputs: ModelInputs, extra_inputs: Ext
         history_lengths = model_inputs.history_lengths.clone()
 
         if not model_inputs.is_chunk:
-            if not model_inputs.is_dummy and not model_inputs.is_decoding:
+            local_is_decoding = model_inputs.is_decoding
+            if model_inputs.dp_meta is not None:
+                local_is_decoding = model_inputs.dp_meta.is_decoding
+            if not model_inputs.is_dummy and not local_is_decoding:
                 # Non-chunk prefill starts an independent stream. Dummy DP
                 # placeholders and interleaved decode must not clear a pending
                 # long-chunk carry.

tests/pytorch/spec_decode/test_spec_agent.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from lmdeploy.pytorch.model_inputs import ModelInputs
+from lmdeploy.pytorch.model_inputs import DPMeta, ModelInputs
 from lmdeploy.pytorch.spec_decode.spec_agent import SpecModelAgent, _expand_sampling_inputs
 from lmdeploy.pytorch.strategies.ar_spec.model_agent import ARSpecExtraInputs
 
@@ -312,7 +312,13 @@ def test_slice_sampling_inputs_prefill():
     assert result is sampling_inputs
 
 
-def _model_inputs(input_ids, *, is_decoding=False, is_chunk=False, is_first_chunk=False, is_last_chunk=False):
+def _model_inputs(input_ids,
+                  *,
+                  is_decoding=False,
+                  is_chunk=False,
+                  is_first_chunk=False,
+                  is_last_chunk=False,
+                  dp_meta=None):
     input_ids = torch.tensor([input_ids])
     seq_length = torch.tensor([input_ids.size(1)])
     history_lengths = torch.tensor([0])
@@ -330,6 +336,7 @@ def _model_inputs(input_ids, *, is_decoding=False, is_chunk=False, is_first_chun
         is_chunk=is_chunk,
         is_first_chunk=is_first_chunk,
         is_last_chunk=is_last_chunk,
+        dp_meta=dp_meta,
     )
 
 
@@ -370,3 +377,16 @@ def test_prepare_inputs_from_main_clears_chunk_carry_on_non_chunk_prefill():
     agent._prepare_inputs_from_main(prefill, _extra([[1, 10], [2, 20], [3, 30]]))
 
     assert agent._prev_chunk_last == {}
+
+
+def test_prepare_inputs_from_main_keeps_chunk_carry_for_dp_local_decode_global_prefill():
+    agent = SpecModelAgent.__new__(SpecModelAgent)
+    saved = torch.ones(1, 1, 2)
+    agent._prev_chunk_last = {'hidden_states': saved.clone()}
+    agent.proposer = _DummyProposer()
+
+    dp_meta = DPMeta(dp_batches=[1, 1], is_decoding=True, dp_is_decoding=False)
+    inputs = _model_inputs([90, 91, 92], is_decoding=False, dp_meta=dp_meta)
+    agent._prepare_inputs_from_main(inputs, _extra([[9, 90], [8, 80], [7, 70]]))
+
+    assert torch.equal(agent._prev_chunk_last['hidden_states'], saved)