Merge remote-tracking branch 'origin/main' into sh/update_dlblas

Author: SHshenhao

lmdeploy/pytorch/engine/model_agent.py

@@ -58,7 +58,7 @@ def to_tensor(self):
         if isinstance(self.vals, torch.Tensor):
             vals = self.vals
         else:
-            vals = torch.from_numpy(vals)
+            vals = torch.from_numpy(self.vals)
         return BatchedLogProbs(vals=vals, indices=torch.from_numpy(self.indices))
 
 
@@ -729,7 +729,7 @@ def __update_inputs(next_token_ids, model_metas, extra_inputs):
                 inputs,
                 return_logits=return_logits,
                 sync_long_context=sync_long_context,
-                return_routed_experts=return_routed_experts and need_output,
+                return_routed_experts=return_routed_experts and self.need_output,
             )
             logits = output['logits'][0]  # [bs, seq, prob] -> [seq, prob]
             seq_length = output.get('seq_length', inputs.seq_length)
@@ -959,8 +959,7 @@ def _build_model(self):
             update_custom_module_map(custom_module_map)
         logger.debug(msg_with_rank(rank, 'build model.'))
         # for router replay
-        need_output = self.dist_ctx.dp > 1 or self.dist_ctx.rank % self.dist_ctx.tp == 0
-        enable_return_routed_experts = self.misc_config.enable_return_routed_experts and need_output
+        enable_return_routed_experts = self.misc_config.enable_return_routed_experts and self.need_output
 
         build_model_ctx = BuildModelContext(
             disable_vision_encoder=self.misc_config.disable_vision_encoder,

lmdeploy/serve/proxy/proxy.py

@@ -587,7 +587,7 @@ async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Reque
         if request.stream is True:
             response = node_manager.stream_generate(request_dict, node_url, '/v1/chat/completions')
             background_task = node_manager.create_background_tasks(node_url, start)
-            return StreamingResponse(response, background=background_task)
+            return StreamingResponse(response, background=background_task, media_type='text/event-stream')
         else:
             response = await node_manager.generate(request_dict, node_url, '/v1/chat/completions')
             node_manager.post_call(node_url, start)
@@ -649,7 +649,7 @@ async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Reque
         if request.stream is True:
             response = node_manager.stream_generate(request_dict, d_url, '/v1/chat/completions')
             background_task = node_manager.create_background_tasks(d_url, start)
-            resp = StreamingResponse(response, background=background_task)
+            resp = StreamingResponse(response, background=background_task, media_type='text/event-stream')
         else:
             response = await node_manager.generate(request_dict, d_url, '/v1/chat/completions')
             node_manager.post_call(d_url, start)
@@ -717,7 +717,7 @@ async def completions_v1(request: CompletionRequest, raw_request: Request = None
         if request.stream is True:
             response = node_manager.stream_generate(request_dict, node_url, '/v1/completions')
             background_task = node_manager.create_background_tasks(node_url, start)
-            return StreamingResponse(response, background=background_task)
+            return StreamingResponse(response, background=background_task, media_type='text/event-stream')
         else:
             response = await node_manager.generate(request_dict, node_url, '/v1/completions')
             node_manager.post_call(node_url, start)
@@ -793,7 +793,7 @@ async def completions_v1(request: CompletionRequest, raw_request: Request = None
         if request.stream is True:
             response = node_manager.stream_generate(request_dict, d_url, '/v1/completions')
             background_task = node_manager.create_background_tasks(d_url, start)
-            resp = StreamingResponse(response, background=background_task)
+            resp = StreamingResponse(response, background=background_task, media_type='text/event-stream')
         else:
             response = await node_manager.generate(request_dict, d_url, '/v1/completions')
             node_manager.post_call(d_url, start)