feat: support separate DiT/VAE service instances for Flux.

xllm/core/distributed_runtime/comm_channel.cpp

@@ -24,9 +24,10 @@ limitations under the License.
 
 namespace xllm {
 
-bool CommChannel::init_brpc(const std::string& server_address) {
+bool CommChannel::init_brpc(const std::string& server_address,
+                            int32_t timeout_ms) {
   options_.connection_type = "pooled";
-  options_.timeout_ms = -1;
+  options_.timeout_ms = timeout_ms;
   options_.connect_timeout_ms = -1;
   options_.max_retry = 3;
 
@@ -53,14 +54,12 @@ bool CommChannel::hello() {
   return true;
 }
 
-bool CommChannel::check_health() {
+bool CommChannel::check_health(int32_t timeout_ms) {
   proto::Status req;
   proto::Status resp;
   brpc::Controller cntl;
 
-  // Set a timeout for health check
-  // check hang status: 10min(magic num)
-  cntl.set_timeout_ms(600000);
+  cntl.set_timeout_ms(timeout_ms);
   stub_->Hello(&cntl, &req, &resp, nullptr);
   if (cntl.Failed()) {
     LOG(WARNING) << "Health check failed: " << cntl.ErrorText();

xllm/core/distributed_runtime/comm_channel.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include <brpc/controller.h>
 #include <folly/futures/Future.h>
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <vector>
@@ -38,7 +39,7 @@ class CommChannel {
   CommChannel() = default;
   virtual ~CommChannel() = default;
 
-  bool init_brpc(const std::string& server_address);
+  bool init_brpc(const std::string& server_address, int32_t timeout_ms = -1);
 
   virtual bool hello();
 
@@ -115,7 +116,7 @@ class CommChannel {
       folly::Promise<int64_t>& promise);
 
   // Check if the connection to worker is healthy
-  virtual bool check_health();
+  virtual bool check_health(int32_t timeout_ms = 600000);
 
   virtual bool sleep(MasterStatus master_status);
 

xllm/core/distributed_runtime/dit_engine.cpp

@@ -18,20 +18,41 @@ limitations under the License.
 
 #include <glog/logging.h>
 #include <sys/sysinfo.h>
+#include <unistd.h>
+
+#include <chrono>
+#include <exception>
+#include <optional>
+#include <sstream>
+#include <unordered_set>
 
 #include "common/device_monitor.h"
 #include "core/common/global_flags.h"
 #include "core/common/metrics.h"
 #include "core/distributed_runtime/master.h"
+#include "core/framework/config/dit_config.h"
 #include "core/framework/config/execution_config.h"
 #include "core/platform/device.h"
+#include "distributed_runtime/comm_channel.h"
+#include "distributed_runtime/remote_worker.h"
 #include "framework/parallel_state/parallel_args.h"
 #include "framework/parallel_state/parallel_state.h"
 #include "runtime/worker.h"
 #include "util/env_var.h"
 #include "util/timer.h"
 
 namespace xllm {
+
+namespace {
+
+int64_t monotonic_time_ms() {
+  return std::chrono::duration_cast<std::chrono::milliseconds>(
+             std::chrono::steady_clock::now().time_since_epoch())
+      .count();
+}
+
+}  // namespace
+
 DiTEngine::DiTEngine(const runtime::Options& options,
                      std::shared_ptr<DistManager> dist_manager)
     : options_(options), dist_manager_(dist_manager) {
@@ -59,6 +80,7 @@ DiTEngine::DiTEngine(const runtime::Options& options,
   // setup all workers and create worker clients in nnode_rank=0 engine side.
   setup_workers(options);
   worker_clients_num_ = worker_clients_.size();
+  setup_vae_workers();
 
   // init thread pool
   threadpool_ = std::make_unique<ThreadPool>(
@@ -74,6 +96,185 @@ void DiTEngine::setup_workers(const runtime::Options& options) {
   worker_clients_ = dist_manager_->get_worker_clients();
 }
 
+void DiTEngine::setup_vae_workers() {
+  const auto& config = DiTConfig::get_instance();
+  if (config.dit_instance_role() != "dit") {
+    return;
+  }
+
+  const std::string& model_id = options_.model_id();
+  const bool is_flux_model = model_id == "flux" || model_id == "flux-dev" ||
+                             model_id.find("flux-dev-") == 0;
+  CHECK(is_flux_model)
+      << "Separate DiT/VAE instances currently support Flux only, got model: "
+      << model_id;
+
+  CHECK(!config.dit_vae_service_addresses().empty())
+      << "dit_vae_service_addresses must be set for a dit instance.";
+
+  std::stringstream addresses(config.dit_vae_service_addresses());
+  std::unordered_set<std::string> configured_addresses;
+  std::string address;
+  int32_t rank = 0;
+  while (std::getline(addresses, address, ',')) {
+    const size_t first = address.find_first_not_of(" \t\n\r");
+    const size_t last = address.find_last_not_of(" \t\n\r");
+    if (first == std::string::npos) {
+      continue;
+    }
+    address = address.substr(first, last - first + 1);
+    CHECK(configured_addresses.insert(address).second)
+        << "Duplicate VAE service address: " << address;
+    auto channel = std::make_unique<CommChannel>();
+    CHECK(channel->init_brpc(address, config.dit_vae_request_timeout_ms()))
+        << "Failed to connect to VAE service: " << address;
+    auto worker_state = std::make_unique<VaeWorkerState>();
+    worker_state->worker = std::make_shared<RemoteWorker>(
+        rank++, address, options_.devices().front(), std::move(channel));
+    vae_workers_.emplace_back(std::move(worker_state));
+  }
+  CHECK(!vae_workers_.empty())
+      << "No valid VAE service address was configured.";
+  const size_t route_seed = static_cast<size_t>(getpid()) ^
+                            static_cast<size_t>(config.dit_worker_port());
+  next_vae_worker_.store(route_seed, std::memory_order_relaxed);
+  LOG(INFO) << "Configured " << vae_workers_.size()
+            << " VAE service instance(s) for DiT routing.";
+}
+
+size_t DiTEngine::select_vae_worker(
+    const std::vector<bool>& attempted_workers) {
+  CHECK_EQ(attempted_workers.size(), vae_workers_.size());
+
+  const size_t start_index =
+      next_vae_worker_.fetch_add(1, std::memory_order_relaxed) %
+      vae_workers_.size();
+  for (size_t offset = 0; offset < vae_workers_.size(); ++offset) {
+    const size_t worker_index = (start_index + offset) % vae_workers_.size();
+    if (attempted_workers[worker_index]) {
+      continue;
+    }
+    return worker_index;
+  }
+
+  LOG(FATAL) << "No untried VAE worker is available.";
+  return 0;
+}
+
+DiTForwardOutput DiTEngine::decode_with_vae(
+    const DiTForwardInput& input,
+    const DiTForwardOutput& latent_output) {
+  if (latent_output.tensors.empty()) {
+    LOG(ERROR) << "DiT instance returned no latent tensors.";
+    return {};
+  }
+
+  DiTForwardInput vae_input = input;
+  vae_input.prompts.clear();
+  vae_input.prompts_2.clear();
+  vae_input.negative_prompts.clear();
+  vae_input.negative_prompts_2.clear();
+  vae_input.prompt_embeds = torch::Tensor();
+  vae_input.pooled_prompt_embeds = torch::Tensor();
+  vae_input.negative_prompt_embeds = torch::Tensor();
+  vae_input.negative_pooled_prompt_embeds = torch::Tensor();
+  vae_input.images = torch::Tensor();
+  vae_input.images_list.clear();
+  vae_input.mask_images = torch::Tensor();
+  vae_input.control_image = torch::Tensor();
+  vae_input.masked_image_latents = torch::Tensor();
+  vae_input.last_images = torch::Tensor();
+  if (latent_output.tensors.size() == 1) {
+    vae_input.latents = latent_output.tensors.front();
+  } else {
+    vae_input.latents = torch::cat(latent_output.tensors, 0);
+  }
+
+  ForwardInput forward_input;
+  forward_input.input_params.dit_forward_input = std::move(vae_input);
+  std::vector<bool> attempted_workers(vae_workers_.size(), false);
+  const bool debug_print = DiTConfig::get_instance().dit_debug_print();
+  Timer decode_timer;
+  for (size_t offset = 0; offset < vae_workers_.size(); ++offset) {
+    const size_t worker_index = select_vae_worker(attempted_workers);
+    attempted_workers[worker_index] = true;
+    auto& worker_state = *vae_workers_[worker_index];
+    const auto& config = DiTConfig::get_instance();
+    const int64_t now_ms = monotonic_time_ms();
+    if (!worker_state.healthy.load(std::memory_order_relaxed) &&
+        now_ms <
+            worker_state.next_health_check_ms.load(std::memory_order_relaxed)) {
+      continue;
+    }
+    if (!worker_state.healthy.load(std::memory_order_relaxed) &&
+        !worker_state.worker->check_health(
+            config.dit_vae_health_check_timeout_ms())) {
+      worker_state.next_health_check_ms.store(
+          now_ms + config.dit_vae_health_check_interval_ms(),
+          std::memory_order_relaxed);
+      continue;
+    }
+    worker_state.healthy.store(true, std::memory_order_relaxed);
+    worker_state.next_health_check_ms.store(0, std::memory_order_relaxed);
+    Timer rpc_timer;
+    std::optional<RawForwardOutput> result;
+    try {
+      result = vae_workers_[worker_index]
+                   ->worker->step_remote_async(forward_input)
+                   .get();
+    } catch (const std::exception& exception) {
+      worker_state.healthy.store(false, std::memory_order_relaxed);
+      worker_state.next_health_check_ms.store(
+          monotonic_time_ms() + config.dit_vae_health_check_interval_ms(),
+          std::memory_order_relaxed);
+      LOG(WARNING) << "VAE worker " << worker_index
+                   << " threw while decoding latent output: "
+                   << exception.what() << "; trying next worker.";
+      continue;
+    } catch (...) {
+      worker_state.healthy.store(false, std::memory_order_relaxed);
+      worker_state.next_health_check_ms.store(
+          monotonic_time_ms() + config.dit_vae_health_check_interval_ms(),
+          std::memory_order_relaxed);
+      LOG(WARNING) << "VAE worker " << worker_index
+                   << " threw an unknown exception while decoding latent "
+                      "output; trying next worker.";
+      continue;
+    }
+    if (!result.has_value()) {
+      worker_state.healthy.store(false, std::memory_order_relaxed);
+      worker_state.next_health_check_ms.store(
+          monotonic_time_ms() + config.dit_vae_health_check_interval_ms(),
+          std::memory_order_relaxed);
+      LOG(WARNING) << "VAE worker " << worker_index
+                   << " failed to decode latent output, trying next worker.";
+      continue;
+    }
+    const auto& output = result->dit_forward_output;
+    if (output.tensors.size() != input.batch_size) {
+      worker_state.healthy.store(false, std::memory_order_relaxed);
+      worker_state.next_health_check_ms.store(
+          monotonic_time_ms() + config.dit_vae_health_check_interval_ms(),
+          std::memory_order_relaxed);
+      LOG(WARNING) << "VAE worker " << worker_index
+                   << " returned an invalid tensor count: "
+                   << output.tensors.size() << ", expected " << input.batch_size
+                   << ".";
+      continue;
+    }
+    if (debug_print) {
+      LOG(INFO) << "VAE worker " << worker_index
+                << " decode rpc latency: " << rpc_timer.elapsed_seconds()
+                << " s, total latency: " << decode_timer.elapsed_seconds()
+                << " s.";
+    }
+    return output;
+  }
+
+  LOG(ERROR) << "All VAE workers failed to decode latent output.";
+  return {};
+}
+
 bool DiTEngine::init() {
   if (!init_model()) {
     LOG(ERROR) << "Failed to init model from: " << options_.model_path();
@@ -134,10 +335,28 @@ DiTForwardOutput DiTEngine::step(std::vector<DiTBatch>& batches) {
   auto results = folly::collectAll(futures).get();
 
   // return the result from the driver
+  for (const auto& result : results) {
+    if (result.hasException() || !result.value().has_value()) {
+      LOG(ERROR) << "At least one DiT worker failed to execute the request.";
+      batches[0].process_forward_error(
+          Status(StatusCode::UNAVAILABLE,
+                 "A DiT worker failed to execute the request."));
+      return {};
+    }
+  }
   auto forward_output = results.front().value();
-  DCHECK(forward_output.has_value()) << "Failed to execute model";
-  batches[0].process_forward_output(forward_output.value().dit_forward_output);
-  return forward_output.value().dit_forward_output;
+  DiTForwardOutput output = forward_output.value().dit_forward_output;
+  if (DiTConfig::get_instance().dit_instance_role() == "dit") {
+    output = decode_with_vae(dit_forward_input, output);
+    if (output.tensors.empty()) {
+      batches[0].process_forward_error(
+          Status(StatusCode::UNAVAILABLE,
+                 "All configured VAE workers failed to decode the request."));
+      return output;
+    }
+  }
+  batches[0].process_forward_output(output);
+  return output;
 }
 
 std::vector<int64_t> DiTEngine::get_active_activation_memory() const {

xllm/core/distributed_runtime/dit_engine.h

@@ -18,10 +18,13 @@ limitations under the License.
 
 #include <gflags/gflags.h>
 
+#include <atomic>
 #include <memory>
+#include <vector>
 
 #include "common/macros.h"
 #include "dist_manager.h"
+#include "distributed_runtime/remote_worker.h"
 #include "engine.h"
 #include "framework/batch/dit_batch.h"
 #include "framework/parallel_state/process_group.h"
@@ -76,6 +79,10 @@ class DiTEngine : public Engine {
  private:
   // setup workers internal
   void setup_workers(const runtime::Options& options);
+  void setup_vae_workers();
+  size_t select_vae_worker(const std::vector<bool>& attempted_workers);
+  DiTForwardOutput decode_with_vae(const DiTForwardInput& input,
+                                   const DiTForwardOutput& latent_output);
   // init models
   bool init_model();
   // options
@@ -84,6 +91,13 @@ class DiTEngine : public Engine {
   int64_t worker_clients_num_;
   // a list of process groups, with each process group handling a single device
   std::vector<std::unique_ptr<ProcessGroup>> process_groups_;
+  struct VaeWorkerState {
+    std::shared_ptr<RemoteWorker> worker;
+    std::atomic<bool> healthy{true};
+    std::atomic<int64_t> next_health_check_ms{0};
+  };
+  std::vector<std::unique_ptr<VaeWorkerState>> vae_workers_;
+  std::atomic<size_t> next_vae_worker_{0};
 };
 
 }  // namespace xllm

xllm/core/distributed_runtime/remote_worker.cpp

@@ -335,7 +335,9 @@ folly::SemiFuture<int64_t> RemoteWorker::get_active_activation_memory_async() {
   return future;
 }
 
-bool RemoteWorker::check_health() { return channel_->check_health(); }
+bool RemoteWorker::check_health(int32_t timeout_ms) {
+  return channel_->check_health(timeout_ms);
+}
 
 folly::SemiFuture<bool> RemoteWorker::sleep_async(MasterStatus master_status) {
   folly::Promise<bool> promise;

xllm/core/distributed_runtime/remote_worker.h

@@ -139,7 +139,7 @@ class RemoteWorker : public WorkerClient {
   folly::SemiFuture<int64_t> get_active_activation_memory_async() override;
 
   // Check if the connection to worker is healthy
-  bool check_health();
+  bool check_health(int32_t timeout_ms = 600000);
 
   // Get worker global rank
   int32_t global_rank() const { return global_rank_; }