shared_exp_dp_v2

Signed-off-by: chenmenglong <[email protected]>

Author: dragondream-chen

vllm_ascend/models/layers/mla.py

@@ -151,17 +151,20 @@ def forward(
             hidden_states: torch.Tensor,
             kv_cache: Optional[torch.Tensor] = None,
             attn_metadata: Optional[AttentionMetadata] = None) -> torch.Tensor:
-        forward_context = get_forward_context()
-        sp_enabled = forward_context.sp_enabled
-        need_gather_q_kv = False
-        if sp_enabled and self.debug_layer_idx < self.layers:
-            need_gather_q_kv = True
-        if not sp_enabled or self.debug_layer_idx < self.layers:
-            output_shape = hidden_states.shape
-        else:
-            # used in deepseek mtp layer
-            output_shape = torch.chunk(hidden_states, self.tp_size,
-                                       dim=0)[0].shape
+        # forward_context = get_forward_context()
+        # sp_enabled = forward_context.sp_enabled
+        # need_gather_q_kv = False
+        # if sp_enabled and self.debug_layer_idx < self.layers:
+        #     need_gather_q_kv = True
+        # if not sp_enabled or self.debug_layer_idx < self.layers:
+        #     output_shape = hidden_states.shape
+        # else:
+        #     # used in deepseek mtp layer
+        #     output_shape = torch.chunk(hidden_states, self.tp_size,
+        #                                dim=0)[0].shape
+
+        need_gather_q_kv = get_forward_context().sp_enabled
+        output_shape = hidden_states.shape
         # FIXME: This does not seem right, should make sure the buffer is fixed
         output = torch.empty(output_shape,
                              dtype=hidden_states.dtype,

vllm_ascend/patch/worker/patch_deepseek_mtp.py

@@ -52,50 +52,31 @@ def predictor_init(self, vllm_config: VllmConfig, prefix: str) -> None:
                                             topk_indices_buffer)
 
 
-def predictor_forward(
-    self,
-    input_ids: torch.Tensor,
-    positions: torch.Tensor,
-    previous_hidden_states: torch.Tensor,
-    inputs_embeds: torch.Tensor,
-    spec_step_index: int = 0,
-) -> torch.Tensor:
-    assert inputs_embeds is not None
-    inputs_embeds = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
-        inputs_embeds, True)
-    # masking inputs at position 0, as not needed by MTP
-    inputs_embeds[positions == 0] = 0
-    inputs_embeds = self.enorm(inputs_embeds)
-    previous_hidden_states = self.hnorm(previous_hidden_states)
-
-    hidden_states = self.eh_proj(
-        torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
-
-    hidden_states, residual = self.mtp_block(positions=positions,
-                                             hidden_states=hidden_states,
-                                             residual=None)
-    hidden_states = residual + hidden_states
-    return hidden_states
+# def predictor_forward(
+#     self,
+#     input_ids: torch.Tensor,
+#     positions: torch.Tensor,
+#     previous_hidden_states: torch.Tensor,
+#     inputs_embeds: torch.Tensor,
+#     spec_step_index: int = 0,
+# ) -> torch.Tensor:
+#     assert inputs_embeds is not None
+#     inputs_embeds = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
+#         inputs_embeds, True)
+#     # masking inputs at position 0, as not needed by MTP
+#     inputs_embeds[positions == 0] = 0
+#     inputs_embeds = self.enorm(inputs_embeds)
+#     previous_hidden_states = self.hnorm(previous_hidden_states)
+
+#     hidden_states = self.eh_proj(
+#         torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
+
+#     hidden_states, residual = self.mtp_block(positions=positions,
+#                                              hidden_states=hidden_states,
+#                                              residual=None)
+#     hidden_states = residual + hidden_states
+#     return hidden_states
 
 
 DeepSeekMultiTokenPredictorLayer.__init__ = predictor_init
-DeepSeekMultiTokenPredictorLayer.forward = predictor_forward
-
-
-def mtp_forward(
-    self,
-    input_ids: torch.Tensor,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    intermediate_tensors: IntermediateTensors,
-    inputs_embeds: torch.Tensor,
-    spec_step_idx: int = 0,
-) -> torch.Tensor:
-    hidden_states = self.model(input_ids, positions, hidden_states,
-                               inputs_embeds, spec_step_idx)
-    hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
-        hidden_states, True)
-    return hidden_states
-
-
-DeepSeekMTP.forward = mtp_forward
+# DeepSeekMultiTokenPredictorLayer.forward = predictor_forward

vllm_ascend/spec_decode/mtp_proposer.py

@@ -190,6 +190,8 @@ def dummy_run(self,
                         kv_caches=self.runner.kv_caches[-1:],
                         spec_step_idx=0)
                 else:
+                    positions = torch.ops.vllm.maybe_pad_and_reduce(positions)
+                    previous_hidden_states = torch.ops.vllm.maybe_pad_and_reduce(previous_hidden_states)
                     self.model(input_ids=input_ids,
                                positions=positions,
                                hidden_states=previous_hidden_states)
@@ -474,10 +476,19 @@ def _propose(
                             spec_step_idx=0,
                             **model_kwargs)
                     else:
+                        input_ids=self.input_ids[:num_input_tokens]
+                        positions=self.positions[:num_input_tokens]
+                        hidden_states=self.hidden_states[:num_input_tokens]
+                        
+                        positions = torch.ops.vllm.maybe_pad_and_reduce(positions)
+                        previous_hidden_states = torch.ops.vllm.maybe_pad_and_reduce(previous_hidden_states)
                         hidden_states = self.model(
-                            input_ids=self.input_ids[:num_input_tokens],
-                            positions=self.positions[:num_input_tokens],
-                            hidden_states=self.hidden_states[:num_input_tokens]
+                            input_ids=input_ids,
+                            positions=positions,
+                            hidden_states=hidden_states
+                        )
+                        hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
+                            hidden_states.contiguous(), True
                         )
 
             num_indices = last_token_indices.shape[0]