feat: support joy-image-edit-plus on NPU device & refact sequence parallel Mixin.

xllm/core/common/global_flags.h

@@ -113,6 +113,8 @@ DECLARE_int64(cfg_size);
 
 DECLARE_int64(vae_size);
 
+DECLARE_int64(text_encoder_tp_size);
+
 DECLARE_bool(enable_mm_encoder_dp);
 
 DECLARE_bool(enable_multi_stream_parallel);

xllm/core/common/options.h

@@ -147,6 +147,8 @@ class Options {
 
   PROPERTY(int32_t, vae_size) = 1;
 
+  PROPERTY(int32_t, text_encoder_tp_size) = 1;
+
   PROPERTY(std::optional<std::string>, instance_name);
 
   PROPERTY(bool, enable_disagg_pd) = false;

xllm/core/distributed_runtime/dist_manager.cpp

@@ -189,12 +189,14 @@ void DistManager::setup_multi_node_workers(
     const int32_t sp_size = options.sp_size();
     const int32_t cfg_size = options.cfg_size();
     const int32_t vae_size = options.vae_size();
+    const int32_t text_encoder_tp_size = options.text_encoder_tp_size();
     LOG(INFO) << "Multi-node serving world_size = " << world_size
               << ", each_node_ranks = " << each_node_ranks
               << ", current node rank = " << options.node_rank()
               << ", nnodes = " << options.nnodes() << ", dp_size = " << dp_size
               << ", tp_size = " << tp_size << ", sp_size = " << sp_size
-              << ", cfg_size = " << cfg_size << ", vae_size = " << vae_size;
+              << ", cfg_size = " << cfg_size << ", vae_size = " << vae_size
+              << ", text_encoder_tp_size = " << text_encoder_tp_size;
   } else {
     LOG(INFO) << "Multi-node serving world_size = " << world_size
               << ", each_node_ranks = " << each_node_ranks

xllm/core/distributed_runtime/master.cpp

@@ -615,6 +615,7 @@ Master::Master(const Options& options, EngineType type)
         .model_id(options.model_id())
         .devices(devices)
         .backend(options.backend())
+        .npu_kernel_backend(options_.npu_kernel_backend())
         .enable_prefix_cache(options_.enable_prefix_cache())
         .enable_chunked_prefill(options_.enable_chunked_prefill())
         .enable_offline_inference(options_.enable_offline_inference())
@@ -634,7 +635,8 @@ Master::Master(const Options& options, EngineType type)
         .tp_size(options_.tp_size())
         .sp_size(options_.sp_size())
         .cfg_size(options_.cfg_size())
-        .vae_size(options_.vae_size());
+        .vae_size(options_.vae_size())
+        .text_encoder_tp_size(options_.text_encoder_tp_size());
 
     auto dit_engine = std::make_unique<DiTEngine>(eng_options);
     engine_ = std::move(dit_engine);

xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h

@@ -21,10 +21,11 @@ limitations under the License.
 
 namespace xllm::spawn_worker_protocol {
 
-inline constexpr int32_t kArgumentCount = 36;
+inline constexpr int32_t kArgumentCount = 37;
 inline constexpr int32_t kMinimumArgumentCount = 34;
 inline constexpr int32_t kIndexerCacheDtypeArgumentIndex = 34;
 inline constexpr int32_t kEnableMtpDraftBodyTp1ArgumentIndex = 35;
+inline constexpr int32_t kTextEncoderTpSizeArgumentIndex = 36;
 inline constexpr char kDefaultIndexerCacheDtype[] = "auto";
 
 inline std::optional<std::string> parse_indexer_cache_dtype(

xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp

@@ -88,6 +88,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr,
                                      int32_t tp_size,
                                      int32_t sp_size,
                                      int32_t cfg_size,
+                                     int32_t text_encoder_tp_size,
                                      int32_t cp_size,
                                      int32_t ep_size,
                                      const InstanceRole& instance_role,
@@ -100,6 +101,8 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr,
   const bool is_dit_backend = backend == "dit";
   const int32_t effective_sp_size = is_dit_backend ? sp_size : 1;
   const int32_t effective_cfg_size = is_dit_backend ? cfg_size : 1;
+  const int32_t effective_text_encoder_tp_size =
+      is_dit_backend ? text_encoder_tp_size : 1;
   runner_options.block_size(block_size)
       .backend(backend)
       .world_size(world_size)
@@ -119,6 +122,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr,
       .tp_size(tp_size)
       .sp_size(effective_sp_size)
       .cfg_size(effective_cfg_size)
+      .text_encoder_tp_size(effective_text_encoder_tp_size)
       .enable_shm(enable_shm)
       .input_shm_size(input_shm_size)
       .output_shm_size(output_shm_size)
@@ -142,6 +146,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr,
       .tp_size(tp_size)
       .sp_size(effective_sp_size)
       .cfg_size(effective_cfg_size)
+      .text_encoder_tp_size(effective_text_encoder_tp_size)
       .communication_backend(communication_backend);
   DistributedConfig::get_instance().master_node_addr(master_node_addr);
   KVCacheConfig::get_instance()

xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h

@@ -59,6 +59,7 @@ class SpawnWorkerServer final {
                              int32_t tp_size,
                              int32_t sp_size,
                              int32_t cfg_size,
+                             int32_t text_encoder_tp_size,
                              int32_t cp_size,
                              int32_t ep_size,
                              const InstanceRole& instance_role,

xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp

@@ -60,6 +60,7 @@ limitations under the License.
 // @instance_role
 // @indexer_cache_dtype
 // @enable_mtp_draft_body_tp1
+// @text_encoder_tp_size
 int main(int argc, char* argv[]) {
   const std::optional<std::string> parsed_indexer_cache_dtype =
       xllm::spawn_worker_protocol::parse_indexer_cache_dtype(argc, argv);
@@ -117,8 +118,14 @@ int main(int argc, char* argv[]) {
       static_cast<int32_t>(
           atoi(argv[xllm::spawn_worker_protocol::
                         kEnableMtpDraftBodyTp1ArgumentIndex])) > 0;
+  const int32_t text_encoder_tp_size =
+      argc > xllm::spawn_worker_protocol::kTextEncoderTpSizeArgumentIndex
+          ? static_cast<int32_t>(
+                atoi(argv[xllm::spawn_worker_protocol::
+                              kTextEncoderTpSizeArgumentIndex]))
+          : 1;
   if (world_size < 1 || global_rank < 0 || global_rank >= world_size ||
-      cp_size < 1 || ep_size < 1 ||
+      cp_size < 1 || ep_size < 1 || text_encoder_tp_size < 1 ||
       (instance_role_str != "DEFAULT" && instance_role_str != "PREFILL" &&
        instance_role_str != "DECODE")) {
     LOG(ERROR) << "Invalid spawn worker topology: global_rank=" << global_rank
@@ -159,6 +166,7 @@ int main(int argc, char* argv[]) {
       << ", max_encoder_cache_size = " << max_encoder_cache_size
       << ", dp_size = " << dp_size << ", tp_size = " << tp_size
       << ", sp_size = " << sp_size << ", cfg_size = " << cfg_size
+      << ", text_encoder_tp_size = " << text_encoder_tp_size
       << ", indexer_cache_dtype = " << indexer_cache_dtype
       << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 << "\n";
 
@@ -193,6 +201,7 @@ int main(int argc, char* argv[]) {
                                  tp_size,
                                  sp_size,
                                  cfg_size,
+                                 text_encoder_tp_size,
                                  cp_size,
                                  ep_size,
                                  instance_role,

xllm/core/distributed_runtime/worker_server.cpp

@@ -166,14 +166,15 @@ void WorkerServer::create_server(const runtime::Options& options,
   const ParallelArgs* parallel_args = nullptr;
   std::unique_ptr<CollectiveCommunicatorBase> comm;
   if (worker_type == WorkerType::DIT) {
-    auto dit_comm =
-        std::make_unique<DiTCollectiveCommunicator>(worker_global_rank,
-                                                    world_size,
-                                                    options.dp_size(),
-                                                    options.tp_size(),
-                                                    options.sp_size(),
-                                                    options.cfg_size(),
-                                                    options.vae_size());
+    auto dit_comm = std::make_unique<DiTCollectiveCommunicator>(
+        worker_global_rank,
+        world_size,
+        options.dp_size(),
+        options.tp_size(),
+        options.sp_size(),
+        options.cfg_size(),
+        options.vae_size(),
+        options.text_encoder_tp_size());
     comm = std::move(dit_comm);
   } else {
     auto common_comm = std::make_unique<CollectiveCommunicator>(
@@ -315,6 +316,9 @@ void WorkerServer::create_spawn_server(int32_t local_rank,
   const char* sp_size_ptr = sp_size_str.c_str();
   std::string cfg_size_str = std::to_string(options.cfg_size());
   const char* cfg_size_ptr = cfg_size_str.c_str();
+  std::string text_encoder_tp_size_str =
+      std::to_string(options.text_encoder_tp_size());
+  const char* text_encoder_tp_size_ptr = text_encoder_tp_size_str.c_str();
   const std::string& indexer_cache_dtype =
       KVCacheConfig::get_instance().indexer_cache_dtype();
   const char* indexer_cache_dtype_ptr = indexer_cache_dtype.c_str();
@@ -372,6 +376,7 @@ void WorkerServer::create_spawn_server(int32_t local_rank,
                         instance_role_ptr,
                         indexer_cache_dtype_ptr,
                         enable_mtp_draft_body_tp1_ptr,
+                        text_encoder_tp_size_ptr,
                         nullptr};
   static_assert(std::size(argv) == spawn_worker_protocol::kArgumentCount + 1);
   pid_t pid;

xllm/core/framework/config/parallel_config.cpp

@@ -44,6 +44,10 @@ DEFINE_int64(cfg_size,
 
 DEFINE_int64(vae_size, 1, "Vae patch parallelism size");
 
+DEFINE_int64(text_encoder_tp_size,
+             1,
+             "Text encoder tensor parallelism size for DiT model.");
+
 DEFINE_string(
     communication_backend,
     "hccl",
@@ -80,6 +84,7 @@ void ParallelConfig::from_flags() {
   XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(cfg_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(vae_size);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(text_encoder_tp_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(communication_backend);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_mm_encoder_dp);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_multi_stream_parallel);
@@ -95,6 +100,7 @@ void ParallelConfig::from_json(const JsonReader& json) {
   XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(vae_size);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(text_encoder_tp_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(communication_backend);
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_mm_encoder_dp);
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_multi_stream_parallel);
@@ -114,6 +120,8 @@ void ParallelConfig::append_config_json(
       config_json, default_config, cfg_size);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, vae_size);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, text_encoder_tp_size);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, communication_backend);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(

xllm/core/framework/config/parallel_config.h

@@ -48,6 +48,7 @@ class ParallelConfig final {
          "sp_size",
          "cfg_size",
          "vae_size",
+         "text_encoder_tp_size",
          "communication_backend",
          "enable_mm_encoder_dp",
          "enable_multi_stream_parallel",
@@ -73,6 +74,8 @@ class ParallelConfig final {
 
   PROPERTY(int64_t, vae_size) = 1;
 
+  PROPERTY(int64_t, text_encoder_tp_size) = 1;
+
   PROPERTY(std::string, communication_backend) = "hccl";
 
   PROPERTY(bool, enable_mm_encoder_dp) = false;

xllm/core/framework/model/model_args.h

@@ -573,6 +573,12 @@ struct ModelArgs {
   PROPERTY(bool, zero_cond_t) = false;
   PROPERTY(bool, use_additional_t_cond) = false;
   PROPERTY(bool, use_layer3d_rope) = false;
+
+  // JoyImage-Edit-Plus dit related args
+  PROPERTY(double, mlp_width_ratio) = 4.0;
+  PROPERTY(int64_t, text_dim) = 4096;
+  PROPERTY(std::vector<int64_t>, rope_dim_list) = { 16, 56, 56 };
+  PROPERTY(int64_t, rope_theta_dit) = 10000;
 };
 
 // Qwen hybrid models may describe full-attention layers explicitly via

xllm/core/framework/parallel_state/dit_collective_communicator.cpp

@@ -35,13 +35,15 @@ limitations under the License.
 #include "util/net.h"
 namespace xllm {
 
-DiTCollectiveCommunicator::DiTCollectiveCommunicator(int32_t global_rank,
-                                                     int32_t world_size,
-                                                     int32_t dit_dp_size,
-                                                     int32_t dit_tp_size,
-                                                     int32_t dit_sp_size,
-                                                     int32_t dit_cfg_size,
-                                                     int32_t dit_vae_size)
+DiTCollectiveCommunicator::DiTCollectiveCommunicator(
+    int32_t global_rank,
+    int32_t world_size,
+    int32_t dit_dp_size,
+    int32_t dit_tp_size,
+    int32_t dit_sp_size,
+    int32_t dit_cfg_size,
+    int32_t dit_vae_size,
+    int32_t dit_text_encoder_tp_size)
     : CollectiveCommunicatorBase(global_rank, world_size) {
   parallel_args_ = std::make_unique<ParallelArgs>(global_rank,
                                                   world_size,
@@ -50,13 +52,15 @@ DiTCollectiveCommunicator::DiTCollectiveCommunicator(int32_t global_rank,
                                                   dit_sp_size,
                                                   dit_cfg_size,
                                                   dit_vae_size,
+                                                  dit_text_encoder_tp_size,
                                                   /*process_group=*/nullptr);
   DiTMapping::Options dit_mapping_options;
   dit_mapping_options.dit_tp_size(dit_tp_size)
       .dit_sp_size(dit_sp_size)
       .dit_cfg_size(dit_cfg_size)
       .dit_dp_size(dit_dp_size)
-      .dit_vae_size(dit_vae_size);
+      .dit_vae_size(dit_vae_size)
+      .dit_text_encoder_tp_size(dit_text_encoder_tp_size);
   dit_mapping_ = std::make_unique<DiTMapping>(
       world_size, global_rank, dit_mapping_options);
 }
@@ -88,14 +92,16 @@ void DiTCollectiveCommunicator::create_process_groups(
       create_process_group_by_type("dp", dit_dp_group_, device);
   parallel_args_->dit_vae_group_ =
       create_process_group_by_type("vae", dit_vae_group_, device);
+  parallel_args_->dit_text_encoder_tp_group_ = create_process_group_by_type(
+      "text_encoder_tp", dit_text_encoder_tp_group_, device);
 }
 
 ProcessGroup* DiTCollectiveCommunicator::create_process_group_by_type(
     const std::string& group_type,
     std::unique_ptr<ProcessGroup>& member_group,
     const torch::Device& device) {
   int32_t group_size = parallel_args_->get_group_size_by_type(group_type);
-  if (group_size > 1 && dit_mapping_) {
+  if (group_size >= 1 && dit_mapping_) {
     auto parallel_info = dit_mapping_->get_parallel_info(group_type);
     auto group_id = parallel_info.current_group_id();
     auto num_group = parallel_info.num_group();

xllm/core/framework/parallel_state/dit_collective_communicator.h

@@ -28,7 +28,8 @@ class DiTCollectiveCommunicator : public CollectiveCommunicatorBase {
                             int32_t dit_tp_size,
                             int32_t dit_sp_size,
                             int32_t dit_cfg_size,
-                            int32_t dit_vae_size);
+                            int32_t dit_vae_size,
+                            int32_t dit_text_encoder_tp_size);
 
   ~DiTCollectiveCommunicator() = default;
 
@@ -54,6 +55,7 @@ class DiTCollectiveCommunicator : public CollectiveCommunicatorBase {
   std::unique_ptr<ProcessGroup> dit_dp_group_;
   std::unique_ptr<ProcessGroup> dit_cfg_group_;
   std::unique_ptr<ProcessGroup> dit_vae_group_;
+  std::unique_ptr<ProcessGroup> dit_text_encoder_tp_group_;
 };
 
 }  // namespace xllm