feat(pd): improve lifecycle supervision, health checks, and multimodal support

docs/CN/source/tutorial/api_server_args.rst

@@ -64,6 +64,24 @@ PD 分离模式参数
 
     当 run_mode 设置为 prefill 或 decode 时需要设置此参数
 
+.. option:: --pd_master_mode
+
+    PD Master 拓扑模式，可选值：
+
+    * ``elastic``：Prefill 和 Decode 节点数量可以动态变化（默认）
+    * ``<P>p<D>d``：Prefill 和 Decode 节点数量固定。例如，``2p4d``
+      表示期望恰好注册 2 个 Prefill 节点和 4 个 Decode 节点。
+
+    当 ``run_mode`` 设置为 ``pd_master`` 时使用此参数。
+    在 ``elastic`` 模式下，至少注册一个 Prefill 和一个 Decode 节点后，PD Master 才会
+    进入 ready 状态，节点数量超过一个时仍然保持 ready。使用固定拓扑模式时，只有已注册
+    节点数量与配置完全一致才进入 ready 状态。节点未 ready 时，``/health`` 和 ``/healthz``
+    以及 ``/readiness`` 返回 HTTP 503。在固定拓扑模式下，PD Master 还会并发请求所有已连接
+    Prefill 和 Decode 节点的 ``/health`` 接口；任一节点请求失败、超时或返回非 HTTP 200 时，
+    PD Master 的健康接口都会返回 HTTP 503。无论使用哪种拓扑模式，PD Master 都会同时执行与普通节点类似的
+    推理进度健康检查：当仍有在途请求，且整个 PD Master 连续 ``HEALTH_TIMEOUT`` 秒
+    没有任何请求成功返回 token 时，接口将返回 HTTP 503。
+
 .. option:: --pd_decode_rpyc_port
 
     PD 模式下解码节点用于 kv move manager rpyc 服务器的端口，默认为 ``42000``
@@ -294,10 +312,14 @@ PD 分离模式参数
 
     当输入图片超过该阈值时，LightLLM 会先自动将其缩放到该像素预算内，再继续后续流程。
 
+    多模态 PD 分离模式下，PD Master 与所有 Prefill 节点必须使用相同值，否则 Prefill 注册会被拒绝。
+
 .. option:: --disable_image_resize
 
     禁用对超过 ``--max_image_pixels`` 的图片的自动缩放。默认开启自动缩放。
 
+    多模态 PD 分离模式下，PD Master 与所有 Prefill 节点必须使用相同值。
+
 .. option:: --visual_infer_batch_size
 
     每次推理批次中处理的图像数量，默认为 ``1``

docs/CN/source/tutorial/deepseek_deployment.rst

@@ -154,6 +154,26 @@ LightLLM 支持以下几种部署模式：
 
 PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以更好地利用硬件资源。
 
+.. warning::
+
+    如果宿主机或容器设置了 ``HTTP_PROXY``、``HTTPS_PROXY`` 或 ``ALL_PROXY``，
+    Prefill/Decode 节点到 PD Master 的 WebSocket 注册连接可能被发送到代理服务器。
+    典型日志为 ``proxy rejected connection: HTTP 403``。此时 P/D 节点自身的
+    HTTP 服务可能已正常启动，但 PD Master 仍显示 0 个已注册节点，
+    ``/readiness``、``/health`` 和 ``/healthz`` 会返回 HTTP 503。
+
+    启动服务前，应将所有 PD Master、Prefill、Decode 和 Config Server 的内网 IP
+    或主机名加入 ``NO_PROXY``，并同时设置小写的 ``no_proxy`` 以兼容不同网络库。
+
+    .. code-block:: bash
+
+        export NO_PROXY="${NO_PROXY:+${NO_PROXY},}${pd_master_ip},${host},127.0.0.1,localhost"
+        export no_proxy="$NO_PROXY"
+
+    Kubernetes 部署还应包含实际使用的 Service DNS 名和 Service 后缀，例如
+    ``pd-master.default.svc,.default.svc,.svc``。不要只配置 ``localhost``：跨 Pod 通信使用的
+    Service DNS 或 Pod IP 也必须绕过代理。
+
 3.1 单 PD Master 模式
 ~~~~~~~~~~~~~~~~~~~~~
 
@@ -175,7 +195,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
 
     # PD prefill 模式 for DeepSeek-R1 (DP+EP) on H200
     # 使用方法: sh pd_prefill.sh <host> <pd_master_ip>
-    # 默认使用 NIXL 传输；如需使用 NCCL 数据面，可设置 LIGHTLLM_PD_KV_TRANSPORT_BACKEND=nccl
+    # 默认使用 NCCL 数据面；如需使用 NIXL，请为 prefill 和 decode 统一设置 --pd_trans_mode nixl
     # nvidia-cuda-mps-control -d，运行MPS(可选, 有mps支持性能会好特别多，但是部分显卡和驱动环境开启mps会容易出现错误，建议升级驱动到较高版本，特别是H系列卡)
 
     export host=$1
@@ -184,6 +204,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
+    --pd_trans_mode nccl \
     --tp 8 \
     --dp 8 \
     --host $host \
@@ -202,13 +223,14 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
 
     # PD decode 模式 for DeepSeek-R1 (DP+EP) on H200
     # 使用方法: sh pd_decode.sh <host> <pd_master_ip>
-    # 默认使用 NIXL 传输；如需使用 NCCL 数据面，可设置 LIGHTLLM_PD_KV_TRANSPORT_BACKEND=nccl
+    # 默认使用 NCCL 数据面；如需使用 NIXL，请为 prefill 和 decode 统一设置 --pd_trans_mode nixl
     export host=$1
     export pd_master_ip=$2
     nvidia-cuda-mps-control -d
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
+    --pd_trans_mode nccl \
     --tp 8 \
     --dp 8 \
     --host $host \
@@ -277,6 +299,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
+    --pd_trans_mode nccl \
     --host $host \
     --port 8019 \
     --tp 8 \
@@ -296,6 +319,7 @@ PD (Prefill-Decode) 分离模式将预填充和解码阶段分离部署，可以
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
+    --pd_trans_mode nccl \
     --host $host \
     --port 8121 \
     --nccl_port 12322 \

docs/CN/source/tutorial/multimodal.rst

@@ -52,6 +52,20 @@ LightLLM支持多种多模态模型的推理，下面以InternVL为例，对多
 
 .. note:: 为了使每一个GPU的显存负载相同，需要visual_dp * visual_tp = tp，例如tp=2，则visual_dp=1, visual_tp=2。
 
+PD 分离模式限制
+---------------
+
+在多模态 PD 分离部署中，``pd_master`` 会在计算输入 token 数前预处理图片，并将结果转发给
+``prefill`` 节点。因此，``pd_master`` 和每个 ``prefill`` 节点必须使用相同的
+``--max_image_pixels`` 和 ``--disable_image_resize`` 有效值。
+
+Prefill 注册时会自动校验这两个参数；任意一项不一致都会拒绝注册，避免 PD Master 的 token
+估算与 Prefill 实际图片预处理结果不一致。Decode 节点不负责图片预处理，不受这项限制。
+
+.. note::
+   启动脚本应为 PD Master 和所有多模态 Prefill 节点显式传入相同参数。比较的是启动配置处理后的
+   最终有效值，而不只是命令行中是否出现该选项。
+
 ViT部署方式
 -----------
 

docs/EN/source/tutorial/api_server_args.rst

@@ -64,6 +64,25 @@ PD disaggregation Mode Parameters
 
     This parameter needs to be set when run_mode is set to prefill or decode
 
+.. option:: --pd_master_mode
+
+    PD master topology mode, optional values:
+
+    * ``elastic``: The number of Prefill and Decode nodes can change dynamically (default)
+    * ``<P>p<D>d``: The numbers of Prefill and Decode nodes are fixed. For example,
+      ``2p4d`` expects exactly 2 Prefill nodes and 4 Decode nodes.
+
+    This parameter is used when ``run_mode`` is set to ``pd_master``.
+    In ``elastic`` mode, PD Master becomes ready after at least one Prefill and one Decode node are registered;
+    additional nodes do not make it unready. In a fixed mode, PD Master is ready only when the registered node
+    counts exactly match the configured topology. ``/health``, ``/healthz``, and ``/readiness`` return HTTP 503
+    while the nodes are not ready. In a fixed topology mode, PD Master also requests ``/health`` concurrently
+    from every connected Prefill and Decode node. A failed or timed-out request, or any response other than
+    HTTP 200, makes the PD Master health endpoints return HTTP 503. Independently of the topology mode,
+    PD Master also applies the regular inference-progress health check: while requests remain in flight,
+    the endpoints return HTTP 503 if no request on the PD Master successfully returns a token for
+    ``HEALTH_TIMEOUT`` consecutive seconds.
+
 .. option:: --pd_decode_rpyc_port
 
     Port used by decode nodes for kv move manager rpyc server in PD mode, default is ``42000``
@@ -293,10 +312,14 @@ Multimodal Parameters
 
     If an input image exceeds this threshold, LightLLM automatically resizes it down to this pixel budget before continuing.
 
+    In multimodal PD disaggregation mode, PD Master and every Prefill node must use the same value; otherwise Prefill registration is rejected.
+
 .. option:: --disable_image_resize
 
     Disable automatic resize for images exceeding ``--max_image_pixels``. Resize is enabled by default.
 
+    In multimodal PD disaggregation mode, PD Master and every Prefill node must use the same value.
+
 .. option:: --visual_infer_batch_size
 
     Number of images processed in each inference batch, default is ``1``

docs/EN/source/tutorial/deepseek_deployment.rst

@@ -154,6 +154,27 @@ Suitable for deploying MoE models across multiple nodes.
 
 PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for deployment, which can better utilize hardware resources.
 
+.. warning::
+
+    If the host or container defines ``HTTP_PROXY``, ``HTTPS_PROXY``, or ``ALL_PROXY``, the WebSocket
+    registration connection from a Prefill/Decode node to PD Master may be sent through that proxy.
+    A typical error is ``proxy rejected connection: HTTP 403``. The P/D HTTP servers may already be
+    running normally, while PD Master still reports zero registered nodes and returns HTTP 503 from
+    ``/readiness``, ``/health``, and ``/healthz``.
+
+    Before starting the services, add the internal IP addresses or hostnames of every PD Master, Prefill,
+    Decode, and Config Server to ``NO_PROXY``. Also set lowercase ``no_proxy`` for compatibility with
+    different networking libraries.
+
+    .. code-block:: bash
+
+        export NO_PROXY="${NO_PROXY:+${NO_PROXY},}${pd_master_ip},${host},127.0.0.1,localhost"
+        export no_proxy="$NO_PROXY"
+
+    For Kubernetes, include the Service DNS names and suffixes actually used by the deployment, for example
+    ``pd-master.default.svc,.default.svc,.svc``. Configuring only ``localhost`` is insufficient: the Service
+    DNS names or Pod IP addresses used for cross-Pod communication must also bypass the proxy.
+
 3.1 Single PD Master Mode
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -175,7 +196,7 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
 
     # PD prefill mode for DeepSeek-R1 (DP+EP) on H200
     # Usage: sh pd_prefill.sh <host> <pd_master_ip>
-    # NIXL is used by default. To use NCCL as the data-plane backend, set LIGHTLLM_PD_KV_TRANSPORT_BACKEND=nccl.
+    # NCCL is the default data-plane backend. Set --pd_trans_mode nixl on both prefill and decode to select NIXL.
     # nvidia-cuda-mps-control -d, run MPS (optional, performance will be much better with mps support, but some GPUs may encounter errors when enabling mps, it's recommended to upgrade to a higher driver version, especially for H-series cards)
 
     export host=$1
@@ -184,6 +205,7 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
+    --pd_trans_mode nccl \
     --tp 8 \
     --dp 8 \
     --host $host \
@@ -199,13 +221,14 @@ PD (Prefill-Decode) disaggregation mode separates prefill and decode stages for
 
     # PD decode mode for DeepSeek-R1 (DP+EP) on H200
     # Usage: sh pd_decode.sh <host> <pd_master_ip>
-    # NIXL is used by default. To use NCCL as the data-plane backend, set LIGHTLLM_PD_KV_TRANSPORT_BACKEND=nccl.
+    # NCCL is the default data-plane backend. Set --pd_trans_mode nixl on both prefill and decode to select NIXL.
     export host=$1
     export pd_master_ip=$2
     nvidia-cuda-mps-control -d
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
+    --pd_trans_mode nccl \
     --tp 8 \
     --dp 8 \
     --host $host \
@@ -274,6 +297,7 @@ Supports multiple PD Master nodes, providing better load balancing and high avai
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "prefill" \
+    --pd_trans_mode nccl \
     --host $host \
     --port 8019 \
     --tp 8 \
@@ -293,6 +317,7 @@ Supports multiple PD Master nodes, providing better load balancing and high avai
     LOADWORKER=18 python -m lightllm.server.api_server \
     --model_dir /path/DeepSeek-R1 \
     --run_mode "decode" \
+    --pd_trans_mode nccl \
     --host $host \
     --port 8121 \
     --nccl_port 12322 \

docs/EN/source/tutorial/multimodal.rst

@@ -52,6 +52,24 @@ Advanced Configuration Parameters
 
 .. note:: To ensure equal memory load on each GPU, visual_dp * visual_tp = tp is required. For example, if tp=2, then visual_dp=1, visual_tp=2.
 
+PD Disaggregation Restrictions
+------------------------------
+
+In a multimodal PD deployment, ``pd_master`` preprocesses images before estimating
+the input token count and forwards the result to the ``prefill`` node. Consequently,
+``pd_master`` and every ``prefill`` node must use the same effective values for
+``--max_image_pixels`` and ``--disable_image_resize``.
+
+These two parameters are checked when a Prefill node registers. A mismatch rejects
+the registration, preventing PD Master's token estimate from differing from the
+image preprocessing performed by Prefill. Decode nodes do not preprocess images and
+are not subject to this restriction.
+
+.. note::
+   Launch scripts should explicitly pass identical values to PD Master and every
+   multimodal Prefill node. Registration compares the final effective startup values,
+   not merely whether an option was present on the command line.
+
 ViT Deployment Methods
 -----------------------
 

lightllm/common/basemodel/basemodel.py

@@ -1050,7 +1050,7 @@ def _check_max_len_infer(self):
         # 模拟最大长度进行 prefill，观察是否出现 OOM
         try:
             logger.info("begin check max_len infer")
-            dummy_input_ids = torch.ones(self.batch_max_tokens, dtype=torch.int32, device="cuda")
+            dummy_input_ids = torch.ones(self.batch_max_tokens, dtype=torch.int64, device="cuda")
             b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda")
             mem_indexes = self.mem_manager.alloc(len(dummy_input_ids)).cuda()
             b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda")
@@ -1126,7 +1126,7 @@ def _autotune_warmup(self):
                 rand_gen = torch.Generator(device="cuda")
                 rand_gen.manual_seed(input_len)
                 dummy_input_ids = torch.randint(
-                    0, 10000, (input_len,), dtype=torch.int32, device="cuda", generator=rand_gen
+                    0, 10000, (input_len,), dtype=torch.int64, device="cuda", generator=rand_gen
                 )
                 b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda")
                 mem_indexes = self.mem_manager.alloc(len(dummy_input_ids)).cuda()
@@ -1188,7 +1188,7 @@ def _init_padded_req(self):
         # prefill init padding req.
         prefill_input_len = 1
         batch_size = 1
-        dummy_input_ids = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
+        dummy_input_ids = torch.ones((batch_size,), dtype=torch.int64, device="cuda")
         b_req_idx = torch.tensor(
             [self.req_manager.HOLD_REQUEST_ID for _ in range(batch_size)], dtype=torch.int32, device="cuda"
         )

lightllm/common/basemodel/batch_objs.py

@@ -57,6 +57,7 @@ class ModelInput:
     mtp_draft_input_hiddens: Optional[torch.Tensor] = None
 
     def to_cuda(self):
+        self.check_input()
         if self.input_ids is not None:
             self.input_ids = self.input_ids.cuda(non_blocking=True)
         if self.mem_indexes is None:
@@ -95,6 +96,10 @@ def __post_init__(self):
 
     def check_input(self):
         assert len(self.multimodal_params) == self.batch_size
+        if self.input_ids is not None:
+            assert (
+                self.input_ids.dtype == torch.int64
+            ), f"model input_ids must use torch.int64, got {self.input_ids.dtype}"
 
 
 @dataclass

lightllm/common/basemodel/cuda_graph.py

@@ -217,7 +217,7 @@ def warmup(self, model):
             seq_len = 2
             total_token_num = batch_size * seq_len
             max_len_in_batch = self.graph_max_len_in_batch
-            input_ids = torch.tensor([1 for _ in range(batch_size)], dtype=torch.int32, device="cuda")
+            input_ids = torch.tensor([1 for _ in range(batch_size)], dtype=torch.int64, device="cuda")
             mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda()
             b_req_idx = torch.tensor(
                 [model.req_manager.HOLD_REQUEST_ID for _ in range(batch_size)], dtype=torch.int32, device="cuda"
@@ -276,7 +276,7 @@ def warmup_overlap(self, model):
                 seq_len = 2
                 total_token_num = batch_size * seq_len
                 max_len_in_batch = self.graph_max_len_in_batch
-                input_ids = torch.tensor([1 for _ in range(batch_size)], dtype=torch.int32, device="cuda")
+                input_ids = torch.tensor([1 for _ in range(batch_size)], dtype=torch.int64, device="cuda")
                 mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda()
                 b_req_idx = torch.tensor(
                     [model.req_manager.HOLD_REQUEST_ID for _ in range(batch_size)], dtype=torch.int32, device="cuda"

lightllm/common/basemodel/prefill_cuda_graph.py

@@ -171,7 +171,7 @@ def warmup(self, model):
         for handle_token_num in self.graph_handle_token_nums[::-1]:
             logger.info(f"Capture prefill cudagraph, handle_token_num: {handle_token_num}")
             total_token_num = handle_token_num
-            input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int32, device="cuda")
+            input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int64, device="cuda")
             mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda()
             b_req_idx = torch.tensor([model.req_manager.HOLD_REQUEST_ID], dtype=torch.int32, device="cuda")
             b_seq_len = torch.empty(1, dtype=torch.int32, device="cuda")
@@ -231,7 +231,7 @@ def warmup_overlap(self, model):
             for micro_batch_index in [0, 1]:
                 # dummy prefill, capture the cudagraph
                 total_token_num = handle_token_num
-                input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int32, device="cuda")
+                input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int64, device="cuda")
                 mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda()
                 b_req_idx = torch.tensor([model.req_manager.HOLD_REQUEST_ID], dtype=torch.int32, device="cuda")
                 b_seq_len = torch.empty(1, dtype=torch.int32, device="cuda")

lightllm/common/basemodel/triton_kernel/mtp_utils.py

@@ -253,13 +253,13 @@ def linear_att_mtp_state_index_update(
 
 def test_mtp_verify():
     req_to_next_token_ids = torch.tensor(
-        [[1, 2, -2, -1, -1], [1, 2, 0, -1, -1], [1, 3, 4, 4, 5]], dtype=torch.int32, device="cuda"
+        [[1, 2, -2, -1, -1], [1, 2, 0, -1, -1], [1, 3, 4, 4, 5]], dtype=torch.int64, device="cuda"
     )
     b_req_idx = torch.tensor([0, 0, 2, 2, 2], dtype=torch.int32, device="cuda")
     b_req_mtp_start_loc = torch.tensor([0, 2], dtype=torch.int32, device="cuda")
-    new_next_token_ids = torch.tensor([1, 4, 2, 4, 13], dtype=torch.int32, device="cuda")
+    new_next_token_ids = torch.tensor([1, 4, 2, 4, 13], dtype=torch.int64, device="cuda")
     all_next_token_ids = torch.tensor(
-        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]], dtype=torch.int32, device="cuda"
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]], dtype=torch.int64, device="cuda"
     )
     mtp_accept_len, accepted_index = mtp_verify(
         req_to_next_token_ids, b_req_mtp_start_loc, new_next_token_ids, b_req_idx