diff --git a/lightllm/server/httpserver/pd_loop.py b/lightllm/server/httpserver/pd_loop.py index f2c7d4b9b..c99ed419f 100644 --- a/lightllm/server/httpserver/pd_loop.py +++ b/lightllm/server/httpserver/pd_loop.py @@ -11,7 +11,12 @@ import sys from typing import Dict, Optional, Union, List from websockets import ClientConnection -from lightllm.server.pd_io_struct import NodeRole, ObjType +from lightllm.server.pd_io_struct import ( + NodeRole, + ObjType, + PD_COMPACT_TOKEN_INFO_LEN, + build_pd_compact_token_info, +) from lightllm.server.httpserver.async_queue import AsyncQueue from lightllm.utils.net_utils import get_hostname_ip from lightllm.utils.log_utils import init_logger @@ -244,10 +249,15 @@ async def _pd_process_generate( pd_upload_websocket=pd_upload_websocket, pd_event=pd_event, ): - metadata["node_mode"] = manager.args.run_mode + if metadata.get("count_output_tokens") == 1: + metadata["node_mode"] = manager.args.run_mode if not return_output_logprobs: for key in ("id", "logprob", "cumlogprob", "special", "logprobs"): metadata.pop(key, None) + compact_token_info = build_pd_compact_token_info(sub_req_id, request_output, metadata, finish_status) + if compact_token_info is not None: + await forwarding_queue.put(compact_token_info) + continue await forwarding_queue.put((sub_req_id, request_output, metadata, finish_status)) except PDPrefillNodeStopGenToken as e: logger.info(f"pd prefill node stop gen token for group_request_id {e.group_request_id}") @@ -262,7 +272,16 @@ async def _up_tokens_to_pd_master(forwarding_queue: AsyncQueue, websocket: Clien if handle_list: load_info: dict = _get_load_info() - await websocket.send(pickle.dumps((ObjType.TOKEN_PACKS, handle_list, load_info))) + group_start = 0 + group_is_compact = len(handle_list[0]) == PD_COMPACT_TOKEN_INFO_LEN + for index in range(1, len(handle_list) + 1): + item_is_compact = index < len(handle_list) and len(handle_list[index]) == PD_COMPACT_TOKEN_INFO_LEN + if index == len(handle_list) or item_is_compact != group_is_compact: + token_list = handle_list[group_start:index] + obj_type = ObjType.TOKEN_PACKS_COMPACT if group_is_compact else ObjType.TOKEN_PACKS + await websocket.send(pickle.dumps((obj_type, token_list, load_info))) + group_start = index + group_is_compact = item_is_compact async def _send_heartbeat_to_pd_master(websocket: ClientConnection): diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index c636caded..5e0649e22 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -12,7 +12,13 @@ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) from typing import Union, List, Tuple, Dict, Optional from lightllm.server.core.objs import FinishStatus -from ..pd_io_struct import PD_Client_Obj, PDUpKVStatus, ObjType, PDDecodeNodeInfo +from ..pd_io_struct import ( + PD_Client_Obj, + PDUpKVStatus, + ObjType, + PDDecodeNodeInfo, + unpack_pd_compact_token_info, +) from lightllm.server.core.objs import SamplingParams, StartArgs from ..multimodal_params import MultimodalParams from ..tokenizer import get_tokenizer @@ -635,12 +641,19 @@ async def handle_loop(self): try: for obj in objs: - if obj[0] == ObjType.TOKEN_PACKS: + if obj[0] in (ObjType.TOKEN_PACKS, ObjType.TOKEN_PACKS_COMPACT): token_list, node_load_info = obj[1], obj[2] self.pd_manager.update_node_load_info(node_load_info) - for sub_req_id, text, metadata, finish_status in token_list: - finish_status: FinishStatus = finish_status + compact_pack = obj[0] == ObjType.TOKEN_PACKS_COMPACT + for token_info in token_list: + if compact_pack: + sub_req_id, text, metadata, finish_status_value = unpack_pd_compact_token_info( + token_info + ) + finish_status = FinishStatus(finish_status_value) + else: + sub_req_id, text, metadata, finish_status = token_info group_req_id = convert_sub_id_to_group_id(sub_req_id) try: req_status: ReqStatus = self.req_id_to_out_inf[group_req_id] diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index c019e90ac..a8bda108b 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -2,7 +2,7 @@ import time import copy from dataclasses import dataclass, field -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple from lightllm.server.req_id_generator import convert_sub_id_to_group_id from fastapi import WebSocket @@ -41,6 +41,100 @@ class ObjType(enum.Enum): PD_UPLOAD_PREFILL_PROMPT_IDS = 4 # prefill 节点上报生成的 prompt ids 信息。 PD_REQ_DECODE_NODE_INFO = 5 # pd master 节点下发给 prefill 节点的请求对应的 decode 节点信息。 HEARTBEAT = 6 # P/D 节点向 pd master 上报的心跳。 + TOKEN_PACKS_COMPACT = 7 # 不含 logprobs 等可选字段的紧凑 token 包。 + + +PD_COMPACT_TOKEN_INFO_LEN = 11 +PDCompactTokenInfo = Tuple[ + int, # sub request id + str, # decoded text + int, # count_output_tokens + int, # prompt_tokens + int, # prompt_cache_len + int, # mtp_accepted_token_num + int, # mtp_verify_token_num + int, # mtp_verify_step_num + int, # finish status + Optional[str], # node_mode, first token only + Optional[Tuple[int, int, int]], # input text/audio/image tokens, first token only +] +_PD_COMPACT_METADATA_KEYS = frozenset( + { + "count_output_tokens", + "prompt_tokens", + "prompt_cache_len", + "mtp_accepted_token_num", + "mtp_verify_token_num", + "mtp_verify_step_num", + "node_mode", + "input_usage", + } +) +_PD_INPUT_USAGE_KEYS = frozenset({"input_text_tokens", "input_audio_tokens", "input_image_tokens"}) + + +def build_pd_compact_token_info(sub_req_id, text, metadata, finish_status) -> Optional[PDCompactTokenInfo]: + """Build the lossless compact form, or return None for optional metadata.""" + if not metadata.keys() <= _PD_COMPACT_METADATA_KEYS: + return None + + input_usage = metadata.get("input_usage") + compact_input_usage = None + if input_usage is not None: + if input_usage.keys() != _PD_INPUT_USAGE_KEYS: + return None + compact_input_usage = ( + input_usage["input_text_tokens"], + input_usage["input_audio_tokens"], + input_usage["input_image_tokens"], + ) + + return ( + sub_req_id, + text, + metadata["count_output_tokens"], + metadata["prompt_tokens"], + metadata["prompt_cache_len"], + metadata["mtp_accepted_token_num"], + metadata["mtp_verify_token_num"], + metadata["mtp_verify_step_num"], + finish_status.status, + metadata.get("node_mode"), + compact_input_usage, + ) + + +def unpack_pd_compact_token_info(token_info: PDCompactTokenInfo): + ( + sub_req_id, + text, + count_output_tokens, + prompt_tokens, + prompt_cache_len, + mtp_accepted_token_num, + mtp_verify_token_num, + mtp_verify_step_num, + finish_status, + node_mode, + input_usage, + ) = token_info + metadata = { + "count_output_tokens": count_output_tokens, + "prompt_tokens": prompt_tokens, + "prompt_cache_len": prompt_cache_len, + "mtp_accepted_token_num": mtp_accepted_token_num, + "mtp_verify_token_num": mtp_verify_token_num, + "mtp_verify_step_num": mtp_verify_step_num, + } + if node_mode is not None: + metadata["node_mode"] = node_mode + if input_usage is not None: + metadata["input_usage"] = { + "input_text_tokens": input_usage[0], + "input_audio_tokens": input_usage[1], + "input_image_tokens": input_usage[2], + } + return sub_req_id, text, metadata, finish_status @dataclass