perf(pd): compact token forwarding packets

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):

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]

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