diff --git a/third_party/torch_npu_ops b/third_party/torch_npu_ops index adc29dc9d2..be50fa1774 160000 --- a/third_party/torch_npu_ops +++ b/third_party/torch_npu_ops @@ -1 +1 @@ -Subproject commit adc29dc9d2e69875524b8a54bfc251463275d66e +Subproject commit be50fa17740dea45d5042b4b8b79ae77f8743429 diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index e1c56661fc..b9bbf45a4b 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -124,6 +124,8 @@ class Options { PROPERTY(int32_t, node_rank) = 0; + PROPERTY(bool, enable_single_process) = false; + PROPERTY(int32_t, dp_size) = 1; PROPERTY(int32_t, cp_size) = 1; diff --git a/xllm/core/distributed_runtime/collective_service.cpp b/xllm/core/distributed_runtime/collective_service.cpp index d76d907483..ce2a6e27e5 100644 --- a/xllm/core/distributed_runtime/collective_service.cpp +++ b/xllm/core/distributed_runtime/collective_service.cpp @@ -20,6 +20,8 @@ limitations under the License. #include +#include "util/net.h" + namespace xllm { CollectiveService::CollectiveService(int dp_group_num, @@ -38,6 +40,27 @@ CollectiveService::CollectiveService(int dp_group_num, root_infos_.push_back(root_info); } #endif + + // Pre-allocate free TCPStore ports for the worst-case communication group + // count. The communicator computes a slot index per group from + // (group_kind, group_index) — see ProcessGroupSlots in + // collective_communicator.cpp — which is bounded by total_num_ slots + // across the world/tp/single_rank/dp_local/moe_tp/moe_ep/encoder_dp + // sections. We over-allocate by a small constant to absorb future group + // additions; the unused ports are released right after binding. + // + // OS-allocated free ports avoid collisions with random outgoing TCP source + // ports inside the local ephemeral range (1024-65535 on common Linux + // hosts), which the previous "master_port + offset" scheme could not. + const int port_slot_count = std::max(1, 4 * total_num_ + 8); + group_ports_.reserve(port_slot_count); + for (int i = 0; i < port_slot_count; ++i) { + const int port = net::get_local_free_port(); + CHECK_GT(port, 0) << "Failed to allocate free TCP port for process-group " + "rendezvous; slot=" + << i; + group_ports_.push_back(port); + } } void CollectiveService::Sync(::google::protobuf::RpcController* controller, @@ -55,6 +78,12 @@ void CollectiveService::Sync(::google::protobuf::RpcController* controller, #if defined(USE_NPU) to_proto_list(root_infos_, response); #endif + // Send the master-allocated ports to every worker. All ranks see the same + // ordered list, so each worker resolves its process-group slots to the + // same TCPStore endpoints the master will bind on. + for (const int port : group_ports_) { + response->add_group_ports(port); + } } std::unordered_map CollectiveService::wait() { diff --git a/xllm/core/distributed_runtime/collective_service.h b/xllm/core/distributed_runtime/collective_service.h index 52eb27c113..fdd975be9b 100644 --- a/xllm/core/distributed_runtime/collective_service.h +++ b/xllm/core/distributed_runtime/collective_service.h @@ -40,6 +40,12 @@ class CollectiveService : public proto::Collective { // wait all worker connected std::unordered_map wait(); + // Master-allocated free TCPStore ports, in the canonical group order shared + // with workers (see CollectiveCommunicator::create_process_groups). The + // master node consumes this list locally; workers receive the same list + // through CommUniqueIdList.group_ports during Sync. + const std::vector& group_ports() const { return group_ports_; } + private: DISALLOW_COPY_AND_ASSIGN(CollectiveService); @@ -57,6 +63,10 @@ class CollectiveService : public proto::Collective { #endif std::mutex mutex_; std::unordered_map addrs_map_; + // Free TCPStore ports the master allocated up-front. Returned verbatim to + // every worker through Sync so each rank uses the same set of ports for + // process-group bootstrap. + std::vector group_ports_; }; } // namespace xllm diff --git a/xllm/core/distributed_runtime/dist_manager.cpp b/xllm/core/distributed_runtime/dist_manager.cpp index eb66ba45d5..265da14fbe 100644 --- a/xllm/core/distributed_runtime/dist_manager.cpp +++ b/xllm/core/distributed_runtime/dist_manager.cpp @@ -30,6 +30,8 @@ limitations under the License. #include "remote_worker.h" #include "runtime/forward_shared_memory_manager.h" #include "runtime/llm_worker_impl.h" +#include "runtime/worker.h" +#include "runtime/worker_client.h" #include "server/xllm_server_registry.h" #include "shm_channel.h" #include "util/net.h" @@ -38,11 +40,22 @@ namespace xllm { DistManager::DistManager(const runtime::Options& options) : server_name_("CollectiveServer") { auto master_node_addr = options.master_node_addr().value_or(""); - if (!master_node_addr.empty()) { + if (options.enable_single_process()) { + // Single-process mode: one OS process drives all local devices (one worker + // thread per device, sharing one HcclCommInitAll world). Independent of + // master_node_addr — single-node single-process needs none. Multi-node + // single-process (nnodes > 1) is not yet supported and is rejected inside + // setup_single_node_workers. + setup_single_node_workers(options); + } else if (!master_node_addr.empty()) { server_name_.append(std::to_string(options.server_idx())); setup_multi_node_workers(options, master_node_addr); } else { - LOG(FATAL) << "master_node_addr is empty."; + LOG(FATAL) + << "master_node_addr is empty and single-process mode is not " + "enabled. Set XLLM_ENABLE_SINGLE_PROCESS=1 for single-process " + "serving, or --master_node_addr for multi-process/multi-node " + "serving."; } } @@ -51,7 +64,7 @@ DistManager::~DistManager() { 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(); @@ -142,6 +155,110 @@ void setup_numa_affinity_and_isolation( } // namespace +void DistManager::setup_single_node_workers(const runtime::Options& options) { + const auto& devices = options.devices(); + CHECK_GT(devices.size(), 0) + << "At least one device is required in single-node mode."; + CHECK_EQ((devices.size() % options.dp_size()), 0) + << "Device size must be divisible by dp size in single-node serving " + "mode."; + CHECK_EQ(options.nnodes(), 1) + << "Multi-node single-process serving (nnodes = " << options.nnodes() + << ") is not yet supported. Run single-node single-process " + "(XLLM_ENABLE_SINGLE_PROCESS=1 with nnodes=1), or use the " + "multi-process launcher (--master_node_addr, one process per device) " + "for multi-node."; + CHECK_EQ(options.node_rank(), 0) + << "Single-process serving requires node_rank=0 (got " + << options.node_rank() << ")."; + + const std::string& backend = options.backend(); + CHECK(backend.empty() || backend == "llm") + << "Single-node serving currently supports only the llm backend, got: " + << backend; + + const int32_t world_size = static_cast(devices.size()); + const int32_t dp_size = options.dp_size(); + const int32_t cp_size = options.cp_size(); + const int32_t ep_size = options.ep_size(); + const int32_t dp_local_tp_size = world_size / dp_size; + + LOG(INFO) << "Single-node serving world_size = " << world_size + << ", dp_size = " << dp_size << ", cp_size = " << cp_size + << ", ep_size = " << ep_size << ", tp_size = " << dp_local_tp_size; + + // initialize process groups if there are multiple devices + if (devices.size() > 1) { +#if defined(USE_NPU) + // Single-process multi-device shares one HCCL world (HcclCommInitAll) and + // routes tensor-parallel collectives through the C++ ProcessGroup. Only the + // TORCH kernel backend uses that path; the ATB backend drives its own + // process-global comm manager, which cannot rendezvous two ranks inside one + // process. Fail fast with actionable guidance instead of hanging. + CHECK_EQ(options.npu_kernel_backend(), "TORCH") + << "Single-node single-process mode with multiple NPU devices requires " + "--npu_kernel_backend=TORCH (got '" + << options.npu_kernel_backend() + << "'). The ATB backend does not support single-process multi-device " + "yet; use the multi-process launcher (set --master_node_addr) for " + "ATB."; +#endif + process_groups_ = parallel_state::create_local_process_groups( + devices, /*options=*/options); + } + + // For data parallelism with dp_size > 1 and dp_size < world_size, build + // dp-local process groups so each dp group has its own communication + // channel. Mirrors the old setup_single_node_workers behaviour. + if (dp_size > 1 && dp_size < world_size) { + dp_local_process_groups_.reserve(dp_size); + for (int32_t dp_rank = 0; dp_rank < dp_size; ++dp_rank) { + auto begin_it = devices.begin(); + std::advance(begin_it, dp_rank * dp_local_tp_size); + auto end_it = devices.begin(); + std::advance(end_it, (dp_rank + 1) * dp_local_tp_size); + std::vector dp_local_devices(begin_it, end_it); + dp_local_process_groups_.emplace_back( + parallel_state::create_local_process_groups(dp_local_devices, + options)); + } + } + + WorkerType worker_type = + (options.task_type() == "generate") ? WorkerType::LLM : WorkerType::ELM; + + for (size_t i = 0; i < devices.size(); ++i) { + const int32_t rank = static_cast(i); + ProcessGroup* pg = (world_size > 1) ? process_groups_[i].get() : nullptr; + ProcessGroup* dp_local_pg = + (dp_size > 1 && dp_size < world_size) + ? dp_local_process_groups_[i / dp_local_tp_size] + [i % dp_local_tp_size] + .get() + : nullptr; + + ParallelArgs parallel_args(rank, + world_size, + dp_size, + cp_size, + /*process_group=*/pg, + ep_size); + parallel_args.dp_local_process_group_ = dp_local_pg; + // For single-node single-process mode the language-model layers expect a + // tensor-parallel group. With dp=ep=cp=1 the global process group also + // serves as the tp/cp/single-rank group. + parallel_args.tp_group_ = pg; + parallel_args.cp_group_ = pg; + parallel_args.single_rank_group_ = pg; + parallel_args.moe_tp_group_ = pg; + + workers_.emplace_back(std::make_unique( + parallel_args, devices[i], options, worker_type)); + worker_clients_.emplace_back( + std::make_shared(workers_.back().get(), options)); + } +} + void DistManager::setup_multi_node_workers( const runtime::Options& options, const std::string& master_node_addr) { diff --git a/xllm/core/distributed_runtime/dist_manager.h b/xllm/core/distributed_runtime/dist_manager.h index f8474bfd4b..cb6f15ce75 100755 --- a/xllm/core/distributed_runtime/dist_manager.h +++ b/xllm/core/distributed_runtime/dist_manager.h @@ -34,6 +34,7 @@ class DistManager { private: DISALLOW_COPY_AND_ASSIGN(DistManager); + void setup_single_node_workers(const runtime::Options& options); void setup_multi_node_workers(const runtime::Options& options, const std::string& master_node_addr); diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 388ef59a6e..cf62419ea9 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -70,9 +70,6 @@ LLMEngine::LLMEngine(const runtime::Options& options, InterruptionBus::get_instance().subscribe([this](bool interrupted) { this->layer_forward_interrupted_ = interrupted; }); - auto master_node_addr = options.master_node_addr().value_or(""); - CHECK(!master_node_addr.empty()) - << " LLM need to set master node addr, Please set --master_node_addr."; const auto& devices = options_.devices(); // initialize device monitor DeviceMonitor::get_instance().initialize(devices); diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 07861b6018..f1adf38ea7 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -199,16 +199,25 @@ Master::Master(const Options& options, EngineType type) master_status_(options.master_status()) { const auto model_path = std::filesystem::path(options_.model_path()).lexically_normal(); + const auto visible_devices = DeviceNameUtils::parse_devices("auto"); + std::vector devices; + if (options_.enable_single_process()) { + // Single-process serving instead drives every visible card from one process + // (one worker thread per device), so it takes all visible devices and its + // world size is the device count rather than the node count. + devices = visible_devices; + } else { // Multi-process serving runs one worker per process. Select one runtime // logical device from the process-visible devices while keeping node_rank as // the global distributed identity. const int32_t visible_device_count = Platform::device_count(); const int32_t device_idx = DeviceNameUtils::get_device_idx( options_.node_rank(), options_.nnodes(), visible_device_count); - const auto visible_devices = DeviceNameUtils::parse_devices("auto"); - const std::vector devices = {visible_devices[device_idx]}; - // World size is the node count (one worker per process). - const int32_t global_world_size = options_.nnodes(); + devices = {visible_devices[device_idx]}; + } + + int32_t global_world_size = static_cast(devices.size()) * options_.nnodes(); + std::string cp_model_type; if (options_.cp_size() > 1 && Platform::uses_model_cp_partition()) { cp_model_type = util::get_model_type(model_path, options_.backend()); @@ -325,6 +334,7 @@ Master::Master(const Options& options, EngineType type) .master_node_addr(options.master_node_addr()) .nnodes(options.nnodes()) .node_rank(options.node_rank()) + .enable_single_process(options.enable_single_process()) .dp_size(options.dp_size()) .ep_size(options.ep_size()) .max_tokens_per_batch(options_.max_tokens_per_batch()) @@ -383,6 +393,7 @@ Master::Master(const Options& options, EngineType type) .master_node_addr(options.master_node_addr()) .nnodes(options.nnodes()) .node_rank(options.node_rank()) + .enable_single_process(options.enable_single_process()) .dp_size(options.dp_size()) .ep_size(options.ep_size()) .cp_size(options_.cp_size()) @@ -437,6 +448,7 @@ Master::Master(const Options& options, EngineType type) .master_node_addr(options_.master_node_addr()) .nnodes(options_.nnodes()) .node_rank(options_.node_rank()) + .enable_single_process(options_.enable_single_process()) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) @@ -503,6 +515,7 @@ Master::Master(const Options& options, EngineType type) .master_node_addr(options_.master_node_addr()) .nnodes(options_.nnodes()) .node_rank(options_.node_rank()) + .enable_single_process(options_.enable_single_process()) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) @@ -540,6 +553,7 @@ Master::Master(const Options& options, EngineType type) .output_shm_size(options_.output_shm_size() * 1024 * 1024) .is_local(options_.is_local()) .node_rank(options_.node_rank()) + .enable_single_process(options_.enable_single_process()) .enable_schedule_overlap(options_.enable_schedule_overlap()) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) diff --git a/xllm/core/distributed_runtime/worker_server.cpp b/xllm/core/distributed_runtime/worker_server.cpp index d1632ff92c..96641c267e 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -176,6 +176,16 @@ void WorkerServer::create_server(const runtime::Options& options, comm = std::move(common_comm); } + // Forward master-allocated TCPStore ports (if any) so the communicator + // binds on OS-chosen free ports instead of fixed master_port + offset + // values that can collide with random outgoing TCP source ports inside + // the local ephemeral range. + std::vector group_ports(uids.group_ports().begin(), + uids.group_ports().end()); + if (!group_ports.empty()) { + comm->set_group_ports(std::move(group_ports)); + } + comm->create_process_groups(master_node_addr, device); parallel_args = comm->parallel_args(); diff --git a/xllm/core/framework/config/distributed_config.cpp b/xllm/core/framework/config/distributed_config.cpp index 26d19816c9..39927c5b03 100644 --- a/xllm/core/framework/config/distributed_config.cpp +++ b/xllm/core/framework/config/distributed_config.cpp @@ -17,11 +17,13 @@ limitations under the License. #include "core/common/global_flags.h" #include "core/framework/config/config_utils.h" +#include "core/util/env_var.h" DEFINE_string(master_node_addr, - "127.0.0.1:19888", + "", "The master address for multi-node distributed serving(e.g. " - "10.18.1.1:9999)."); + "10.18.1.1:9999). Required for multi-process/multi-node serving; " + "leave empty only when XLLM_ENABLE_SINGLE_PROCESS is set."); DEFINE_string( xtensor_master_node_addr, @@ -53,6 +55,11 @@ void DistributedConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(xtensor_master_node_addr); XLLM_CONFIG_ASSIGN_FROM_FLAG(nnodes); XLLM_CONFIG_ASSIGN_FROM_FLAG(node_rank); + // Single-process mode is opt-in via environment variable rather than a + // gflag, so it can be set from the launch environment without threading a + // CLI flag through every entrypoint. + enable_single_process(util::get_bool_env("XLLM_ENABLE_SINGLE_PROCESS", + enable_single_process())); XLLM_CONFIG_ASSIGN_FROM_FLAG(etcd_addr); XLLM_CONFIG_ASSIGN_FROM_FLAG(etcd_namespace); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_service_routing); @@ -66,6 +73,7 @@ void DistributedConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(nnodes); // don't read rank-related config // XLLM_CONFIG_ASSIGN_FROM_JSON(node_rank); + // enable_single_process is sourced from the environment, not JSON. XLLM_CONFIG_ASSIGN_FROM_JSON(etcd_addr); XLLM_CONFIG_ASSIGN_FROM_JSON(etcd_namespace); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_service_routing); @@ -84,6 +92,8 @@ void DistributedConfig::append_config_json( // don't dump rank-related config // APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( // config_json, default_config, node_rank); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_single_process); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, etcd_addr); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/distributed_config.h b/xllm/core/framework/config/distributed_config.h index 1856c25315..5b3659a5b6 100644 --- a/xllm/core/framework/config/distributed_config.h +++ b/xllm/core/framework/config/distributed_config.h @@ -52,7 +52,7 @@ class DistributedConfig final { return kOptionCategory; } - PROPERTY(std::string, master_node_addr) = "127.0.0.1:19888"; + PROPERTY(std::string, master_node_addr) = ""; PROPERTY(std::string, xtensor_master_node_addr) = "127.0.0.1:19889"; @@ -60,6 +60,13 @@ class DistributedConfig final { PROPERTY(int32_t, node_rank) = 0; + // Run all local devices inside a single OS process (one worker thread per + // device, sharing one HcclCommInitAll world) instead of one process per + // device. Independent of master_node_addr: single-node single-process needs + // no master_node_addr, while a future multi-node single-process deployment + // would still use master_node_addr for inter-node rendezvous. + PROPERTY(bool, enable_single_process) = false; + PROPERTY(std::string, etcd_addr); PROPERTY(std::string, etcd_namespace); diff --git a/xllm/core/framework/parallel_state/CMakeLists.txt b/xllm/core/framework/parallel_state/CMakeLists.txt index 26bdab465d..f286d13180 100644 --- a/xllm/core/framework/parallel_state/CMakeLists.txt +++ b/xllm/core/framework/parallel_state/CMakeLists.txt @@ -15,6 +15,7 @@ cc_library( process_group.h $<$:npu_rank_table_env.h> $<$:npu_process_group.h> + $<$:single_process/hccl_process_group.h> $<$:mlu_process_group.h> $<$:musa_process_group.h> $<$,$>:cuda_process_group.h> @@ -33,6 +34,7 @@ cc_library( process_group.cpp $<$:npu_rank_table_env.cpp> $<$:npu_process_group.cpp> + $<$:single_process/hccl_process_group.cpp> collective_communicator.cpp dit_collective_communicator.cpp DEPS diff --git a/xllm/core/framework/parallel_state/collective_communicator.cpp b/xllm/core/framework/parallel_state/collective_communicator.cpp index ab8eb11e2e..648746598d 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.cpp +++ b/xllm/core/framework/parallel_state/collective_communicator.cpp @@ -304,22 +304,63 @@ void CollectiveCommunicator::create_process_groups( int32_t port; net::parse_host_port_from_addr(master_addr, host, port); + // Compute section bases for the master-allocated port list. Every rank + // computes the same offsets given (world_size, dp_size, ep_size, + // enable_mm_encoder_dp), so the (section, group_index) -> slot mapping is + // identical on master and workers. group_ports_[slot] is then used as the + // TCPStore port for the corresponding process group, which sidesteps the + // EADDRINUSE race that the old "master_port + offset" scheme hits inside + // the OS ephemeral range. + // + // When group_ports_ is empty (legacy callers / DiT) we fall back to the + // master_port + offset formula via use_master_ports below, preserving + // byte-for-byte behaviour for those paths. + const bool use_master_ports = !group_ports_.empty(); + const bool encoder_dp_enabled = + ::xllm::ParallelConfig::get_instance().enable_mm_encoder_dp(); + const int32_t tp_size_pre = world_size / dp_size; + const int32_t moe_tp_size_pre = world_size / ep_size; + + const int32_t kEncoderDpBase = 0; + const int32_t kWorldBase = encoder_dp_enabled ? dp_size : 0; + const int32_t kTpBase = kWorldBase + 1; + const int32_t kSingleRankBase = kTpBase + dp_size; + const int32_t kDpLocalBase = + kSingleRankBase + (tp_size_pre > 1 ? world_size : 0); + const int32_t kMoeTpBase = kDpLocalBase + (dp_size > 1 ? tp_size_pre : 0); + const int32_t kMoeEpBase = kMoeTpBase + (ep_size > 1 ? ep_size : 0); + + auto resolve_port = + [&](int32_t section_base, int32_t group_index, int32_t legacy_port) { + if (!use_master_ports) { + return legacy_port; + } + const int32_t slot = section_base + group_index; + CHECK_LT(slot, static_cast(group_ports_.size())) + << "Master-allocated TCPStore port slot " << slot + << " out of range (size " << group_ports_.size() << ")."; + return group_ports_[slot]; + }; + int32_t port_offset = 0; // Encoder DP is used by multi-modal models to parallelize vision encoder // work inside each language-model TP group. The rank set matches the TP // group, but each rank runs a full encoder on different multi-modal items. - if (::xllm::ParallelConfig::get_instance().enable_mm_encoder_dp()) { + if (encoder_dp_enabled) { const int32_t encoder_dp_size = world_size / dp_size; - port_offset = global_rank / encoder_dp_size + 1; - encoder_dp_group_ = create_process_group(global_rank, - world_size, - encoder_dp_size, - port + port_offset, - false, - host, - "encoder_dp_group", - device); + const int32_t encoder_dp_group_index = global_rank / encoder_dp_size; + port_offset = encoder_dp_group_index + 1; + encoder_dp_group_ = create_process_group( + global_rank, + world_size, + encoder_dp_size, + resolve_port( + kEncoderDpBase, encoder_dp_group_index, port + port_offset), + /*trans=*/false, + host, + "encoder_dp_group", + device); parallel_args_->encoder_dp_group_ = encoder_dp_group_.get(); port += dp_size; } @@ -333,8 +374,8 @@ void CollectiveCommunicator::create_process_groups( process_group_ = create_process_group(global_rank, world_size, world_size, - ++port, - false, + resolve_port(kWorldBase, 0, ++port), + /*trans=*/false, host, "world_group", device); @@ -342,7 +383,8 @@ void CollectiveCommunicator::create_process_groups( int32_t tp_size = world_size / dp_size; CHECK_EQ(tp_size * dp_size, world_size); - port_offset = global_rank / tp_size + 1; + const int32_t tp_group_index = global_rank / tp_size; + port_offset = tp_group_index + 1; std::string tp_host = host; #if defined(USE_NPU) if (::xllm::KernelConfig::get_instance().npu_kernel_backend() == "TORCH" && @@ -351,14 +393,15 @@ void CollectiveCommunicator::create_process_groups( tp_host = get_rank_table_server_host(tp_group_start, host); } #endif - tp_group_ = create_process_group(global_rank, - world_size, - tp_size, - port + port_offset, - false, - tp_host, - "tp_group", - device); + tp_group_ = create_process_group( + global_rank, + world_size, + tp_size, + resolve_port(kTpBase, tp_group_index, port + port_offset), + /*trans=*/false, + tp_host, + "tp_group", + device); parallel_args_->tp_group_ = tp_group_.get(); // Single-rank group is used for modules that don't need tensor parallel (TP) // communication. This avoids unnecessary communication. When tp_size > 1, @@ -375,8 +418,11 @@ void CollectiveCommunicator::create_process_groups( global_rank, world_size, 1, - port + dp_size + single_rank_group_port_gap + global_rank + 1, - false, + resolve_port( + kSingleRankBase, + global_rank, + port + dp_size + single_rank_group_port_gap + global_rank + 1), + /*trans=*/false, host, "single_rank_group", device); @@ -392,15 +438,17 @@ void CollectiveCommunicator::create_process_groups( port += dp_size + single_rank_group_port_gap + single_rank_group_count; if (dp_size > 1) { - port_offset = global_rank % tp_size + 1; - dp_local_process_group_ = create_process_group(global_rank, - world_size, - dp_size, - port + port_offset, - true, - host, - "dp_group", - device); + const int32_t dp_local_group_index = global_rank % tp_size; + port_offset = dp_local_group_index + 1; + dp_local_process_group_ = create_process_group( + global_rank, + world_size, + dp_size, + resolve_port(kDpLocalBase, dp_local_group_index, port + port_offset), + /*trans=*/true, + host, + "dp_group", + device); parallel_args_->dp_local_process_group_ = dp_local_process_group_.get(); port += tp_size; } @@ -410,7 +458,8 @@ void CollectiveCommunicator::create_process_groups( if (ep_size == 1) { parallel_args_->moe_tp_group_ = process_group_.get(); } else { - port_offset = global_rank / moe_tp_size + 1; + const int32_t moe_tp_group_index = global_rank / moe_tp_size; + port_offset = moe_tp_group_index + 1; std::string moe_tp_host = host; #if defined(USE_NPU) if (::xllm::KernelConfig::get_instance().npu_kernel_backend() == "TORCH") { @@ -419,30 +468,32 @@ void CollectiveCommunicator::create_process_groups( moe_tp_host = get_rank_table_server_host(moe_tp_group_start, host); } #endif - moe_tp_group_ = create_process_group(global_rank, - world_size, - moe_tp_size, - port + port_offset, - false, - moe_tp_host, - "moe_tp_group", - device); + moe_tp_group_ = create_process_group( + global_rank, + world_size, + moe_tp_size, + resolve_port(kMoeTpBase, moe_tp_group_index, port + port_offset), + /*trans=*/false, + moe_tp_host, + "moe_tp_group", + device); parallel_args_->moe_tp_group_ = moe_tp_group_.get(); port += ep_size; - port_offset = global_rank % moe_tp_size + 1; - moe_ep_group_ = create_process_group(global_rank, - world_size, - ep_size, - port + port_offset, - true, - host, - "moe_ep_group", - device); + const int32_t moe_ep_group_index = global_rank % moe_tp_size; + port_offset = moe_ep_group_index + 1; + moe_ep_group_ = create_process_group( + global_rank, + world_size, + ep_size, + resolve_port(kMoeEpBase, moe_ep_group_index, port + port_offset), + /*trans=*/true, + host, + "moe_ep_group", + device); parallel_args_->moe_ep_group_ = moe_ep_group_.get(); port += moe_tp_size; } - const int32_t tp_group_index = global_rank / tp_size; parallel_args_->python_tp_rendezvous_host_ = host; parallel_args_->python_tp_rendezvous_port_ = port + tp_group_index + 1; CHECK_LE(parallel_args_->python_tp_rendezvous_port_, 65535) diff --git a/xllm/core/framework/parallel_state/collective_communicator_base.h b/xllm/core/framework/parallel_state/collective_communicator_base.h index e86b986250..3e1f08c4fd 100644 --- a/xllm/core/framework/parallel_state/collective_communicator_base.h +++ b/xllm/core/framework/parallel_state/collective_communicator_base.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include "parallel_args.h" #include "process_group.h" @@ -37,12 +38,23 @@ class CollectiveCommunicatorBase { virtual const ParallelArgs* parallel_args() = 0; + // Master-broadcast TCPStore ports for process-group rendezvous. When set, + // create_process_groups consumes these instead of computing ports from + // master_addr (which lives inside the OS ephemeral range and can collide + // with random outgoing TCP source ports on a busy host). Empty list is + // legal — implementations fall back to the legacy offset scheme so + // single-node-single-process and DiT remain unchanged. + void set_group_ports(std::vector ports) { + group_ports_ = std::move(ports); + } + int get_global_rank() const { return global_rank_; } int get_world_size() const { return world_size_; } protected: int global_rank_; int world_size_; + std::vector group_ports_; }; } // namespace xllm diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 4f1162859b..0461c10ad9 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -20,6 +20,9 @@ limitations under the License. #include "util/net.h" #if defined(USE_NPU) +#include + +#include "framework/parallel_state/single_process/hccl_process_group.h" #include "hccl/hccl.h" #include "npu_process_group.h" #endif @@ -456,9 +459,36 @@ std::vector> create_local_process_groups( process_groups.reserve(devices.size()); #if defined(USE_NPU) + // Single-process multi-device path: use HcclCommInitAll so every device + // shares one HCCL world. Two c10d_npu::ProcessGroupHCCL instances coexisting + // inside the same process can dispatch collectives in mismatched orders and + // deadlock the local_tp_group under load; HcclCommInitAll-backed groups + // sidestep that by sharing a single HCCL bootstrap. + std::vector device_idxs; + device_idxs.reserve(devices.size()); + for (const auto& device : devices) { + device_idxs.push_back(device.index()); + } std::vector comms(devices.size()); + // HcclCommInitAll requires a device bound on the calling thread to refresh + // the runtime device logic id; the engine thread that runs this has not + // bound one yet (workers bind their own device on their own threads later). + // Bind through c10_npu::SetDevice rather than raw aclrtSetDevice so + // torch_npu's thread-local device cache stays in sync; a raw aclrtSetDevice + // would desync that cache and make later c10_npu::SetDevice calls skip the + // real device switch, leaving worker streams in the wrong context. + int32_t prev_device = c10_npu::current_device(); + CHECK_EQ(c10_npu::SetDevice(static_cast(device_idxs[0])), + ACL_SUCCESS) + << "c10_npu::SetDevice failed for device " << device_idxs[0] + << " before HcclCommInitAll"; + HCCLCHECK(HcclCommInitAll(world_size, device_idxs.data(), comms.data())); + if (prev_device >= 0) { + CHECK_EQ(c10_npu::SetDevice(prev_device), ACL_SUCCESS) + << "c10_npu::SetDevice failed to restore device " << prev_device; + } for (int32_t i = 0; i < world_size; ++i) { - process_groups.emplace_back(std::make_unique( + process_groups.emplace_back(std::make_unique( /*rank=*/i, world_size, devices[i], comms[i])); } #elif defined(USE_CUDA) || defined(USE_MLU) || defined(USE_ILU) || \ diff --git a/xllm/core/framework/parallel_state/single_process/hccl_process_group.cpp b/xllm/core/framework/parallel_state/single_process/hccl_process_group.cpp new file mode 100644 index 0000000000..6d7bb6b965 --- /dev/null +++ b/xllm/core/framework/parallel_state/single_process/hccl_process_group.cpp @@ -0,0 +1,374 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "framework/parallel_state/single_process/hccl_process_group.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "platform/device.h" + +namespace xllm { + +namespace { + +#define XLLM_HCCLCHECK(cmd) \ + do { \ + HcclResult r = (cmd); \ + if (r != HCCL_SUCCESS) { \ + LOG(FATAL) << "HCCL error " << static_cast(r) << " at " \ + << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +// HcclCommInitAll binds each comm to the device context active on the thread +// that created it (the engine thread, with device 0 current). The comms are +// then used from each worker's own thread, where the device must be bound +// eagerly before launching an HCCL op; torch_npu's set_device is lazy and the +// raw HCCL launch is not a torch op that would trigger the deferred bind, so we +// bind explicitly via c10_npu::SetDevice. Without this the runtime reports +// "stream not in current context" and aborts. +void bind_device(const torch::Device& device) { + CHECK_EQ(c10_npu::SetDevice(static_cast(device.index())), + ACL_SUCCESS) + << "c10_npu::SetDevice failed for device " << device.index(); +} + +HcclDataType to_hccl_data_type(const torch::Tensor& input) { + switch (input.scalar_type()) { + case torch::kFloat: + return HCCL_DATA_TYPE_FP32; + case torch::kHalf: + return HCCL_DATA_TYPE_FP16; + case torch::kDouble: + return HCCL_DATA_TYPE_FP64; + case torch::kLong: + return HCCL_DATA_TYPE_INT64; + case torch::kInt: + return HCCL_DATA_TYPE_INT32; + case torch::kChar: + return HCCL_DATA_TYPE_INT8; + case torch::kByte: + return HCCL_DATA_TYPE_UINT8; + case torch::kBool: + return HCCL_DATA_TYPE_UINT8; + case torch::kBFloat16: + return HCCL_DATA_TYPE_BFP16; + default: + LOG(FATAL) << "Unsupported tensor dtype for HCCL: " + << input.scalar_type(); + } +} + +// Lightweight c10d::Work that waits on a single NPU event recorded after the +// HCCL launch. The compute path always queues HCCL on an NPU stream and then +// issues blocking waits via `wait()`, so this is enough to surface completion +// to the caller without taking on the full ProcessGroupHCCL state machine. +// Mirrors CudaNcclWork from the CUDA single-process group. +class HcclNpuWork : public c10d::Work { + public: + HcclNpuWork(c10d::OpType op_type, + c10_npu::NPUStream stream, + const torch::Device& device) + : c10d::Work(/*rank=*/-1, op_type), device_(device) { + event_.record(stream); + } + + bool isCompleted() override { + c10_npu::NPUGuard guard(device_); + return event_.query(); + } + + bool wait(std::chrono::milliseconds /*timeout*/ = kNoTimeout) override { + c10_npu::NPUGuard guard(device_); + // Make the current stream wait for completion of the HCCL work that was + // recorded on the issuing stream. This matches the semantics callers get + // from c10d_npu::ProcessGroupHCCL where wait() is a stream-side join, not a + // host-side synchronize. + event_.block(c10_npu::getCurrentNPUStream(device_.index())); + return true; + } + + private: + c10_npu::NPUEvent event_; + torch::Device device_; +}; + +c10::intrusive_ptr make_work(c10d::OpType op_type, + c10_npu::NPUStream stream, + const torch::Device& device) { + return c10::make_intrusive(op_type, stream, device); +} + +} // namespace + +HcclProcessGroup::HcclProcessGroup(int32_t rank, + int32_t world_size, + const torch::Device& device, + HcclComm comm) + : ProcessGroup(rank, world_size, device), comm_(comm) { + CHECK(comm != nullptr) << "HcclComm must be non-null"; +} + +HcclProcessGroup::~HcclProcessGroup() { + // We own the comm; release it before the base class destructor runs so any + // device state can be torn down cleanly. + bind_device(device()); + HcclCommDestroy(comm_); + comm_ = nullptr; + Device::empty_cache(device().index()); +} + +c10_npu::NPUStream HcclProcessGroup::hccl_stream() { + return c10_npu::getCurrentNPUStream(device().index()); +} + +void HcclProcessGroup::allreduce(torch::Tensor& input) { + allreduce_async(input)->wait(); +} + +c10::intrusive_ptr HcclProcessGroup::allreduce_async( + torch::Tensor& input) { + CHECK_EQ(input.device(), device()) + << "allreduce input must live on the process group's device"; + CHECK(input.is_contiguous()) << "allreduce input must be contiguous"; + + bind_device(device()); + c10_npu::NPUStream stream = hccl_stream(); + XLLM_HCCLCHECK(HcclAllReduce(input.data_ptr(), + input.data_ptr(), + static_cast(input.numel()), + to_hccl_data_type(input), + HCCL_REDUCE_SUM, + comm_, + stream.stream())); + return make_work(c10d::OpType::ALLREDUCE, stream, device()); +} + +void HcclProcessGroup::allgather(const torch::Tensor& input, + std::vector& outputs) { + allgather_async(input, outputs)->wait(); +} + +c10::intrusive_ptr HcclProcessGroup::allgather_async( + const torch::Tensor& input, + std::vector& outputs) { + CHECK_EQ(static_cast(outputs.size()), world_size()) + << "allgather output count must equal world_size"; + CHECK_EQ(input.device(), device()) + << "allgather input must live on the process group's device"; + CHECK(input.is_contiguous()) << "allgather input must be contiguous"; + + bind_device(device()); + c10_npu::NPUStream stream = hccl_stream(); + // HcclAllGather requires a single contiguous receive buffer; build one + // sized [world_size, *input.shape] then copy into the per-rank outputs + // afterwards. This mirrors the ProcessGroup::allgather + cat semantics + // callers expect. + std::vector stacked_shape; + stacked_shape.reserve(input.dim() + 1); + stacked_shape.push_back(world_size()); + for (int64_t s : input.sizes()) { + stacked_shape.push_back(s); + } + torch::Tensor stacked = torch::empty(stacked_shape, input.options()); + XLLM_HCCLCHECK(HcclAllGather(input.data_ptr(), + stacked.data_ptr(), + static_cast(input.numel()), + to_hccl_data_type(input), + comm_, + stream.stream())); + // Slice stacked into the supplied outputs. Each slice shares storage with + // stacked, so the gather kernel only writes once. + for (int32_t i = 0; i < world_size(); ++i) { + if (!outputs[i].defined()) { + outputs[i] = stacked.select(0, i); + } else { + outputs[i].copy_(stacked.select(0, i), /*non_blocking=*/true); + } + } + return make_work(c10d::OpType::ALLGATHER, stream, device()); +} + +c10::intrusive_ptr HcclProcessGroup::allgather_base_async( + const torch::Tensor& input, + torch::Tensor& output) { + CHECK_EQ(input.device(), device()) + << "allgather_base input must live on the process group's device"; + CHECK(output.defined()) << "allgather_base output must be preallocated"; + CHECK_EQ(output.device(), device()) + << "allgather_base output must live on the process group's device"; + CHECK(output.is_contiguous()) << "allgather_base output must be contiguous"; + + torch::Tensor input_buf = input.contiguous(); + CHECK_EQ(output.numel(), input_buf.numel() * world_size()) + << "allgather_base output size must equal world_size * input size"; + + bind_device(device()); + c10_npu::NPUStream stream = hccl_stream(); + XLLM_HCCLCHECK(HcclAllGather(input_buf.data_ptr(), + output.data_ptr(), + static_cast(input_buf.numel()), + to_hccl_data_type(input_buf), + comm_, + stream.stream())); + return make_work(c10d::OpType::_ALLGATHER_BASE, stream, device()); +} + +torch::Tensor HcclProcessGroup::allgather_base_sync( + const torch::Tensor& input) { + CHECK_EQ(input.device(), device()) + << "allgather_base input must live on the process group's device"; + std::vector out_shape; + out_shape.reserve(input.dim() + 1); + out_shape.push_back(world_size()); + for (int64_t s : input.sizes()) { + out_shape.push_back(s); + } + torch::Tensor output = torch::empty(out_shape, input.options()); + allgather_base_async(input, output)->wait(); + return output; +} + +void HcclProcessGroup::reduce_scatter(const torch::Tensor& input, + torch::Tensor& output) { + CHECK(input.is_contiguous()) << "reduce_scatter input must be contiguous"; + CHECK_EQ(input.device(), device()) + << "reduce_scatter input must live on the process group's device"; + CHECK(output.defined()) << "reduce_scatter output must be defined"; + CHECK_EQ(output.device(), device()) + << "reduce_scatter output must live on the process group's device"; + CHECK_EQ(input.numel(), output.numel() * world_size()) + << "reduce_scatter input size must equal world_size * output size"; + + bind_device(device()); + c10_npu::NPUStream stream = hccl_stream(); + XLLM_HCCLCHECK(HcclReduceScatter(input.data_ptr(), + output.data_ptr(), + static_cast(output.numel()), + to_hccl_data_type(input), + HCCL_REDUCE_SUM, + comm_, + stream.stream())); + // Block the caller until the scatter completes so the semantics match the + // base class (which calls ->wait()). + make_work(c10d::OpType::REDUCE_SCATTER, stream, device())->wait(); +} + +void HcclProcessGroup::all_to_all_single( + torch::Tensor output, + torch::Tensor input, + std::vector output_split_sizes, + std::vector input_split_sizes, + bool async_op, + c10::intrusive_ptr* async_work) { + CHECK(output.defined()) << "all_to_all_single output must be defined"; + CHECK(input.defined()) << "all_to_all_single input must be defined"; + CHECK_EQ(input.device(), device()) + << "all_to_all_single input must live on the process group's device"; + CHECK_EQ(output.device(), device()) + << "all_to_all_single output must live on the process group's device"; + + // Treat complex tensors the same way the base class does: split each + // complex element into real+imag along the last dim before sending. + if (input.is_complex()) { + input = torch::view_as_real(input); + } + if (output.is_complex()) { + output = torch::view_as_real(output); + } + CHECK(input.is_contiguous()) << "all_to_all_single input must be contiguous"; + CHECK(output.is_contiguous()) + << "all_to_all_single output must be contiguous"; + + const int32_t ws = world_size(); + std::vector in_splits = input_split_sizes; + std::vector out_splits = output_split_sizes; + if (in_splits.empty()) { + CHECK_EQ(input.size(0) % ws, 0) + << "input dim 0 must be divisible by world_size for equal-split a2a"; + in_splits.assign(ws, input.size(0) / ws); + } + if (out_splits.empty()) { + CHECK_EQ(output.size(0) % ws, 0) + << "output dim 0 must be divisible by world_size for equal-split a2a"; + out_splits.assign(ws, output.size(0) / ws); + } + CHECK_EQ(static_cast(in_splits.size()), ws); + CHECK_EQ(static_cast(out_splits.size()), ws); + + // HCCL has no direct alltoall-with-splits primitive; emulate via a batch of + // send/recv items submitted together through HcclBatchSendRecv (the NPU + // analogue of ncclGroupStart/ncclSend/ncclRecv/ncclGroupEnd). + const int64_t in_inner = input.numel() / std::max(1, input.size(0)); + const int64_t out_inner = + output.numel() / std::max(1, output.size(0)); + const HcclDataType dtype = to_hccl_data_type(input); + const size_t elem_size = static_cast(input.element_size()); + char* input_ptr = static_cast(input.data_ptr()); + char* output_ptr = static_cast(output.data_ptr()); + + std::vector items; + items.reserve(static_cast(ws) * 2); + int64_t in_offset = 0; + int64_t out_offset = 0; + for (int32_t r = 0; r < ws; ++r) { + const uint64_t send_count = static_cast(in_splits[r] * in_inner); + const uint64_t recv_count = + static_cast(out_splits[r] * out_inner); + if (send_count > 0) { + items.push_back(HcclSendRecvItem{HCCL_SEND, + input_ptr + in_offset * elem_size, + send_count, + dtype, + static_cast(r)}); + } + if (recv_count > 0) { + items.push_back(HcclSendRecvItem{HCCL_RECV, + output_ptr + out_offset * elem_size, + recv_count, + dtype, + static_cast(r)}); + } + in_offset += in_splits[r] * in_inner; + out_offset += out_splits[r] * out_inner; + } + + bind_device(device()); + c10_npu::NPUStream stream = hccl_stream(); + XLLM_HCCLCHECK(HcclBatchSendRecv(items.data(), + static_cast(items.size()), + comm_, + stream.stream())); + + c10::intrusive_ptr work = + make_work(c10d::OpType::ALLTOALL_BASE, stream, device()); + if (async_op) { + CHECK(async_work != nullptr) << "async_work must be provided for async_op"; + *async_work = work; + } else { + work->wait(); + } +} + +} // namespace xllm diff --git a/xllm/core/framework/parallel_state/single_process/hccl_process_group.h b/xllm/core/framework/parallel_state/single_process/hccl_process_group.h new file mode 100644 index 0000000000..b4baa21d02 --- /dev/null +++ b/xllm/core/framework/parallel_state/single_process/hccl_process_group.h @@ -0,0 +1,77 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include "framework/parallel_state/process_group.h" +#include "hccl/hccl.h" + +namespace xllm { + +// Single-process multi-device process group backed by a raw HcclComm +// (created via HcclCommInitAll). Owns the comm and destroys it in the +// destructor. This is the NPU analogue of NcclProcessGroup: it deliberately +// bypasses c10d_npu::ProcessGroupHCCL because two ProcessGroupHCCL instances +// coexisting inside one process can dispatch collectives in mismatched orders +// and deadlock under load. All collectives are queued on the device's current +// NPU stream so they chain naturally with model-forward kernels. +class HcclProcessGroup final : public ProcessGroup { + public: + explicit HcclProcessGroup(int32_t rank, + int32_t world_size, + const torch::Device& device, + HcclComm comm); + + ~HcclProcessGroup() override; + + void allreduce(torch::Tensor& input) override; + c10::intrusive_ptr allreduce_async(torch::Tensor& input) override; + + void allgather(const torch::Tensor& input, + std::vector& outputs) override; + c10::intrusive_ptr allgather_async( + const torch::Tensor& input, + std::vector& outputs) override; + + c10::intrusive_ptr allgather_base_async( + const torch::Tensor& input, + torch::Tensor& output) override; + torch::Tensor allgather_base_sync(const torch::Tensor& input) override; + + void reduce_scatter(const torch::Tensor& input, + torch::Tensor& output) override; + + void all_to_all_single( + torch::Tensor output, + torch::Tensor input, + std::vector output_split_sizes = {}, + std::vector input_split_sizes = {}, + bool async_op = false, + c10::intrusive_ptr* async_work = nullptr) override; + + private: + // Returns the NPU stream that HCCL ops should run on. We piggy-back on the + // current stream of the device so that ops chain naturally with model + // forward kernels (matches c10d_npu::ProcessGroupHCCL behaviour from the + // caller's perspective). + c10_npu::NPUStream hccl_stream(); + + HcclComm comm_ = nullptr; +}; + +} // namespace xllm diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index dd18133f8d..236b0d1ef1 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -232,7 +232,24 @@ std::optional LLMWorkerImpl::step_for_schedule_overlap( mutable_params.parallel.has_initial_state); } #endif - return execute_no_sync_on_stream(input, *compute_stream_); + std::optional output = + execute_no_sync_on_stream(input, *compute_stream_); +#if defined(USE_NPU) + // Single-process schedule overlap publishes this step's sample outputs to + // last_step_output_ for consumption by other threads (the engine-facing + // device->host copy in forward_output_to_raw runs on a folly executor thread + // with no NPU stream binding, and cannot be ordered against compute_stream_). + // The forward ran NO_SYNC, so without a barrier here those consumers race the + // still-running compute-stream kernels and read reused device memory, which + // aborts with an NPU MTE out-of-range error. Multi-process serializes through + // the RPC/shm response so it never hit this and keeps its overlap unchanged. + // Sync compute_stream_ before the output leaves this device-bound worker + // thread; the next step's prepare still overlaps on prepare_stream_. + if (options_.enable_single_process() && compute_stream_ != nullptr) { + compute_stream_->synchronize(); + } +#endif + return output; } ForwardInput diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 97c8c7d732..c06e1b3158 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -108,6 +108,10 @@ struct Options { // the node_rank of current worker process at. PROPERTY(int32_t, node_rank) = 0; + // run all local devices inside a single OS process (one worker thread per + // device) instead of one process per device. + PROPERTY(bool, enable_single_process) = false; + // data parallelism size, currently mainly used for MoE model // default set as 1 for non-MoE model PROPERTY(int32_t, dp_size) = 1; diff --git a/xllm/core/runtime/params_utils.cpp b/xllm/core/runtime/params_utils.cpp index ffcacdf98b..b9f6bf8fe3 100644 --- a/xllm/core/runtime/params_utils.cpp +++ b/xllm/core/runtime/params_utils.cpp @@ -77,6 +77,154 @@ void proto_to_forward_output(const proto::ForwardOutput& pb_output, COUNTER_ADD(proto_latency_seconds_proto2o, timer.elapsed_seconds()); } +namespace { +// Decode one already-CPU-resident Token plus its optional embedding vector into +// a RawToken. Mirrors the field-by-field mapping the brpc path performs across +// forward_output_to_proto -> proto_to_forward_output (logprob optionality, +// top-k slices, per-token hidden-state embeddings), but without the proto. +RawToken decode_raw_token(const Token& token, + const torch::Tensor& token_embeddings) { + RawToken raw_token; + raw_token.id = token.id; + raw_token.logprob = token.logprob; + raw_token.top_tokens.assign(token.top_tokens.begin(), token.top_tokens.end()); + raw_token.top_logprobs.assign(token.top_logprobs.begin(), + token.top_logprobs.end()); + if (token_embeddings.defined()) { + const float* ptr = token_embeddings.data_ptr(); + raw_token.embeddings.assign( + ptr, ptr + static_cast(token_embeddings.size(0))); + } + return raw_token; +} +} // namespace + +void forward_output_to_raw(const ForwardOutput& forward_output, + RawForwardOutput& raw_forward_output) { + // Convert an in-process ForwardOutput directly into a RawForwardOutput. + // Single-node single-process mode has no transport between worker and engine, + // so instead of the wasteful forward_output_to_proto -> + // proto_to_forward_output serialize/deserialize round-trip, we decode the + // host-side tensors straight into the raw vectors. The decode reuses + // build_token (the same primitive the brpc path calls internally), so the + // result is byte-identical to what a RemoteWorker would have produced over + // brpc. + const auto& sample_output = forward_output.sample_output; + const auto& beam_search_output = forward_output.beam_search_output; + + auto to_cpu = [](const torch::Tensor& t) -> torch::Tensor { + if (!t.defined() || t.device().is_cpu()) { + return t; + } + return t.to(torch::kCPU); + }; + + torch::Tensor next_tokens = to_cpu(sample_output.next_tokens); + torch::Tensor logprobs = to_cpu(sample_output.logprobs); + torch::Tensor top_tokens = to_cpu(sample_output.top_tokens); + torch::Tensor top_logprobs = to_cpu(sample_output.top_logprobs); + torch::Tensor embeddings = to_cpu(sample_output.embeddings); + torch::Tensor expert_load_data = to_cpu(forward_output.expert_load_data); + torch::Tensor src_seq_idxes = to_cpu(beam_search_output.src_seq_idxes); + torch::Tensor out_tokens = to_cpu(beam_search_output.out_tokens); + torch::Tensor out_logprobs = to_cpu(beam_search_output.out_logprobs); + + // Per-sequence token outputs. num_seqs and the 1D/2D branch selection mirror + // forward_output_to_proto exactly. + int32_t num_seqs = next_tokens.size(0); + if (embeddings.defined() && embeddings.numel() > 0) { + num_seqs = std::max(num_seqs, static_cast(embeddings.size(0))); + } + raw_forward_output.outputs.reserve(num_seqs); + for (int32_t output_idx = 0; output_idx < num_seqs; ++output_idx) { + RawSampleOutput raw_seq_out; + if (next_tokens.dim() == 2) { + // Speculative decoding: one row per sequence, multiple tokens per row. + const auto curr_next_tokens = next_tokens[output_idx]; + const auto curr_logprobs = + logprobs.defined() ? logprobs[output_idx] : logprobs; + const auto curr_top_tokens = + top_tokens.defined() ? top_tokens[output_idx] : top_tokens; + const auto curr_top_logprobs = + top_logprobs.defined() ? top_logprobs[output_idx] : top_logprobs; + const auto curr_embeddings = + embeddings.defined() ? embeddings[output_idx] : embeddings; + + const int32_t num_tokens = curr_next_tokens.size(0); + raw_seq_out.tokens.reserve(num_tokens); + for (int32_t i = 0; i < num_tokens; ++i) { + const auto token = build_token(i, + curr_next_tokens, + curr_logprobs, + curr_top_tokens, + curr_top_logprobs); + if (token.id == -1) { + break; + } + const auto token_embeddings = + curr_embeddings.defined() ? curr_embeddings[i] : curr_embeddings; + raw_seq_out.tokens.emplace_back( + decode_raw_token(token, token_embeddings)); + } + } else { + // Single token per sequence; embedding-only requests emit a -1 + // placeholder token that still carries its embedding (no early break, + // matching proto). + Token token(-1); + if (next_tokens.defined() && next_tokens.numel() > 0) { + token = build_token( + output_idx, next_tokens, logprobs, top_tokens, top_logprobs); + } + const auto token_embeddings = + embeddings.defined() ? embeddings[output_idx] : embeddings; + raw_seq_out.tokens.emplace_back( + decode_raw_token(token, token_embeddings)); + } + raw_forward_output.outputs.emplace_back(std::move(raw_seq_out)); + } + + // EPLB expert load + prepared_layer_id. The proto path only serializes these + // under enable_eplb(), so prepared_layer_id decodes to the proto3 default (0) + // when EPLB is off — preserve that exactly. + if (::xllm::EPLBConfig::get_instance().enable_eplb()) { + raw_forward_output.prepared_layer_id = forward_output.prepared_layer_id; + if (expert_load_data.defined()) { + const auto flattened = expert_load_data.view({-1}).contiguous(); + const int64_t* ptr = flattened.data_ptr(); + raw_forward_output.expert_load_data.assign( + ptr, ptr + static_cast(flattened.size(0))); + } + } else { + raw_forward_output.prepared_layer_id = 0; + } + + // Beam search kernel output (flattened 1D tensors, assumed contiguous as in + // the proto path). + if (src_seq_idxes.defined() && src_seq_idxes.numel() > 0) { + const int32_t* ptr = src_seq_idxes.data_ptr(); + raw_forward_output.src_seq_idxes.assign( + ptr, ptr + static_cast(src_seq_idxes.numel())); + } + if (out_tokens.defined() && out_tokens.numel() > 0) { + const int32_t* ptr = out_tokens.data_ptr(); + raw_forward_output.out_tokens.assign( + ptr, ptr + static_cast(out_tokens.numel())); + } + if (out_logprobs.defined() && out_logprobs.numel() > 0) { + const float* ptr = out_logprobs.data_ptr(); + raw_forward_output.out_logprobs.assign( + ptr, ptr + static_cast(out_logprobs.numel())); + } + + // DiT images are already torch tensors on both sides, so the proto round-trip + // was pure overhead: just move them to host. + raw_forward_output.dit_forward_output.tensors.reserve( + forward_output.dit_forward_output.tensors.size()); + for (const auto& t : forward_output.dit_forward_output.tensors) { + raw_forward_output.dit_forward_output.tensors.emplace_back(to_cpu(t)); + } +} + void forward_output_to_proto(const torch::Tensor& next_tokens, const torch::Tensor& logprobs, const torch::Tensor& top_tokens, diff --git a/xllm/core/runtime/params_utils.h b/xllm/core/runtime/params_utils.h index 6af85f5452..09d96445ff 100644 --- a/xllm/core/runtime/params_utils.h +++ b/xllm/core/runtime/params_utils.h @@ -39,6 +39,14 @@ void packed_proto_to_forward_input( void proto_to_forward_output(const proto::ForwardOutput& pb_output, RawForwardOutput& raw_forward_output); +// Convert an in-process ForwardOutput into a RawForwardOutput by reading +// the host-side tensors of the SampleOutput. Mirrors the logic that +// WorkerService::ExecuteModel would have produced for a remote worker, but +// avoids the proto round-trip when the engine and worker live in the same +// process (single-node single-process mode). +void forward_output_to_raw(const ForwardOutput& forward_output, + RawForwardOutput& raw_forward_output); + void forward_output_to_proto(const torch::Tensor& next_tokens, const torch::Tensor& logprobs, const torch::Tensor& top_tokens, diff --git a/xllm/core/runtime/worker_client.cpp b/xllm/core/runtime/worker_client.cpp index 0759598e4e..98ea1cc14b 100644 --- a/xllm/core/runtime/worker_client.cpp +++ b/xllm/core/runtime/worker_client.cpp @@ -16,6 +16,7 @@ limitations under the License. #include "worker_client.h" #include +#include #include #include #include @@ -28,10 +29,31 @@ limitations under the License. #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_input_params.h" #include "framework/state_dict/state_dict.h" +#include "runtime/params_utils.h" #include "util/timer.h" +#if defined(USE_CUDA) || defined(USE_DCU) +#include +#endif + namespace xllm { +WorkerClient::WorkerClient(Worker* w, const runtime::Options& options) + : worker_(w), + options_(options), + dispatch_threadpool_( + std::make_unique(/*num_threads=*/1, + /*cpu_binding=*/false, + /*pool_name=*/"WorkerClient.dispatch")) { + // Pin the dispatch thread to the worker's device so any CUDA driver calls + // issued during prepare (cudaHostAlloc, stream guards, etc.) operate on + // the right context without competing with the engine thread. + if (worker_ != nullptr) { + const Device dev = Device(worker_->device()); + dev.set_device(); + } +} + bool WorkerClient::init_model(const std::string& model_weights_path, int32_t random_seed, MasterStatus master_status) { @@ -106,11 +128,107 @@ folly::SemiFuture> WorkerClient::step_async( return worker_->step_async(input); } +void WorkerClient::build_fake_overlap_output( + const ForwardInput& input, + RawForwardOutput& raw_output) const { + // Mirror WorkerService::step's overlap branch fake-token construction. + // The number of placeholder tokens equals the number of sampled sequences + // for this step (decode + any prefill-sampled seqs), read from + // sampling_params.sample_idxes. Packed engine inputs carry the layout in + // the host buffer, so unpack first when sample_idxes is not materialized. + int32_t num_samples = 0; + if (input.sampling_params.sample_idxes.defined()) { + num_samples = + static_cast(input.sampling_params.sample_idxes.size(0)); + } else if (input.input_host_buffer_has_layout) { + ForwardInput unpacked_input; + const bool unpacked = detail::unpack_from_input_host_buffer( + input, torch::Device(torch::kCPU), unpacked_input); + if (unpacked && unpacked_input.sampling_params.sample_idxes.defined()) { + num_samples = static_cast( + unpacked_input.sampling_params.sample_idxes.size(0)); + } + } + + raw_output.outputs.clear(); + raw_output.outputs.reserve(num_samples); + for (int32_t i = 0; i < num_samples; ++i) { + RawSampleOutput sample_output; + RawToken token; + // Negative 1-based index placeholder: -(i+1). On the next step, + // update_input_by_last_step_output replaces it with next_tokens[i]. + token.id = -(static_cast(i) + 1); + sample_output.tokens.emplace_back(std::move(token)); + raw_output.outputs.emplace_back(std::move(sample_output)); + } + raw_output.prepared_layer_id = -1; +} + folly::SemiFuture> WorkerClient::step_remote_async(const ForwardInput& input) { - LOG(FATAL) << "WorkerClient Method step_remote_async with ForwardInput " - "param is UnImplemented."; - return folly::makeSemiFuture(std::optional(std::nullopt)); + // Single-node single-process path: dispatch the step on the worker and + // convert the resulting ForwardOutput into a RawForwardOutput so callers + // (e.g. LLMEngine::step) get the same shape they would receive from a + // RemoteWorker over brpc. + if (worker_ == nullptr) { + LOG(FATAL) << "WorkerClient Method step_remote_async with ForwardInput " + "param is UnImplemented."; + return folly::makeSemiFuture(std::optional(std::nullopt)); + } + + // Schedule-overlap path: the engine pipelines step N+1 while step N's GPU + // work is still in flight, and fetches step N's real result on the next + // iteration via get_last_step_result_async. We MUST NOT block the returned + // future on the worker's forward completion here: LLMEngine::step does + // collectAll(futures).get(), and the consumer that unblocks the worker's + // producer-consumer cv (update_last_step_result -> get_last_step_result) is + // only dispatched by the scheduler AFTER engine_->step() returns. Waiting on + // the forward here would deadlock the worker's cv against the engine thread. + // + // Instead, mirror multi-process WorkerService::step: kick the forward off + // fire-and-forget on the dispatch thread (it runs, records its result, and + // satisfies the cv handshake against the separate get_last_step path), and + // immediately resolve with a fake-token RawForwardOutput. The real tokens + // are picked up next iteration. + if (options_.enable_schedule_overlap()) { + RawForwardOutput fake_output; + build_fake_overlap_output(input, fake_output); + dispatch_threadpool_->schedule([this, input]() mutable { + worker_->step_async(input) + .via(folly::getGlobalCPUExecutor()) + .thenValue([](std::optional&& /*unused*/) {}); + }); + return folly::makeSemiFuture( + std::optional(std::move(fake_output))); + } + + folly::Promise> promise; + auto future = promise.getSemiFuture(); + // Non-overlap path: run the entire prepare+step kickoff on the per-worker + // dispatch thread. WorkerImpl::step_async runs prepare_work_before_execute + // synchronously on the calling thread, which can issue blocking CUDA driver + // calls. If we ran it on the engine thread, worker N+1's prepare would not + // start until worker N's prepare finished — meanwhile worker N's step kernel + // is already busy-waiting on NCCL collectives that need worker N+1, which + // deadlocks both GPUs. Dispatching here lets every worker's prepare run in + // parallel on its own thread. + dispatch_threadpool_->schedule([this, + input, + promise = std::move(promise)]() mutable { + worker_->step_async(input) + .via(folly::getGlobalCPUExecutor()) + .thenValue([promise = std::move(promise)]( + std::optional&& forward_output) mutable { + if (!forward_output.has_value()) { + promise.setValue(std::nullopt); + return; + } + RawForwardOutput raw; + forward_output_to_raw(forward_output.value(), raw); + promise.setValue(std::optional(std::move(raw))); + }); + }); + return future; } folly::SemiFuture WorkerClient::process_group_test_async() { @@ -198,7 +316,29 @@ const torch::Device& WorkerClient::device() const { return worker_->device(); } folly::SemiFuture> WorkerClient::get_last_step_result_async() { - return folly::makeSemiFuture(std::optional(std::nullopt)); + // Same single-node single-process pattern as step_remote_async: bridge the + // worker's ForwardOutput into the engine-facing RawForwardOutput shape. + if (worker_ == nullptr) { + return folly::makeSemiFuture(std::optional(std::nullopt)); + } + folly::Promise> promise; + auto future = promise.getSemiFuture(); + dispatch_threadpool_->schedule([this, + promise = std::move(promise)]() mutable { + worker_->get_last_step_result_async() + .via(folly::getGlobalCPUExecutor()) + .thenValue([promise = std::move(promise)]( + std::optional&& forward_output) mutable { + if (!forward_output.has_value()) { + promise.setValue(std::nullopt); + return; + } + RawForwardOutput raw; + forward_output_to_raw(forward_output.value(), raw); + promise.setValue(std::optional(std::move(raw))); + }); + }); + return future; } folly::SemiFuture> diff --git a/xllm/core/runtime/worker_client.h b/xllm/core/runtime/worker_client.h index 8186527ccc..e56c326f0e 100644 --- a/xllm/core/runtime/worker_client.h +++ b/xllm/core/runtime/worker_client.h @@ -37,7 +37,7 @@ namespace xllm { class WorkerClient { public: WorkerClient() = default; - explicit WorkerClient(Worker* w) : worker_(w) {} + WorkerClient(Worker* w, const runtime::Options& options); virtual ~WorkerClient() = default; // initialize model, cache manager. blocking call @@ -153,7 +153,33 @@ class WorkerClient { virtual folly::SemiFuture get_active_activation_memory_async(); private: + // Build the fake-token RawForwardOutput that schedule-overlap expects in + // place of the real (still in-flight) forward result. Mirrors the + // multi-process WorkerService::step overlap branch: one output per sampled + // sequence carrying a single negative placeholder token (id = -(i+1)), + // which WorkerImpl::update_input_by_last_step_output replaces with the real + // token of the previous step on the next iteration. + void build_fake_overlap_output(const ForwardInput& input, + RawForwardOutput& raw_output) const; + Worker* worker_ = nullptr; // not owend + + // Runtime options for the in-process worker. Used to read + // enable_schedule_overlap() from the SAME source WorkerImpl branches on + // (options_.enable_schedule_overlap_), so the client and worker never + // disagree about whether the producer-consumer overlap pipeline is active. + runtime::Options options_; + + // Dedicated single-thread pool for dispatching the in-process worker step + // pipeline. The engine fans out one step per worker; without a per-worker + // dispatch thread, the engine thread would synchronously serialise each + // worker's `prepare_work_before_execute` (which can issue blocking CUDA + // calls like `cudaHostAlloc`). With multi-device single-process NCCL the + // first worker's step starts its collectives and busy-waits for the + // partner; if the partner is still queued behind the engine-thread + // prepare, the GPUs deadlock. Running prepare+dispatch on this dedicated + // thread per WorkerClient lets all workers begin their steps in parallel. + std::unique_ptr dispatch_threadpool_; }; } // namespace xllm diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 5fa33ae766..9ccf85d045 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -1165,31 +1165,39 @@ void WorkerImpl::execute_cuda_block_copy_kernel( folly::SemiFuture> WorkerImpl::step_async( const ForwardInput& input) { - ForwardInput input_on_device; - - prepare_work_before_execute(input, input_on_device); - folly::Promise> promise; auto future = promise.getSemiFuture(); - threadpool_.schedule([this, - input = std::move(input_on_device), - promise = std::move(promise)]() mutable { + // Prepare runs on the worker's own threadpool (the same thread that + // executes step). Doing prepare on the calling thread instead serialises + // every worker's input H2D copy through that caller, and in single-node + // single-process mode each prepare can issue blocking device driver calls + // through the caching host allocator. While one worker's step is busy-waiting + // on a collective for its partner, those driver calls share the driver's + // internal locks with the running collective and the caller never returns to + // schedule the second worker's step — both devices deadlock. Pinning prepare + // to the worker's device thread keeps all device work on a single thread. + threadpool_.schedule([this, input, promise = std::move(promise)]() mutable { + ForwardInput input_on_device; + prepare_work_before_execute(input, input_on_device); + if (hierarchy_kv_cache_transfer_ != nullptr) { - hierarchy_kv_cache_transfer_->set_layer_synchronizer(input.input_params); + hierarchy_kv_cache_transfer_->set_layer_synchronizer( + input_on_device.input_params); } // run the model on the given input in working thread if (!enable_schedule_overlap()) { - const auto output = this->step(input); + const auto output = this->step(input_on_device); promise.setValue(output); } else { - if (last_step_output_valid_ && input.token_ids.numel() > 0 && - input.input_params.meta.batch_forward_type.has_decode()) { + if (last_step_output_valid_ && input_on_device.token_ids.numel() > 0 && + input_on_device.input_params.meta.batch_forward_type.has_decode()) { // replace step i model input with true output of step i-1 - input = update_input_by_last_step_output_for_schedule_overlap(input); + input_on_device = update_input_by_last_step_output_for_schedule_overlap( + input_on_device); } - const auto output = this->step_for_schedule_overlap(input); + const auto output = this->step_for_schedule_overlap(input_on_device); if (output.has_value()) { if (is_driver() || ::xllm::EPLBConfig::get_instance().enable_eplb()) { std::unique_lock lock(mtx_); diff --git a/xllm/proto/common.proto b/xllm/proto/common.proto index f86f70222e..f9d71483bf 100644 --- a/xllm/proto/common.proto +++ b/xllm/proto/common.proto @@ -119,6 +119,14 @@ message CommUniqueId { message CommUniqueIdList { repeated CommUniqueId comm_unique_ids = 1; + + // Master-allocated TCPStore ports for each communication group, in a + // canonical order shared between the master and every worker. Workers + // consume this list when bootstrapping their process groups instead of + // deriving ports from master_addr, so the bind targets sit on OS-chosen + // free ports rather than fixed offsets that may collide with random + // outgoing TCP source ports on a busy host. + repeated int32 group_ports = 2; } enum MasterStatus { diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 24bd4937a5..fa88bd9fdf 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -164,6 +164,7 @@ Options create_options(const std::string& instance_name, bool is_local) { static_cast(disagg_pd_config.transfer_listen_port())) .nnodes(distributed_config.nnodes()) .node_rank(distributed_config.node_rank()) + .enable_single_process(distributed_config.enable_single_process()) .dp_size(parallel_config.dp_size()) .cp_size(parallel_config.cp_size()) .ep_size(parallel_config.ep_size()) @@ -369,10 +370,12 @@ int run() { service_config.host(net::get_local_ip_addr()); } + const bool single_process_mode = distributed_config.enable_single_process(); const bool is_local = - !service_config.host().empty() && - net::extract_ip(distributed_config.master_node_addr()) == - service_config.host(); + single_process_mode || + (!service_config.host().empty() && + net::extract_ip(distributed_config.master_node_addr()) == + service_config.host()); LOG(INFO) << "set worker role to " << (is_local ? "local worker" : "remote worker"); @@ -451,6 +454,9 @@ int run() { std::unique_ptr master; // working node + // Assistant masters serve non-zero node ranks in a multi-process multi-node + // deployment. Single-process serving always runs at node_rank 0, so it takes + // the master path below. if (options.node_rank() != 0) { if (model_config.backend() == "dit") { master = std::make_unique(options);