feat: support multi machine for offline inference.

docs/en/getting_started/offline_service.md

@@ -31,6 +31,62 @@ llm.finish()
 
 For LLM Beam Search, use `beam_width` as the switch. `top_logprobs` controls the top-k candidate count used for beam expansion at each decode step. If `top_logprobs` is left at its default value, xLLM uses `beam_width` as the top logprob count. Set `top_logprobs` to a value greater than `beam_width` when you want each beam to consider more candidate tokens. This beam-search top-k is different from the sampling cutoff parameter `top_k`. `best_of` is not the Beam Search switch, and this offline LLM guide does not use `num_return_sequences` to control the returned beams.
 
+## Multi-Machine Offline Inference
+
+When a single machine does not have enough devices to load the model, you can use multi-machine offline inference to distribute the model across multiple machines. All machines run the same script and are differentiated by the `nnodes`, `node_rank`, and `master_node_addr` parameters:
+
+- The machine with `node_rank=0` is the **driver**, responsible for submitting inference requests and collecting results
+- Machines with `node_rank>0` are **assistants**, which enter a wait loop after initialization and are scheduled by the driver
+
+### Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `nnodes` | Total number of machines, defaults to 1 (single machine) |
+| `node_rank` | Current machine index, range `[0, nnodes)`, driver is 0 |
+| `master_node_addr` | Communication address of the driver machine in `host:port` format; must be reachable from all machines |
+| `rank_tablefile` | Path to the HCCL collective communication topology file (required for multi-machine) |
+
+### Example
+
+Using 2 machines with 16 devices total, assuming the driver IP is `11.87.191.99`:
+
+**Machine 0 (driver):**
+
+```bash
+python examples/generate.py \
+    --model='/path/models/GLM-5-final-w8a8' \
+    --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \
+    --max_seqs_per_batch=300 \
+    --max_memory_utilization=0.9 \
+    --nnodes=2 \
+    --node_rank=0 \
+    --rank_tablefile=hccl_tools/hccl_2s_16p.json \
+    --master_node_addr="11.87.191.99:19888"
+```
+
+**Machine 1 (assistant):**
+
+```bash
+python examples/generate.py \
+    --model='/path/models/GLM-5-final-w8a8' \
+    --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \
+    --max_seqs_per_batch=300 \
+    --max_memory_utilization=0.9 \
+    --nnodes=2 \
+    --node_rank=1 \
+    --rank_tablefile=hccl_tools/hccl_2s_16p.json \
+    --master_node_addr="11.87.191.99:19888"
+```
+
+### Notes
+
+1. All machines must use the same model path, `devices` configuration, and `nnodes` value
+2. `master_node_addr` must be the actual reachable IP of the driver; loopback addresses (`127.0.0.1`, `localhost`, etc.) are not allowed
+3. `rank_tablefile` must be pre-generated via `hccl_tools` to describe the multi-machine collective communication topology
+4. Assistant machines only need to construct the `LLM` object — there is no need to call `generate()` or define prompts; during `LLM` initialization, when `node_rank>0` is detected, the process automatically enters a wait loop and is shut down remotely when the driver calls `llm.finish()`
+5. Machines can be started in any order; the framework waits internally for all machines to be ready
+
 ## Embedding
 
 Generate embedding example: [:simple-github: https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py](https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py)

docs/zh/getting_started/offline_service.md

@@ -31,6 +31,62 @@ llm.finish()
 
 LLM Beam Search 使用 `beam_width` 作为开启参数。`top_logprobs` 控制每个解码步用于 beam 扩展的 top-k 候选数量。如果 `top_logprobs` 保持默认值，xLLM 会使用 `beam_width` 作为 top logprob 数量。如果希望每个 beam 考虑更多候选 token，可以将 `top_logprobs` 设置为大于 `beam_width` 的值。这里的 beam-search top-k 不同于采样截断参数 `top_k`。`best_of` 不是 Beam Search 开关，本文档也不使用 `num_return_sequences` 来控制 LLM 返回的 beam 数。
 
+## 多机离线推理
+
+当单台机器的设备数量不足以加载模型时，可以使用多机离线推理将模型分布在多个节点上。所有节点运行相同的脚本，通过 `nnodes`、`node_rank` 和 `master_node_addr` 参数区分角色：
+
+- `node_rank=0` 的节点为驱动节点（driver），负责发起推理请求和收集结果
+- `node_rank>0` 的节点为辅助节点（assistant），启动后自动进入等待状态，由驱动节点调度执行
+
+### 参数说明
+
+| 参数 | 说明 |
+|------|------|
+| `nnodes` | 总节点数，默认为 1（单机） |
+| `node_rank` | 当前节点编号，范围 `[0, nnodes)`，驱动节点为 0 |
+| `master_node_addr` | 驱动节点的通信地址，格式为 `host:port`，所有节点必须能访问该地址 |
+| `rank_tablefile` | HCCL 集合通信拓扑文件路径（多机场景必须提供） |
+
+### 使用示例
+
+以 2 机 16 卡为例，假设驱动节点 IP 为 `11.87.191.99`：
+
+**节点 0（驱动节点）：**
+
+```bash
+python examples/generate.py \
+    --model='/path/models/GLM-5-final-w8a8' \
+    --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \
+    --max_seqs_per_batch=300 \
+    --max_memory_utilization=0.9 \
+    --nnodes=2 \
+    --node_rank=0 \
+    --rank_tablefile=hccl_tools/hccl_2s_16p.json \
+    --master_node_addr="11.87.191.99:19888"
+```
+
+**节点 1（辅助节点）：**
+
+```bash
+python examples/generate.py \
+    --model='/path/models/GLM-5-final-w8a8' \
+    --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \
+    --max_seqs_per_batch=300 \
+    --max_memory_utilization=0.9 \
+    --nnodes=2 \
+    --node_rank=1 \
+    --rank_tablefile=hccl_tools/hccl_2s_16p.json \
+    --master_node_addr="11.87.191.99:19888"
+```
+
+### 注意事项
+
+1. 所有节点必须使用相同的模型路径、`devices` 配置和 `nnodes` 参数
+2. `master_node_addr` 必须为驱动节点的真实可达 IP，不能使用 `127.0.0.1`、`localhost` 等回环地址
+3. `rank_tablefile` 需预先通过 `hccl_tools` 生成，描述多机集合通信拓扑
+4. 辅助机器只需构造 `LLM` 对象即可，无需调用 `generate()` 或定义 prompts；`LLM` 初始化时检测到 `node_rank>0` 会自动进入等待循环，直到驱动节点调用 `llm.finish()` 时被远程关闭
+5. 各节点启动顺序无要求，框架内部会自动等待所有节点就绪
+
 ## Embedding
 
 生成Embedding示例：[:simple-github: https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py](https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py)

xllm/__init__.py

@@ -61,8 +61,10 @@ def _load_xllm_export() -> ModuleType:
     "ArgumentParser",
     "Embedding",
     "LLM",
+    "LLMAssistantMaster",
     "LLMMaster",
     "VLM",
+    "VLMAssistantMaster",
     "VLMMaster",
     "Options",
     "SamplingParams",
@@ -104,8 +106,10 @@ def _load_public_api() -> None:
             "ArgumentParser": ArgumentParser,
             "Embedding": Embedding,
             "LLM": LLM,
+            "LLMAssistantMaster": xllm_export.LLMAssistantMaster,
             "LLMMaster": xllm_export.LLMMaster,
             "VLM": VLM,
+            "VLMAssistantMaster": xllm_export.VLMAssistantMaster,
             "VLMMaster": xllm_export.VLMMaster,
             "Options": xllm_export.Options,
             "SamplingParams": SamplingParams,
@@ -139,8 +143,10 @@ def __dir__() -> list[str]:
     "ArgumentParser",
     "Embedding",
     "LLM",
+    "LLMAssistantMaster",
     "LLMMaster",
     "VLM",
+    "VLMAssistantMaster",
     "VLMMaster",
     "Options",
     "SamplingParams",

xllm/core/distributed_runtime/comm_channel.cpp

@@ -70,6 +70,20 @@ bool CommChannel::check_health() {
   return resp.ok();
 }
 
+bool CommChannel::shutdown() {
+  proto::Empty req;
+  proto::Status resp;
+  brpc::Controller cntl;
+
+  cntl.set_timeout_ms(5000);
+  stub_->Shutdown(&cntl, &req, &resp, nullptr);
+  if (cntl.Failed() || !resp.ok()) {
+    LOG(ERROR) << "Shutdown failed: " << cntl.ErrorText();
+    return false;
+  }
+  return true;
+}
+
 bool CommChannel::allocate_kv_cache(const KVCacheShape& kv_cache_shape) {
   if (!kv_cache_shape.has_key_cache_shape() ||
       kv_cache_shape.key_cache_shape().empty()) {

xllm/core/distributed_runtime/comm_channel.h

@@ -122,6 +122,10 @@ class CommChannel {
 
   virtual bool stop_profile();
 
+  // Request the remote worker to terminate. Currently only used by
+  // multi-machine offline inference for coordinated shutdown of all machines.
+  virtual bool shutdown();
+
  protected:
   bool execute_model_with_brpc(
       const ForwardInput& input,

xllm/core/distributed_runtime/dist_manager.cpp

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "distributed_runtime/dist_manager.h"
 
+#include <folly/futures/Future.h>
 #include <glog/logging.h>
 
 #include "comm_channel.h"
@@ -47,11 +48,13 @@ DistManager::DistManager(const runtime::Options& options)
 }
 
 DistManager::~DistManager() {
+  shutdown_remote_workers();
+
   // Stop health check
   HealthCheckManager::instance().stop_health_check_thread();
 
   XllmServer* collective_server =
-      ServerRegistry::get_instance().get_server(server_name_);
+      ServerRegistry::get_instance().try_get_server(server_name_);
   if (collective_server != nullptr) {
     collective_server->stop();
 
@@ -177,6 +180,8 @@ void DistManager::setup_multi_node_workers(
   const int32_t each_node_ranks = static_cast<int32_t>(devices.size());
   const int32_t world_size = each_node_ranks * options.nnodes();
   const int32_t base_rank = options.node_rank() * each_node_ranks;
+  local_rank_start_ = base_rank;
+  local_rank_count_ = each_node_ranks;
   const int32_t dp_size = options.dp_size();
   const int32_t cp_size = options.cp_size();
   const int32_t ep_size = options.ep_size();
@@ -335,4 +340,42 @@ void DistManager::setup_multi_node_workers(
     }
   }
 }
+
+bool DistManager::shutdown_remote_workers() {
+  if (remote_workers_shutdown_) {
+    return true;
+  }
+  remote_workers_shutdown_ = true;
+
+  if (worker_clients_.empty()) {
+    return true;
+  }
+
+  std::vector<folly::SemiFuture<bool>> futures;
+  futures.reserve(worker_clients_.size());
+  const int32_t local_rank_end = local_rank_start_ + local_rank_count_;
+  for (size_t worker_rank = 0; worker_rank < worker_clients_.size();
+       ++worker_rank) {
+    const int32_t global_rank = static_cast<int32_t>(worker_rank);
+    if (global_rank >= local_rank_start_ && global_rank < local_rank_end) {
+      continue;
+    }
+    if (local_rank_count_ > 1 && global_rank % local_rank_count_ != 0) {
+      continue;
+    }
+    futures.emplace_back(worker_clients_[worker_rank]->shutdown_async());
+  }
+  if (futures.empty()) {
+    return true;
+  }
+
+  bool success = true;
+  auto results = folly::collectAll(futures).get();
+  for (const auto& result : results) {
+    if (!result.hasValue() || !result.value()) {
+      success = false;
+    }
+  }
+  return success;
+}
 }  // namespace xllm

xllm/core/distributed_runtime/dist_manager.h

@@ -32,6 +32,11 @@ class DistManager {
     return worker_clients_;
   };
 
+  // Shut down all remote (non-local) worker processes via the Shutdown RPC.
+  // Currently only used by multi-machine offline inference. Sends shutdown
+  // only to workers on other machines; skips local ranks. Idempotent.
+  bool shutdown_remote_workers();
+
  private:
   DISALLOW_COPY_AND_ASSIGN(DistManager);
   void setup_multi_node_workers(const runtime::Options& options,
@@ -57,5 +62,8 @@ class DistManager {
   std::vector<std::unique_ptr<WorkerServer>> servers_;
 
   std::string server_name_;
+  int32_t local_rank_start_ = 0;
+  int32_t local_rank_count_ = 0;
+  bool remote_workers_shutdown_ = false;
 };
 }  // namespace xllm

xllm/core/distributed_runtime/dit_master.cpp

@@ -165,10 +165,8 @@ DiTAssistantMaster::DiTAssistantMaster(const Options& options)
 }
 
 DiTAssistantMaster::~DiTAssistantMaster() {
-  // wait for the loop thread to finish
-  if (loop_thread_.joinable()) {
-    loop_thread_.join();
-  }
+  running_ = false;
+  wait();
 }
 
 void DiTAssistantMaster::run() {
@@ -182,4 +180,10 @@ void DiTAssistantMaster::run() {
   });
 }
 
+void DiTAssistantMaster::wait() {
+  if (loop_thread_.joinable()) {
+    loop_thread_.join();
+  }
+}
+
 }  // namespace xllm

xllm/core/distributed_runtime/dit_master.h

@@ -73,6 +73,7 @@ class DiTAssistantMaster : public Master {
   DiTAssistantMaster(const Options& options);
   ~DiTAssistantMaster();
   void run() override;
+  void wait() override;
 
   static void handle_signal(int signum) { running_ = false; }
 

xllm/core/distributed_runtime/engine.h

@@ -171,6 +171,9 @@ class Engine {
     return false;
   };
 
+  // Shut down remote workers (multi-machine offline inference only).
+  virtual bool shutdown_remote_workers() { return true; };
+
   // XTensor mode: get GlobalXTensor offsets for allocated blocks
   // Returns per-layer K/V offsets for each block
   // Output: offsets[layer_id] = {k_offsets, v_offsets}

xllm/core/distributed_runtime/llm_engine.cpp

@@ -910,6 +910,13 @@ bool LLMEngine::unlink_p2p(const std::vector<std::string>& remote_addrs) {
   return true;
 }
 
+bool LLMEngine::shutdown_remote_workers() {
+  if (!dist_manager_) {
+    return true;
+  }
+  return dist_manager_->shutdown_remote_workers();
+}
+
 ForwardOutput LLMEngine::step(std::vector<Batch>& batch) {
   if (worker_clients_.empty()) {
     // empty worker, return

xllm/core/distributed_runtime/llm_engine.h

@@ -117,6 +117,8 @@ class LLMEngine : public Engine {
 
   bool unlink_p2p(const std::vector<std::string>& remote_addrs) override;
 
+  bool shutdown_remote_workers() override;
+
   std::shared_ptr<DistManager> get_dist_manager() { return dist_manager_; };
 
   bool sleep(MasterStatus master_status) override;

xllm/core/distributed_runtime/llm_master.cpp

@@ -21,6 +21,7 @@ limitations under the License.
 
 #include <atomic>
 #include <boost/algorithm/string.hpp>
+#include <chrono>
 #include <csignal>
 #include <memory>
 #include <thread>
@@ -572,6 +573,10 @@ bool LLMMaster::unlink_p2p(const std::vector<std::string>& remote_addrs) {
   return engine_->unlink_p2p(remote_addrs);
 }
 
+bool LLMMaster::shutdown_remote_workers() {
+  return engine_->shutdown_remote_workers();
+}
+
 LLMAssistantMaster::LLMAssistantMaster(const Options& options)
     : Master(
           options,
@@ -589,10 +594,8 @@ LLMAssistantMaster::LLMAssistantMaster(const Options& options)
 }
 
 LLMAssistantMaster::~LLMAssistantMaster() {
-  // wait for the loop thread to finish
-  if (loop_thread_.joinable()) {
-    loop_thread_.join();
-  }
+  running_ = false;
+  wait();
 }
 
 void LLMAssistantMaster::run() {
@@ -606,6 +609,12 @@ void LLMAssistantMaster::run() {
   });
 }
 
+void LLMAssistantMaster::wait() {
+  if (loop_thread_.joinable()) {
+    loop_thread_.join();
+  }
+}
+
 // ============== Async RL training support: Pause/Resume ==============
 void LLMMaster::pause_scheduler(const std::string& mode) {
   LOG(INFO) << "LLMMaster: pausing scheduler (mode=" << mode << ")";