diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h
index 50cf1c6fc9..df93aa5b2e 100755
--- a/xllm/core/common/global_flags.h
+++ b/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);
diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h
index c9b3ac67e3..e0839c9863 100644
--- a/xllm/core/common/options.h
+++ b/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;
diff --git a/xllm/core/distributed_runtime/dist_manager.cpp b/xllm/core/distributed_runtime/dist_manager.cpp
index eb66ba45d5..4ca3878761 100644
--- a/xllm/core/distributed_runtime/dist_manager.cpp
+++ b/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
diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp
index dd48dc3894..73603ba87d 100644
--- a/xllm/core/distributed_runtime/master.cpp
+++ b/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);
diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h
index b51c5d84b4..f8aafd95e7 100644
--- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h
+++ b/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(
diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp
index 5d621351fb..b875fc0be2 100644
--- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp
+++ b/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()
diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h
index 3b4935f2a2..3ce9240d55 100644
--- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h
+++ b/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,
diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp
index a2cdbc2c18..e676810f66 100644
--- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp
+++ b/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,
diff --git a/xllm/core/distributed_runtime/worker_server.cpp b/xllm/core/distributed_runtime/worker_server.cpp
index 42598464c5..9b25a3851c 100644
--- a/xllm/core/distributed_runtime/worker_server.cpp
+++ b/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;
diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp
index d29e1e2786..f333d53993 100644
--- a/xllm/core/framework/config/parallel_config.cpp
+++ b/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(
diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h
index 22159a1c2a..bae219a705 100644
--- a/xllm/core/framework/config/parallel_config.h
+++ b/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;
diff --git a/xllm/core/framework/model/model_args.h b/xllm/core/framework/model/model_args.h
index 4f6b7cc570..4a0a371473 100644
--- a/xllm/core/framework/model/model_args.h
+++ b/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
diff --git a/xllm/core/framework/parallel_state/dit_collective_communicator.cpp b/xllm/core/framework/parallel_state/dit_collective_communicator.cpp
index cfbaa13248..428d8efb0f 100644
--- a/xllm/core/framework/parallel_state/dit_collective_communicator.cpp
+++ b/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,6 +92,8 @@ 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(
@@ -95,7 +101,7 @@ ProcessGroup* DiTCollectiveCommunicator::create_process_group_by_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();
diff --git a/xllm/core/framework/parallel_state/dit_collective_communicator.h b/xllm/core/framework/parallel_state/dit_collective_communicator.h
index 83672c7b2d..2a84721f2a 100644
--- a/xllm/core/framework/parallel_state/dit_collective_communicator.h
+++ b/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
diff --git a/xllm/core/framework/parallel_state/dit_mapping.cpp b/xllm/core/framework/parallel_state/dit_mapping.cpp
index eec1b0b97a..f81571a879 100644
--- a/xllm/core/framework/parallel_state/dit_mapping.cpp
+++ b/xllm/core/framework/parallel_state/dit_mapping.cpp
@@ -28,6 +28,7 @@ DiTMapping::DiTMapping(const int32_t world_size,
   cfg_.backend("hccl");
   dp_.backend("hccl");
   vae_.backend("hccl");
+  text_encoder_tp_.backend("hccl");
   parse_parallel_info();
   validate();
   RankGenerator rank_generator(world_size);
@@ -47,6 +48,15 @@ DiTMapping::DiTMapping(const int32_t world_size,
   auto ranks_mapping_vae =
       rank_generator.get_ranks_mapping(vae_group_ranks, vae_group_order);
   set_group_by_type(vae_, "vae", ranks_mapping_vae.at("vae"));
+
+  std::vector<int32_t> text_encoder_tp_group_ranks = {
+      text_encoder_tp_.group_size()};
+  std::vector<std::string> text_encoder_tp_group_order = {"text_encoder_tp"};
+  auto ranks_mapping_text_encoder_tp = rank_generator.get_ranks_mapping(
+      text_encoder_tp_group_ranks, text_encoder_tp_group_order);
+  set_group_by_type(text_encoder_tp_,
+                    "text_encoder_tp",
+                    ranks_mapping_text_encoder_tp.at("text_encoder_tp"));
 }
 
 void DiTMapping::parse_parallel_info() {
@@ -65,6 +75,9 @@ void DiTMapping::parse_parallel_info() {
   if (options_.dit_vae_size() != -1) {
     vae_.group_size(options_.dit_vae_size());
   }
+  if (options_.dit_text_encoder_tp_size() != -1) {
+    text_encoder_tp_.group_size(options_.dit_text_encoder_tp_size());
+  }
 }
 
 void DiTMapping::validate() {
@@ -107,6 +120,21 @@ void DiTMapping::validate() {
              std::to_string(vae_.group_size()) + ", world_size is " +
              std::to_string(world_size_) +
              ". Please check `vae` and 'world_size'.";
+
+  CHECK(text_encoder_tp_.group_size() >= 1 &&
+        text_encoder_tp_.group_size() <= world_size_)
+      << "text_encoder_tp_size must be between 1 and world_size. "
+         "text_encoder_tp_size is " +
+             std::to_string(text_encoder_tp_.group_size()) +
+             ", world_size is " + std::to_string(world_size_) +
+             ". Please check `text_encoder_tp_size` and 'world_size'.";
+
+  CHECK(world_size_ % text_encoder_tp_.group_size() == 0)
+      << "world_size could not be divided by text_encoder_tp_size. "
+         "text_encoder_tp_size is " +
+             std::to_string(text_encoder_tp_.group_size()) +
+             ", world_size is " + std::to_string(world_size_) +
+             ". Please check `text_encoder_tp_size` and 'world_size'.";
 }
 
 void DiTMapping::set_group_by_type(
@@ -150,6 +178,8 @@ const ParallelInfo& DiTMapping::get_parallel_info(
     return dp_;
   } else if (group_type == "vae") {
     return vae_;
+  } else if (group_type == "text_encoder_tp") {
+    return text_encoder_tp_;
   } else {
     LOG(FATAL) << "get unexpected group_type: " << group_type;
   }
@@ -161,6 +191,7 @@ nlohmann::json DiTMapping::to_json() {
   data["SpSize"] = options_.dit_sp_size();
   data["TpSize"] = options_.dit_tp_size();
   data["CfgSize"] = options_.dit_cfg_size();
+  data["TextEncoderTpSize"] = options_.dit_text_encoder_tp_size();
   data["worldSize"] = world_size_;
   data["rank"] = rank_;
   data["sp"] = sp_.to_json();
@@ -168,6 +199,7 @@ nlohmann::json DiTMapping::to_json() {
   data["cfg"] = cfg_.to_json();
   data["dp"] = dp_.to_json();
   data["vae"] = vae_.to_json();
+  data["text_encoder_tp"] = text_encoder_tp_.to_json();
   return data;
 }
 
diff --git a/xllm/core/framework/parallel_state/dit_mapping.h b/xllm/core/framework/parallel_state/dit_mapping.h
index 536e6cd469..b227e14574 100644
--- a/xllm/core/framework/parallel_state/dit_mapping.h
+++ b/xllm/core/framework/parallel_state/dit_mapping.h
@@ -36,6 +36,8 @@ class DiTMapping final {
     PROPERTY(int32_t, dit_dp_size) = -1;
     // vae size
     PROPERTY(int32_t, dit_vae_size) = -1;
+    // text encoder tensor parallel size
+    PROPERTY(int32_t, dit_text_encoder_tp_size) = -1;
   };
 
   DiTMapping(const int32_t world_size,
@@ -72,5 +74,6 @@ class DiTMapping final {
   ParallelInfo cfg_ = ParallelInfo();
   ParallelInfo dp_ = ParallelInfo();
   ParallelInfo vae_ = ParallelInfo();
+  ParallelInfo text_encoder_tp_ = ParallelInfo();
 };
 }  // namespace xllm
diff --git a/xllm/core/framework/parallel_state/parallel_args.h b/xllm/core/framework/parallel_state/parallel_args.h
index e2e3b0edc4..845929065d 100644
--- a/xllm/core/framework/parallel_state/parallel_args.h
+++ b/xllm/core/framework/parallel_state/parallel_args.h
@@ -101,6 +101,7 @@ struct ParallelArgs {
                int32_t sp_size,
                int32_t cfg_size,
                int32_t vae_size,
+               int32_t text_encoder_tp_size,
                ProcessGroup* process_group)
       : rank_(rank),
         world_size_(world_size),
@@ -109,6 +110,7 @@ struct ParallelArgs {
         sp_size_(sp_size),
         cfg_size_(cfg_size),
         vae_size_(vae_size),
+        text_encoder_tp_size_(text_encoder_tp_size),
         process_group_(process_group) {}
 
   int32_t get_group_size_by_type(const std::string& group_type) const {
@@ -124,6 +126,8 @@ struct ParallelArgs {
       return ep_size();
     } else if (group_type == "vae") {
       return vae_size();
+    } else if (group_type == "text_encoder_tp") {
+      return text_encoder_tp_size();
     } else if (group_type == "cp") {
       return cp_size();
     } else {
@@ -185,6 +189,9 @@ struct ParallelArgs {
   // cfg size
   PROPERTY(int32_t, vae_size) = 1;
 
+  // text encoder tensor parallel size
+  PROPERTY(int32_t, text_encoder_tp_size) = 1;
+
   // atb hccl mapping json data
   PROPERTY(nlohmann::json, mapping_data);
 
@@ -225,6 +232,7 @@ struct ParallelArgs {
   ProcessGroup* dit_cfg_group_ = nullptr;
   ProcessGroup* dit_dp_group_ = nullptr;
   ProcessGroup* dit_vae_group_ = nullptr;
+  ProcessGroup* dit_text_encoder_tp_group_ = nullptr;
 };
 
 }  // namespace xllm
diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp
index 313ff262b6..91c20b4c86 100644
--- a/xllm/core/framework/parallel_state/parallel_state.cpp
+++ b/xllm/core/framework/parallel_state/parallel_state.cpp
@@ -292,7 +292,6 @@ torch::Tensor scatter(torch::Tensor input,
       << "dim_size " << dim_size << " cannot be divided by world_size "
       << world_size;
 
-  // torch::split does not create contiguous tensors by default.
   const auto tensor_list = input.split(dim_size / world_size, dim);
   const int32_t rank = process_group->rank();
   return tensor_list[rank];
@@ -312,12 +311,10 @@ std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input,
     return [input]() { return input; };
   }
 
-  auto rank = process_group->rank();
-
   TORCH_CHECK(input.dim() == 4,
               "all_to_all_4D: input must be 4D, got dim=",
               input.dim());
-  auto send_input = input;
+  torch::Tensor send_input = input;
 
   if (scatter_idx == 2 && gather_idx == 1) {
     // branch A : from "sequence shard" -> "head shard"
@@ -369,12 +366,10 @@ std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input,
               shard_head_num,
               head_size]() mutable -> torch::Tensor {
         all2all_work->wait();
-        auto comm_output =
-            output.reshape({seqlen, bs, shard_head_num, head_size})
-                .transpose(0, 1)
-                .contiguous()
-                .reshape({bs, seqlen, shard_head_num, head_size});
-        return comm_output;
+        return output.reshape({seqlen, bs, shard_head_num, head_size})
+            .transpose(0, 1)
+            .contiguous()
+            .reshape({bs, seqlen, shard_head_num, head_size});
       };
     }
   } else if (scatter_idx == 1 && gather_idx == 2) {
diff --git a/xllm/core/framework/parallel_state/parallel_state.h b/xllm/core/framework/parallel_state/parallel_state.h
index b40ddfa34d..38e8b3aed2 100644
--- a/xllm/core/framework/parallel_state/parallel_state.h
+++ b/xllm/core/framework/parallel_state/parallel_state.h
@@ -15,6 +15,8 @@ limitations under the License.
 
 #pragma once
 
+#include <functional>
+
 #include "parallel_args.h"
 #include "process_group.h"
 
@@ -78,11 +80,11 @@ torch::Tensor scatter(torch::Tensor input,
                       ProcessGroup* process_group,
                       int dim = -1);
 
-std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input_,
+std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input,
                                              int32_t scatter_idx,
                                              int32_t gather_idx,
-                                             bool is_sync,
-                                             ProcessGroup* pg);
+                                             bool async_ops,
+                                             ProcessGroup* process_group);
 
 // Create a process group where each process has a single device
 // devices: list of devices to create process groups on.
diff --git a/xllm/core/kernels/npu/active.cpp b/xllm/core/kernels/npu/active.cpp
index 01bff64f07..17107a3d31 100644
--- a/xllm/core/kernels/npu/active.cpp
+++ b/xllm/core/kernels/npu/active.cpp
@@ -22,9 +22,13 @@ limitations under the License.
 namespace xllm::kernel::npu {
 
 torch::Tensor active(const torch::Tensor& input, const std::string& act_mode) {
-  if (act_mode != "silu" && act_mode != "swiglu") {
-    LOG(FATAL) << "Only swiglu activation is supported in NPU active";
+  if (act_mode == "gelu" || act_mode == "gelu_pytorch_tanh") {
+    const auto approximate = act_mode == "gelu_pytorch_tanh" ? "tanh" : "none";
+    return at_npu::native::custom_ops::npu_gelu(input, approximate);
   }
-  return at_npu::native::custom_ops::npu_swiglu(input);
+  if (act_mode == "silu" || act_mode == "swiglu") {
+    return at_npu::native::custom_ops::npu_swiglu(input);
+  }
+  LOG(FATAL) << "Unsupported NPU activation: " << act_mode;
 }
-}  // namespace xllm::kernel::npu
\ No newline at end of file
+}  // namespace xllm::kernel::npu
diff --git a/xllm/core/kernels/npu/attention.cpp b/xllm/core/kernels/npu/attention.cpp
index cd6a0e9f13..663cf44dcd 100644
--- a/xllm/core/kernels/npu/attention.cpp
+++ b/xllm/core/kernels/npu/attention.cpp
@@ -23,6 +23,17 @@ void reshape_paged_cache(torch::Tensor& key,
                          torch::Tensor& k_cache,
                          std::optional<torch::Tensor>& v_cache,
                          const torch::Tensor& slot_mapping) {
+  // KV-cache-free full prefill does not need to populate paged caches.
+  if (!k_cache.defined()) {
+    CHECK(!v_cache.has_value() || !v_cache->defined())
+        << "NPU key and value caches must either both be defined or both be "
+           "empty.";
+    return;
+  }
+  CHECK(value.has_value() && value->defined())
+      << "NPU value tensor must be defined when key cache is defined.";
+  CHECK(v_cache.has_value() && v_cache->defined())
+      << "NPU value cache must be defined when key cache is defined.";
   atb::npu_reshape_and_cache(
       key, value.value(), k_cache, v_cache.value(), slot_mapping);
 }
diff --git a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp
index 3c85b1eacf..8d1ba3bcc2 100644
--- a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp
+++ b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp
@@ -107,7 +107,8 @@ std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
     int64_t block_size,
     int64_t sparse_mode,
     const std::string& input_layout,
-    bool softmax_lse_flag) {
+    bool softmax_lse_flag,
+    bool is_causal) {
   check_tensor(query, "query", "npu_fused_infer_attention");
   check_tensor(key, "key", "npu_fused_infer_attention");
   check_tensor(value, "value", "npu_fused_infer_attention");
@@ -143,7 +144,7 @@ std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
   std::string layout = input_layout;
   char* input_layout_ptr = const_cast<char*>(layout.c_str());
   int64_t pre_tokens = kSwaIntMax;
-  int64_t next_tokens = 0;
+  int64_t next_tokens = is_causal ? 0 : kSwaIntMax;
   int64_t inner_precise = 0;
   int64_t antiquant_mode = 0;
   int64_t key_antiquant_mode = 0;
diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h
index b9a40f2c4f..fac965877c 100644
--- a/xllm/core/kernels/npu/npu_ops_api.h
+++ b/xllm/core/kernels/npu/npu_ops_api.h
@@ -51,6 +51,12 @@ void batch_decode(const torch::Tensor& query,
                   const torch::Tensor& seq_lens,
                   torch::Tensor& output);
 
+void apply_rotary(torch::Tensor& query,
+                  torch::Tensor& key,
+                  const torch::Tensor& cos,
+                  const torch::Tensor& sin,
+                  const std::string& input_layout);
+
 std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
     const torch::Tensor& query,
     const torch::Tensor& key,
@@ -65,7 +71,8 @@ std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
     int64_t block_size,
     int64_t sparse_mode,
     const std::string& input_layout,
-    bool softmax_lse_flag = false);
+    bool softmax_lse_flag = false,
+    bool is_causal = true);
 
 void batch_chunked_paged_prefill(const torch::Tensor& query,
                                  const torch::Tensor& k_cache,
diff --git a/xllm/core/kernels/npu/rope.cpp b/xllm/core/kernels/npu/rope.cpp
index ac9b90f582..06ec8aabf6 100644
--- a/xllm/core/kernels/npu/rope.cpp
+++ b/xllm/core/kernels/npu/rope.cpp
@@ -20,6 +20,41 @@ limitations under the License.
 
 namespace xllm::kernel::npu {
 
+void apply_rotary(torch::Tensor& query,
+                  torch::Tensor& key,
+                  const torch::Tensor& cos,
+                  const torch::Tensor& sin,
+                  const std::string& input_layout) {
+  CHECK(cos.defined() && sin.defined());
+  CHECK(cos.sizes() == sin.sizes());
+  CHECK_GT(cos.dim(), 0);
+  const int64_t rotary_dim = cos.size(-1);
+  CHECK_GT(rotary_dim, 0);
+  const int64_t num_tokens = cos.numel() / rotary_dim;
+  CHECK_GT(num_tokens, 0);
+  CHECK_EQ(query.numel() % (num_tokens * rotary_dim), 0);
+  CHECK_EQ(key.numel() % (num_tokens * rotary_dim), 0);
+
+  const std::vector<int64_t> query_shape = query.sizes().vec();
+  const std::vector<int64_t> key_shape = key.sizes().vec();
+  const int64_t num_query_heads = query.numel() / (num_tokens * rotary_dim);
+  const int64_t num_key_heads = key.numel() / (num_tokens * rotary_dim);
+
+  torch::Tensor query_view =
+      query.contiguous().view({1, num_tokens, num_query_heads, rotary_dim});
+  torch::Tensor key_view =
+      key.contiguous().view({1, num_tokens, num_key_heads, rotary_dim});
+  torch::Tensor cos_view =
+      cos.contiguous().view({1, num_tokens, 1, rotary_dim});
+  torch::Tensor sin_view =
+      sin.contiguous().view({1, num_tokens, 1, rotary_dim});
+
+  auto rotary_result = at_npu::native::custom_ops::npu_apply_rotary_pos_emb(
+      query_view, key_view, cos_view, sin_view, input_layout);
+  query = std::get<0>(rotary_result).view(query_shape);
+  key = std::get<1>(rotary_result).view(key_shape);
+}
+
 void apply_rotary(torch::Tensor& q,
                   torch::Tensor& k,
                   const torch::Tensor& cos_sin_cache,
diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp
index d383aa4788..56497a39b5 100644
--- a/xllm/core/kernels/ops_api.cpp
+++ b/xllm/core/kernels/ops_api.cpp
@@ -115,8 +115,17 @@ void apply_rotary(RotaryParams& params) {
                     params.dynamic_ntk,
                     params.max_query_len);
 #elif defined(USE_NPU)
-  npu::apply_rotary(
-      params.q, params.k, params.cos_sin, params.position_ids.value());
+  if (!params.position_ids.has_value() && params.cos.defined() &&
+      params.sin.defined()) {
+    npu::apply_rotary(params.q, params.k, params.cos, params.sin, "BSND");
+  } else {
+    CHECK(params.position_ids.has_value())
+        << "NPU rotary embedding requires position_ids when precomputed "
+           "cos/sin "
+           "are unavailable";
+    npu::apply_rotary(
+        params.q, params.k, params.cos_sin, params.position_ids.value());
+  }
 #elif defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_DCU)
   bool is_neox = !params.interleaved;
   torch::Tensor pos_ids;
diff --git a/xllm/core/layers/common/CMakeLists.txt b/xllm/core/layers/common/CMakeLists.txt
index 5314192fb4..920f414ad8 100755
--- a/xllm/core/layers/common/CMakeLists.txt
+++ b/xllm/core/layers/common/CMakeLists.txt
@@ -7,6 +7,7 @@ cc_library(
     oxygen_vision_attention.h
     qwen2_attention.h
     qwen2_vision_attention.h
+    ../npu/npu_rope_layer_impl.h
     qwen3_next_rms_norm.h
     deepseek_v4_rotary_embedding.h
     partial_rotary_embedding.h
@@ -33,6 +34,7 @@ cc_library(
     oxygen_vision_attention.cpp
     qwen2_attention.cpp
     qwen2_vision_attention.cpp
+    ../npu/npu_rope_layer_impl.cpp
     qwen3_next_rms_norm.cpp
     deepseek_v4_rotary_embedding.cpp
     partial_rotary_embedding.cpp
diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp
index 59ddaa0c37..2c394a65c1 100644
--- a/xllm/core/layers/common/dense_mlp.cpp
+++ b/xllm/core/layers/common/dense_mlp.cpp
@@ -113,6 +113,7 @@ torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
       return down_proj_->forward(gate_up,
                                  row_parallel_reduce_mode_for_fc1(*fc1_ctx));
     }
+
     return down_proj_->forward(gate_up);
   }
 
diff --git a/xllm/core/layers/common/qwen2_vision_attention.cpp b/xllm/core/layers/common/qwen2_vision_attention.cpp
index d7f3d89cf1..cb885d0bdf 100644
--- a/xllm/core/layers/common/qwen2_vision_attention.cpp
+++ b/xllm/core/layers/common/qwen2_vision_attention.cpp
@@ -18,6 +18,9 @@ limitations under the License.
 #include <algorithm>
 
 #include "framework/parallel_state/parallel_state.h"
+#if defined(USE_NPU)
+#include "kernels/npu/npu_ops_api.h"
+#endif
 #if defined(USE_MLU)
 #include "kernels/mlu/mlu_ops_api.h"
 #endif
@@ -63,6 +66,9 @@ Qwen2VisionAttentionImpl::Qwen2VisionAttentionImpl(const ModelContext& context,
                                             quant_args,
                                             parallel_args.tp_group_,
                                             options));
+#if defined(USE_NPU)
+  rope_layer_ = register_module("rope", NpuRopeLayer(context));
+#endif
 }
 
 std::vector<torch::Tensor> Qwen2VisionAttentionImpl::split_qkv(
@@ -157,6 +163,44 @@ void compute_qwen2_vision_attention_torch(
 }
 #endif  // defined(USE_CUDA) || defined(USE_DCU)
 
+#if defined(USE_NPU)
+void compute_qwen_vision_attention_fused(
+    const torch::Tensor& query,
+    const torch::Tensor& key,
+    const torch::Tensor& value,
+    torch::Tensor& output,
+    const std::vector<int32_t>& cu_seq_len_vec,
+    float scale) {
+  CHECK_GE(cu_seq_len_vec.size(), static_cast<size_t>(2));
+  CHECK_EQ(cu_seq_len_vec.front(), 0);
+
+  std::vector<int64_t> actual_seq_lengths;
+  actual_seq_lengths.reserve(cu_seq_len_vec.size() - 1);
+  for (size_t index = 1; index < cu_seq_len_vec.size(); ++index) {
+    CHECK_GE(cu_seq_len_vec[index], cu_seq_len_vec[index - 1]);
+    actual_seq_lengths.emplace_back(cu_seq_len_vec[index]);
+  }
+
+  auto attention_result =
+      xllm::kernel::npu::npu_fused_infer_attention(query,
+                                                   key,
+                                                   value,
+                                                   /*atten_mask=*/std::nullopt,
+                                                   /*block_table=*/std::nullopt,
+                                                   actual_seq_lengths,
+                                                   actual_seq_lengths,
+                                                   query.size(1),
+                                                   key.size(1),
+                                                   scale,
+                                                   /*block_size=*/0,
+                                                   /*sparse_mode=*/0,
+                                                   "TND",
+                                                   /*softmax_lse_flag=*/false,
+                                                   /*is_causal=*/false);
+  output.copy_(std::get<0>(attention_result).view_as(output));
+}
+#endif  // defined(USE_NPU)
+
 }  // namespace
 
 torch::Tensor Qwen2VisionAttentionImpl::forward(
@@ -190,6 +234,16 @@ torch::Tensor Qwen2VisionAttentionImpl::forward(
   k = k.reshape({B * S, num_attention_heads_per_partition_, head_dim});
 
   // Apply rotary position embedding to both q and k in a single call.
+#if defined(USE_NPU)
+  auto q_atb = q.reshape({B * S, -1});
+  auto k_atb = k.reshape({B * S, -1});
+  auto rotary_result = rope_layer_->forward(
+      q_atb, k_atb, m_cos_pos, m_sin_pos, cu_seq_len, cu_seq_len_vec);
+  q = std::get<0>(rotary_result)
+          .reshape({B * S, num_attention_heads_per_partition_, head_dim});
+  k = std::get<1>(rotary_result)
+          .reshape({B * S, num_attention_heads_per_partition_, head_dim});
+#else
   // NOTE: Do NOT call apply_rotary twice; the first call already handles both
   // q and k. A second call would incorrectly apply RoPE to k a second time.
   xllm::kernel::RotaryParams rotary_params;
@@ -215,6 +269,7 @@ torch::Tensor Qwen2VisionAttentionImpl::forward(
       {B * S, num_attention_heads_per_partition_, head_dim});
   k = rotary_params.k.reshape(
       {B * S, num_attention_heads_per_partition_, head_dim});
+#endif
 
   // q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
   // q and k are already [B*S, H, D] after the reshape above; just
@@ -254,6 +309,8 @@ torch::Tensor Qwen2VisionAttentionImpl::forward(
   // AOT kernels in this project do not support head_dim=80, so we intentionally
   // do not call FlashInfer here and run attention entirely in PyTorch instead.
   compute_qwen2_vision_attention_torch(q, k, v, output, cu_seq_len_vec, scale_);
+#elif defined(USE_NPU)
+  compute_qwen_vision_attention_fused(q, k, v, output, cu_seq_len_vec, scale_);
 #endif
 
   // context_layer = rearrange(output, "(b s) h d -> s b (h d)", b=batch_size)
diff --git a/xllm/core/layers/common/qwen2_vision_attention.h b/xllm/core/layers/common/qwen2_vision_attention.h
index 3ed8ae4b54..d7f458f782 100644
--- a/xllm/core/layers/common/qwen2_vision_attention.h
+++ b/xllm/core/layers/common/qwen2_vision_attention.h
@@ -25,6 +25,9 @@ limitations under the License.
 #include "framework/quant_args.h"
 #include "framework/state_dict/state_dict.h"
 #include "linear.h"
+#if defined(USE_NPU)
+#include "layers/npu/npu_rope_layer_impl.h"
+#endif
 
 namespace xllm {
 namespace layer {
@@ -55,6 +58,9 @@ class Qwen2VisionAttentionImpl : public torch::nn::Module {
 
   QKVParallelLinear qkv_proj_{nullptr};
   RowParallelLinear proj_{nullptr};
+#if defined(USE_NPU)
+  NpuRopeLayer rope_layer_{nullptr};
+#endif
 };
 TORCH_MODULE(Qwen2VisionAttention);
 
diff --git a/xllm/core/layers/common/rms_norm.cpp b/xllm/core/layers/common/rms_norm.cpp
index b98b29e249..978b955d4f 100644
--- a/xllm/core/layers/common/rms_norm.cpp
+++ b/xllm/core/layers/common/rms_norm.cpp
@@ -46,6 +46,21 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>> RMSNormImpl::forward(
     std::optional<torch::Tensor> residual,
     std::optional<torch::Tensor> inplace_output) {
   auto org_shape = input.sizes().vec();
+
+  // Qwen3-VL vision blocks use LayerNorm, while the NPU fused normalization
+  // kernel only supports RMSNorm. Use Torch LayerNorm for the NPU TORCH path.
+  if (Platform::is_npu() && mode_ == kLayerNormMode) {
+    torch::Tensor norm_input = input;
+    std::optional<torch::Tensor> residual_out;
+    if (residual.has_value()) {
+      norm_input = input + residual.value();
+      residual_out = norm_input;
+    }
+    torch::Tensor output =
+        torch::layer_norm(norm_input, {norm_dim_}, weight_, bias_, eps_);
+    return std::make_tuple(output, residual_out);
+  }
+
   input = input.reshape({-1, norm_dim_});
 
   torch::Tensor output;
@@ -148,6 +163,13 @@ void RMSNormImpl::load_state_dict(const StateDict& state_dict) {
 #endif
 }
 
+void RMSNormImpl::verify_loaded_weights(const std::string& prefix) const {
+  CHECK(weight_is_loaded_) << "weight is not loaded for " << prefix + "weight";
+  if (bias_.defined()) {
+    CHECK(bias_is_loaded_) << "bias is not loaded for " << prefix + "bias";
+  }
+}
+
 void RMSNormImpl::set_layernorm_mode() {
   mode_ = kLayerNormMode;
   bias_ = register_parameter(
diff --git a/xllm/core/layers/common/rms_norm.h b/xllm/core/layers/common/rms_norm.h
index 32ee5fb3fd..0209267aa8 100644
--- a/xllm/core/layers/common/rms_norm.h
+++ b/xllm/core/layers/common/rms_norm.h
@@ -47,6 +47,8 @@ class RMSNormImpl : public torch::nn::Module {
 
   void load_state_dict(const StateDict& state_dict);
 
+  void verify_loaded_weights(const std::string& prefix) const;
+
   torch::Tensor weight() const { return weight_; }
   torch::Tensor bias() const { return bias_; }
   double eps() const { return eps_; }
diff --git a/xllm/core/layers/npu/npu_rope_layer_impl.cpp b/xllm/core/layers/npu/npu_rope_layer_impl.cpp
new file mode 100644
index 0000000000..b5a000232a
--- /dev/null
+++ b/xllm/core/layers/npu/npu_rope_layer_impl.cpp
@@ -0,0 +1,169 @@
+/* Copyright 2026 The xLLM Authors.
+
+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 "npu_rope_layer_impl.h"
+
+#include <glog/logging.h>
+
+#include "atb_speed/utils/tensor_util.h"
+
+namespace xllm {
+namespace layer {
+
+NpuRopeLayerImpl::NpuRopeLayerImpl(const ModelContext& context)
+    : BaseLayer(context) {
+  const auto& args = context.get_model_args();
+  head_dim_ = args.mm_head_dim();
+  if (head_dim_ <= 0) {
+    head_dim_ = args.mm_hidden_size() / args.mm_num_attention_heads();
+  }
+  CHECK_GT(head_dim_, 0);
+  output_tensors_.resize(2);
+  init_layer();
+}
+
+int64_t NpuRopeLayerImpl::init_layer() {
+  name_ = "rope_layer";
+  return init_node(rope_node_);
+}
+
+int64_t NpuRopeLayerImpl::init_node(atb_speed::Model::Node& node) {
+  atb::infer::RopeParam rope_param;
+  rope_param.rotaryCoeff = 2;
+
+  atb::Operation* operation = nullptr;
+  const atb::Status status = atb::CreateOperation(rope_param, &operation);
+  if (status != atb::NO_ERROR || operation == nullptr) {
+    LOG(ERROR) << "Failed to create RopeOperation, status=" << status;
+    return status == atb::NO_ERROR ? -1 : status;
+  }
+
+  node.operation.reset(operation);
+  CHECK_EQ(node.operation->GetInputNum(), 5);
+  CHECK_EQ(node.operation->GetOutputNum(), 2);
+  node.inTensors.resize(node.operation->GetInputNum());
+  node.outTensors.resize(node.operation->GetOutputNum());
+  node.variantPack.inTensors.resize(node.operation->GetInputNum());
+  node.variantPack.outTensors.resize(node.operation->GetOutputNum());
+  return atb::NO_ERROR;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> NpuRopeLayerImpl::forward(
+    const torch::Tensor& query,
+    const torch::Tensor& key,
+    const torch::Tensor& cos,
+    const torch::Tensor& sin,
+    const torch::Tensor& seq_len,
+    std::vector<int32_t>& seq_len_host,
+    int32_t node_id) {
+  build_node_variant_pack(
+      rope_node_, query, key, cos, sin, seq_len, seq_len_host);
+  const atb::Status status = execute_node(rope_node_, node_id);
+  LOG_IF(FATAL, status != atb::NO_ERROR)
+      << "RopeOperation execute failed, status=" << status;
+  return {output_tensors_.at(0), output_tensors_.at(1)};
+}
+
+void NpuRopeLayerImpl::build_node_variant_pack(
+    atb_speed::Model::Node& node,
+    const torch::Tensor& query,
+    const torch::Tensor& key,
+    const torch::Tensor& cos,
+    const torch::Tensor& sin,
+    const torch::Tensor& seq_len,
+    std::vector<int32_t>& seq_len_host) {
+  CHECK(query.dim() == 2 || query.dim() == 3);
+  CHECK_EQ(query.dim(), key.dim());
+  CHECK_EQ(query.sizes(), key.sizes());
+  CHECK_EQ(cos.sizes(), sin.sizes());
+  CHECK_EQ(query.size(0), cos.size(0));
+  CHECK_EQ(head_dim_ % 2, 0);
+  CHECK_EQ(seq_len_host.size(), static_cast<size_t>(seq_len.numel()));
+
+  if (query.dim() == 3) {
+    CHECK_EQ(query.size(-1), head_dim_);
+  } else {
+    CHECK_EQ(query.size(-1) % head_dim_, 0);
+  }
+  const int64_t rotary_dim = head_dim_ / 2;
+  CHECK(cos.size(-1) == head_dim_ || cos.size(-1) == rotary_dim);
+  auto to_nd = [](const torch::Tensor& tensor) {
+    auto contiguous = tensor.is_contiguous() ? tensor : tensor.contiguous();
+#ifdef TORCH_HIGHER_THAN_PTA6
+    if (at_npu::native::get_npu_format(contiguous) != ACL_FORMAT_ND) {
+      return at_npu::native::npu_format_cast(contiguous, ACL_FORMAT_ND);
+    }
+#else
+    if (at_npu::native::NPUNativeFunctions::get_npu_format(contiguous) !=
+        ACL_FORMAT_ND) {
+      return at_npu::native::NPUNativeFunctions::npu_format_cast(contiguous,
+                                                                 ACL_FORMAT_ND);
+    }
+#endif
+    return contiguous;
+  };
+  auto query_input = to_nd(query);
+  auto key_input = to_nd(key);
+  CHECK(query_input.scalar_type() == torch::kFloat16 ||
+        query_input.scalar_type() == torch::kBFloat16)
+      << "ATB Rope only supports float16/bfloat16 query and key";
+  CHECK_EQ(key_input.scalar_type(), query_input.scalar_type());
+  internal_query_ = atb_speed::Utils::AtTensor2Tensor(query_input);
+  internal_key_ = atb_speed::Utils::AtTensor2Tensor(key_input);
+  CHECK_EQ(cos.scalar_type(), query_input.scalar_type());
+  CHECK_EQ(sin.scalar_type(), query_input.scalar_type());
+  CHECK_EQ(seq_len.scalar_type(), torch::kInt);
+  auto cos_atb = to_nd(cos);
+  auto sin_atb = to_nd(sin);
+  auto seq_len_atb = to_nd(seq_len);
+  CHECK_EQ(cos_atb.scalar_type(), query_input.scalar_type());
+  CHECK_EQ(sin_atb.scalar_type(), query_input.scalar_type());
+  CHECK_EQ(seq_len_atb.scalar_type(), torch::kInt);
+  internal_cos_ = atb_speed::Utils::AtTensor2Tensor(cos_atb);
+  internal_sin_ = atb_speed::Utils::AtTensor2Tensor(sin_atb);
+  internal_seq_len_ = atb_speed::Utils::AtTensor2Tensor(seq_len_atb);
+
+  node.variantPack.inTensors.at(0) = internal_query_;
+  node.variantPack.inTensors.at(1) = internal_key_;
+  node.variantPack.inTensors.at(2) = internal_cos_;
+  node.variantPack.inTensors.at(3) = internal_sin_;
+  node.variantPack.inTensors.at(4) = internal_seq_len_;
+  node.variantPack.inTensors.at(4).hostData = seq_len_host.data();
+
+  atb::SVector<atb::TensorDesc> input_descs;
+  input_descs.resize(node.operation->GetInputNum());
+  input_descs.at(0) = internal_query_.desc;
+  input_descs.at(1) = internal_key_.desc;
+  input_descs.at(2) = internal_cos_.desc;
+  input_descs.at(3) = internal_sin_.desc;
+  input_descs.at(4) = internal_seq_len_.desc;
+  atb::SVector<atb::TensorDesc> output_descs;
+  output_descs.resize(node.operation->GetOutputNum());
+  const atb::Status status =
+      node.operation->InferShape(input_descs, output_descs);
+  CHECK_EQ(status, atb::NO_ERROR);
+
+  output_tensors_.at(0) =
+      atb_speed::Utils::CreateAtTensorFromTensorDesc(output_descs.at(0));
+  output_tensors_.at(1) =
+      atb_speed::Utils::CreateAtTensorFromTensorDesc(output_descs.at(1));
+  node.variantPack.outTensors.at(0) =
+      atb_speed::Utils::AtTensor2Tensor(output_tensors_.at(0));
+  node.variantPack.outTensors.at(1) =
+      atb_speed::Utils::AtTensor2Tensor(output_tensors_.at(1));
+}
+
+}  // namespace layer
+}  // namespace xllm
diff --git a/xllm/core/layers/npu/npu_rope_layer_impl.h b/xllm/core/layers/npu/npu_rope_layer_impl.h
new file mode 100644
index 0000000000..8457660e1a
--- /dev/null
+++ b/xllm/core/layers/npu/npu_rope_layer_impl.h
@@ -0,0 +1,72 @@
+/* Copyright 2026 The xLLM Authors.
+
+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 <torch/torch.h>
+
+#include <cstdint>
+#include <tuple>
+#include <vector>
+
+#include "atb/atb_infer.h"
+#include "atb_speed/base/model.h"
+#include "npu_base_layer.h"
+
+namespace xllm {
+namespace layer {
+
+class NpuRopeLayerImpl : public BaseLayer {
+ public:
+  explicit NpuRopeLayerImpl(const ModelContext& context);
+
+  ~NpuRopeLayerImpl() override = default;
+
+  std::tuple<torch::Tensor, torch::Tensor> forward(
+      const torch::Tensor& query,
+      const torch::Tensor& key,
+      const torch::Tensor& cos,
+      const torch::Tensor& sin,
+      const torch::Tensor& seq_len,
+      std::vector<int32_t>& seq_len_host,
+      int32_t node_id = 0);
+
+ private:
+  int64_t init_layer() override;
+
+  int64_t init_node(atb_speed::Model::Node& node);
+
+  void build_node_variant_pack(atb_speed::Model::Node& node,
+                               const torch::Tensor& query,
+                               const torch::Tensor& key,
+                               const torch::Tensor& cos,
+                               const torch::Tensor& sin,
+                               const torch::Tensor& seq_len,
+                               std::vector<int32_t>& seq_len_host);
+
+  atb_speed::Model::Node rope_node_;
+  std::vector<at::Tensor> output_tensors_;
+  atb::Tensor internal_query_;
+  atb::Tensor internal_key_;
+  atb::Tensor internal_cos_;
+  atb::Tensor internal_sin_;
+  atb::Tensor internal_seq_len_;
+  int64_t head_dim_ = 0;
+};
+
+TORCH_MODULE(NpuRopeLayer);
+
+}  // namespace layer
+}  // namespace xllm
diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp
index df42b89b44..df5fe7283b 100644
--- a/xllm/core/layers/npu_torch/attention.cpp
+++ b/xllm/core/layers/npu_torch/attention.cpp
@@ -49,14 +49,13 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
     return std::make_tuple(output, output_lse);
   }
 
-  bool only_prefill =
+  const bool only_prefill =
       attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
 
   torch::Tensor k_cache = kv_cache.get_k_cache();
   torch::Tensor v = value.view({-1, num_kv_heads_, head_size_});
   std::optional<torch::Tensor> v_cache = kv_cache.get_v_cache();
 
-  // Reshape and cache key/value
   xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
   reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_});
   reshape_paged_cache_params.value = v;
diff --git a/xllm/core/layers/qwen3_vision_layer.cpp b/xllm/core/layers/qwen3_vision_layer.cpp
index afd1783399..23013ea69a 100644
--- a/xllm/core/layers/qwen3_vision_layer.cpp
+++ b/xllm/core/layers/qwen3_vision_layer.cpp
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "qwen3_vision_layer.h"
+#include "core/layers/qwen3_vision_layer.h"
 
 namespace xllm {
 namespace layer {
@@ -22,12 +22,57 @@ Qwen3_VisionLayerImpl::Qwen3_VisionLayerImpl(const ModelContext& context)
     : Qwen2_5_VisionLayerImpl(context, true) {}
 
 void Qwen3_VisionLayerImpl::load_state_dict(const StateDict& state_dict) {
-  attention_->load_state_dict(state_dict.get_dict_with_prefix("attn."));
-  mlp_->load_state_dict(
-      state_dict.get_dict_with_prefix("mlp."), {"linear_fc1."}, "linear_fc2.");
-  norm1_->load_state_dict(state_dict.get_dict_with_prefix("norm1."));
-  norm2_->load_state_dict(state_dict.get_dict_with_prefix("norm2."));
+  const auto attention_dict = state_dict.get_dict_with_prefix("attn.");
+  const auto mlp_dict = state_dict.get_dict_with_prefix("mlp.");
+  const auto norm1_dict = state_dict.get_dict_with_prefix("norm1.");
+  const auto norm2_dict = state_dict.get_dict_with_prefix("norm2.");
+
+  attention_->load_state_dict(attention_dict);
+  mlp_->load_state_dict(mlp_dict, {"linear_fc1."}, "linear_fc2.");
+  norm1_->load_state_dict(norm1_dict);
+  norm2_->load_state_dict(norm2_dict);
+
+  attention_qkv_weight_loaded_ |= attention_dict.has("qkv.weight");
+  attention_qkv_bias_loaded_ |= attention_dict.has("qkv.bias");
+  attention_proj_weight_loaded_ |= attention_dict.has("proj.weight");
+  attention_proj_bias_loaded_ |= attention_dict.has("proj.bias");
+  mlp_fc1_weight_loaded_ |= mlp_dict.has("linear_fc1.weight");
+  mlp_fc1_bias_loaded_ |= mlp_dict.has("linear_fc1.bias");
+  mlp_fc2_weight_loaded_ |= mlp_dict.has("linear_fc2.weight");
+  mlp_fc2_bias_loaded_ |= mlp_dict.has("linear_fc2.bias");
+  norm1_weight_loaded_ |= norm1_dict.has("weight");
+  norm1_bias_loaded_ |= norm1_dict.has("bias");
+  norm2_weight_loaded_ |= norm2_dict.has("weight");
+  norm2_bias_loaded_ |= norm2_dict.has("bias");
+}
+
+void Qwen3_VisionLayerImpl::verify_loaded_weights(
+    const std::string& prefix) const {
+  CHECK(attention_qkv_weight_loaded_)
+      << "weight is not loaded for " << prefix + "attn.qkv.weight";
+  CHECK(attention_qkv_bias_loaded_)
+      << "bias is not loaded for " << prefix + "attn.qkv.bias";
+  CHECK(attention_proj_weight_loaded_)
+      << "weight is not loaded for " << prefix + "attn.proj.weight";
+  CHECK(attention_proj_bias_loaded_)
+      << "bias is not loaded for " << prefix + "attn.proj.bias";
+  CHECK(mlp_fc1_weight_loaded_)
+      << "weight is not loaded for " << prefix + "mlp.linear_fc1.weight";
+  CHECK(mlp_fc1_bias_loaded_)
+      << "bias is not loaded for " << prefix + "mlp.linear_fc1.bias";
+  CHECK(mlp_fc2_weight_loaded_)
+      << "weight is not loaded for " << prefix + "mlp.linear_fc2.weight";
+  CHECK(mlp_fc2_bias_loaded_)
+      << "bias is not loaded for " << prefix + "mlp.linear_fc2.bias";
+  CHECK(norm1_weight_loaded_)
+      << "weight is not loaded for " << prefix + "norm1.weight";
+  CHECK(norm1_bias_loaded_)
+      << "bias is not loaded for " << prefix + "norm1.bias";
+  CHECK(norm2_weight_loaded_)
+      << "weight is not loaded for " << prefix + "norm2.weight";
+  CHECK(norm2_bias_loaded_)
+      << "bias is not loaded for " << prefix + "norm2.bias";
 }
 
 }  // namespace layer
-}  // namespace xllm
\ No newline at end of file
+}  // namespace xllm
diff --git a/xllm/core/layers/qwen3_vision_layer.h b/xllm/core/layers/qwen3_vision_layer.h
index aa13cb2716..ee6d614cdc 100644
--- a/xllm/core/layers/qwen3_vision_layer.h
+++ b/xllm/core/layers/qwen3_vision_layer.h
@@ -14,18 +14,37 @@ limitations under the License.
 ==============================================================================*/
 
 #pragma once
-#include "qwen2_5_vision_layer.h"
+
+#include <string>
+
+#include "core/layers/qwen2_5_vision_layer.h"
 
 namespace xllm {
 namespace layer {
 
-class Qwen3_VisionLayerImpl : public Qwen2_5_VisionLayerImpl {
+class Qwen3_VisionLayerImpl final : public Qwen2_5_VisionLayerImpl {
  public:
-  Qwen3_VisionLayerImpl(const ModelContext& context);
+  explicit Qwen3_VisionLayerImpl(const ModelContext& context);
 
   void load_state_dict(const StateDict& state_dict);
+
+  void verify_loaded_weights(const std::string& prefix) const;
+
+ private:
+  bool norm1_weight_loaded_ = false;
+  bool norm1_bias_loaded_ = false;
+  bool norm2_weight_loaded_ = false;
+  bool norm2_bias_loaded_ = false;
+  bool attention_qkv_weight_loaded_ = false;
+  bool attention_qkv_bias_loaded_ = false;
+  bool attention_proj_weight_loaded_ = false;
+  bool attention_proj_bias_loaded_ = false;
+  bool mlp_fc1_weight_loaded_ = false;
+  bool mlp_fc1_bias_loaded_ = false;
+  bool mlp_fc2_weight_loaded_ = false;
+  bool mlp_fc2_bias_loaded_ = false;
 };
 TORCH_MODULE(Qwen3_VisionLayer);
 
 }  // namespace layer
-}  // namespace xllm
\ No newline at end of file
+}  // namespace xllm
diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h
index 0d89b738b2..70ee0deb50 100644
--- a/xllm/core/runtime/options.h
+++ b/xllm/core/runtime/options.h
@@ -135,6 +135,10 @@ struct Options {
   // Default set as 1
   PROPERTY(int32_t, vae_size) = 1;
 
+  // text encoder tensor parallelism size
+  // Default set as 1
+  PROPERTY(int32_t, text_encoder_tp_size) = 1;
+
   // enable enable_schedule_overlap to improve runtime execution efficiency.
   PROPERTY(bool, enable_schedule_overlap) = true;
 
diff --git a/xllm/models/dit/autoencoders/autoencoder_kl_wan.h b/xllm/models/dit/autoencoders/autoencoder_kl_wan.h
index 2d98d5ece7..770085898d 100644
--- a/xllm/models/dit/autoencoders/autoencoder_kl_wan.h
+++ b/xllm/models/dit/autoencoders/autoencoder_kl_wan.h
@@ -268,9 +268,17 @@ class WanRMSNormImpl : public torch::nn::Module {
 
   torch::Tensor forward(const torch::Tensor& x) {
     int64_t norm_dim = channel_first_ ? 1 : -1;
-    auto normed = torch::nn::functional::normalize(
-        x,
-        torch::nn::functional::NormalizeFuncOptions().dim(norm_dim).eps(1e-12));
+    // Match diffusers: normalize fp16/bf16 inputs in fp32 for numerical
+    // stability, then cast the normalized activations back to the input dtype.
+    const bool needs_fp32_normalize =
+        x.dtype() == torch::kFloat16 || x.dtype() == torch::kBFloat16;
+    auto norm_input = needs_fp32_normalize ? x.to(torch::kFloat32) : x;
+    auto normed =
+        torch::nn::functional::normalize(
+            norm_input,
+            torch::nn::functional::NormalizeFuncOptions().dim(norm_dim).eps(
+                1e-12))
+            .to(x.dtype());
     auto out = normed * scale_ * gamma_;
     if (bias_enabled_) {
       out = out + bias_;
@@ -313,9 +321,11 @@ class WanUpsampleImpl : public torch::nn::Module {
       : options_(options) {}
 
   torch::Tensor forward(const torch::Tensor& x) {
+    // Match diffusers WanUpsample: interpolation is performed in fp32, then
+    // cast back to the input dtype.
     auto result =
         torch::nn::functional::interpolate(x.to(torch::kFloat32), options_);
-    return result;
+    return result.to(x.dtype());
   }
 
  private:
@@ -1633,8 +1643,10 @@ class AutoencoderKLWanImpl : public torch::nn::Module {
   }
 
   torch::Tensor encode_(const torch::Tensor& videos) {
-    auto orig_dtype = videos.dtype();
-    auto x = videos.to(torch::kFloat32);
+    // The Wan VAE follows the dtype selected by the owning pipeline's
+    // TensorOptions. Individual modules (RMSNorm, upsample) still upcast their
+    // numerically-sensitive intermediate math to fp32 as diffusers does.
+    auto x = videos.to(device_, dtype_);
     if (parallel_ctx_.is_parallel()) {
       x = parallel_ctx_.split(x);
     }
@@ -1677,7 +1689,7 @@ class AutoencoderKLWanImpl : public torch::nn::Module {
       out = parallel_ctx_.merge(out);
     }
     clear_cache();
-    return out.to(orig_dtype);
+    return out;
   }
 
   AutoencoderKLOutput encode(const torch::Tensor& videos) {
@@ -1687,8 +1699,9 @@ class AutoencoderKLWanImpl : public torch::nn::Module {
   }
 
   DecoderOutput decode_(const torch::Tensor& latents) {
-    auto orig_dtype = latents.dtype();
-    torch::Tensor processed_latents = latents.to(torch::kFloat32);
+    // Follow the pipeline-selected dtype for the VAE module; fp32-only
+    // intermediate ops are handled locally in those modules.
+    torch::Tensor processed_latents = latents.to(device_, dtype_);
     if (parallel_ctx_.is_parallel()) {
       processed_latents = parallel_ctx_.split(processed_latents);
     }
@@ -1716,7 +1729,7 @@ class AutoencoderKLWanImpl : public torch::nn::Module {
     auto dec = torch::clamp(out, -1.0f, 1.0f);
 
     clear_cache();
-    return DecoderOutput(dec.to(orig_dtype));
+    return DecoderOutput(dec);
   }
 
   DecoderOutput decode(
@@ -1740,10 +1753,10 @@ class AutoencoderKLWanImpl : public torch::nn::Module {
   }
 
   void load_model(std::unique_ptr<DiTFolderLoader> loader) {
-    encoder_->to(torch::kFloat32);
-    decoder_->to(torch::kFloat32);
-    quant_conv_->to(torch::kFloat32);
-    post_quant_conv_->to(torch::kFloat32);
+    encoder_->to(device_, dtype_);
+    decoder_->to(device_, dtype_);
+    quant_conv_->to(device_, dtype_);
+    post_quant_conv_->to(device_, dtype_);
 
     for (const auto& state_dict : loader->get_state_dicts()) {
       encoder_->load_state_dict(state_dict->get_dict_with_prefix("encoder."));
diff --git a/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h
new file mode 100644
index 0000000000..118c5a0ae8
--- /dev/null
+++ b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h
@@ -0,0 +1,917 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+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 <torch/torch.h>
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <limits>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "core/framework/config/kernel_config.h"
+#include "core/framework/dit_cache/dit_cache.h"
+#include "core/framework/dit_model_loader.h"
+#include "core/framework/kv_cache/kv_cache.h"
+#include "core/framework/model/model_input_params.h"
+#include "core/framework/model_context.h"
+#include "core/framework/multimodal/mm_input.h"
+#include "core/framework/multimodal/mm_visitor.h"
+#include "core/framework/parallel_state/process_group.h"
+#include "core/framework/state_dict/state_dict.h"
+#include "core/framework/tokenizer/tokenizer.h"
+#include "core/runtime/dit_forward_params.h"
+#include "models/dit/autoencoders/autoencoder_kl_wan.h"
+#include "models/dit/processors/vae_image_processor.h"
+#include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h"
+#include "models/dit/transformers/transformer_joyimage_edit_plus.h"
+#include "models/dit/utils/dit_parallel_mixin.h"
+#include "models/dit/utils/util.h"
+#include "models/model_registry.h"
+#include "models/vlm/mposition/mposition.h"
+#include "models/vlm/qwen3_vl.h"
+#include "processors/multimodal_processor.h"
+#include "util/tensor_helper.h"
+
+namespace xllm {
+
+using JoyImageEditTextEncoder = Qwen3_VLForConditionalGeneration;
+
+class JoyImageEditPlusPipelineImpl : public torch::nn::Module,
+                                     public dit::CFGParallelMixin {
+ public:
+  using JoyImageShapeList = std::vector<std::vector<std::array<int64_t, 3>>>;
+
+  struct TransformerForwardContext {
+    torch::Tensor rope_cos;
+    torch::Tensor rope_sin;
+    torch::Tensor attention_mask;
+  };
+
+  JoyImageEditPlusPipelineImpl(const DiTModelContext& context)
+      : dit::CFGParallelMixin(context),
+        context_(context),
+        parallel_args_(context.get_parallel_args()),
+        vae_model_args_(context.get_model_args("vae")) {
+    options_ = context.get_tensor_options();
+    dtype_ = options_.dtype().toScalarType();
+    device_ = options_.device();
+
+    const ModelArgs& transformer_model_args =
+        context.get_model_args("transformer");
+    in_channels_ = transformer_model_args.in_channels();
+    num_layers_ = transformer_model_args.num_layers();
+    const int64_t hidden_size = transformer_model_args.hidden_size();
+    const int64_t num_heads = transformer_model_args.num_attention_heads();
+    CHECK_EQ(hidden_size % num_heads, 0)
+        << "hidden_size must be divisible by num_attention_heads";
+    head_dim_ = hidden_size / num_heads;
+    rope_dim_list_ = transformer_model_args.rope_dim_list();
+    theta_ = transformer_model_args.rope_theta_dit();
+    auto patch = transformer_model_args.wan_patch_size();
+    patch_t_ = patch[0];
+    patch_h_ = patch[1];
+    patch_w_ = patch[2];
+
+    latent_channels_ = vae_model_args_.z_dim();
+    vae_scale_factor_spatial_ = vae_model_args_.vae_scale_factor_spatial();
+    if (vae_scale_factor_spatial_ <= 0) {
+      vae_scale_factor_spatial_ = 8;
+    }
+
+    vae_ = AutoencoderKLWan(context.get_model_context("vae"));
+    transformer_ = joyimage::JoyImageEditPlusTransformer3DModel(
+        context.get_model_context("transformer"), parallel_args_);
+    scheduler_ =
+        FlowMatchEulerDiscreteScheduler(context.get_model_context("scheduler"));
+
+    if (context.has_component("text_encoder")) {
+      text_encoder_model_args_ = context.get_model_args("text_encoder");
+#if defined(USE_NPU)
+      CHECK_EQ(KernelConfig::get_instance().npu_kernel_backend(), "TORCH")
+          << "JoyImageEditPlus in-process Qwen3-VL text encoding requires "
+             "--npu_kernel_backend=TORCH.";
+#endif
+      ProcessGroup* tp_group = parallel_args_.dit_text_encoder_tp_group_;
+      CHECK(tp_group != nullptr)
+          << "JoyImageEditPlus requires an injected text encoder TP group.";
+      ParallelArgs vlm_parallel_args(
+          tp_group->rank(), tp_group->world_size(), tp_group);
+      vlm_parallel_args.tp_size(tp_group->world_size());
+      vlm_parallel_args.tp_group_ = tp_group;
+      text_encoder_empty_kv_caches_.resize(
+          static_cast<size_t>(text_encoder_model_args_.n_layers()));
+      const ModelContext& source_vlm_context =
+          context.get_model_context("text_encoder");
+      ModelContext vlm_context(vlm_parallel_args,
+                               source_vlm_context.get_model_args(),
+                               source_vlm_context.get_quant_args(),
+                               source_vlm_context.get_tensor_options());
+      text_encoder_ = JoyImageEditTextEncoder(vlm_context);
+      CHECK(!text_encoder_.is_empty())
+          << "Failed to create JoyImageEditPlus Qwen3-VL text encoder";
+      mposition_generator_ =
+          MPositionGeneratorFactory::get_instance().create_mposition_generator(
+              "qwen3_vl");
+    }
+
+    vae_image_processor_ =
+        xllm::VAEImageProcessor(context.get_model_context("vae"),
+                                /*do_resize=*/true,
+                                /*do_normalize=*/true,
+                                /*do_binarize=*/false,
+                                /*do_convert_rgb=*/false,
+                                /*do_convert_grayscale=*/false,
+                                latent_channels_,
+                                /*scale_factor=*/vae_scale_factor_spatial_);
+
+    register_module("vae", vae_);
+    register_module("scheduler", scheduler_);
+    register_module("transformer", transformer_);
+    if (!text_encoder_.is_empty()) {
+      register_module("text_encoder", text_encoder_);
+    }
+    register_module("vae_image_processor", vae_image_processor_);
+
+    latents_mean_ = vae_model_args_.latents_mean();
+    latents_std_ = vae_model_args_.latents_std();
+  }
+
+  torch::Tensor latents_mean_tensor(const torch::Tensor& ref) const {
+    return torch::tensor(latents_mean_, torch::kFloat32)
+        .view({1, latent_channels_, 1, 1, 1})
+        .to(ref.device(), ref.dtype());
+  }
+  torch::Tensor latents_std_tensor(const torch::Tensor& ref) const {
+    return torch::tensor(latents_std_, torch::kFloat32)
+        .view({1, latent_channels_, 1, 1, 1})
+        .to(ref.device(), ref.dtype());
+  }
+
+  // Patchify a [C, T, H, W] latent into [num_patches, C, pt, ph, pw] and
+  // return {patches, (lt, lh, lw)}.
+  std::pair<torch::Tensor, std::array<int64_t, 3>> patchify(
+      const torch::Tensor& item) {
+    int64_t c = item.size(0), t = item.size(1), h = item.size(2),
+            w = item.size(3);
+    int64_t lt = t / patch_t_, lh = h / patch_h_, lw = w / patch_w_;
+    auto p = item.reshape({c, lt, patch_t_, lh, patch_h_, lw, patch_w_});
+    p = p.permute({1, 3, 5, 0, 2, 4, 6})
+            .reshape({-1, c, patch_t_, patch_h_, patch_w_});
+    return {p, {lt, lh, lw}};
+  }
+
+  // Build the 6D padded latent tensor: target noise + reference latents.
+  // Returns {padded_latents[B,N,C,pt,ph,pw], target_mask[B,N],
+  //          shape_list (per-sample list of (t,h,w))}.
+  std::tuple<torch::Tensor, torch::Tensor, JoyImageShapeList> prepare_latents(
+      int64_t batch_size,
+      int64_t num_channels_latents,
+      int64_t height,
+      int64_t width,
+      int64_t seed,
+      const std::vector<std::vector<torch::Tensor>>& reference_images,
+      const torch::Tensor& provided_latents) {
+    std::vector<torch::Tensor> all_patches;
+    std::vector<torch::Tensor> all_target_masks;
+    std::vector<std::vector<std::array<int64_t, 3>>> all_shapes;
+    int64_t max_patches = 0;
+
+    int64_t h_target = height / vae_scale_factor_spatial_;
+    int64_t w_target = width / vae_scale_factor_spatial_;
+
+    for (int64_t b = 0; b < batch_size; ++b) {
+      std::vector<torch::Tensor> items;
+      // Target noise: [C, 1, h', w'].
+      torch::Tensor noise;
+      if (provided_latents.defined()) {
+        noise = provided_latents[b].to(device_, dtype_);
+      } else {
+        noise = xllm::dit::randn_tensor(
+            {num_channels_latents, 1, h_target, w_target}, seed + b, options_);
+      }
+      items.push_back(noise);
+
+      // References: VAE-encode each, normalize, squeeze to [C, 1, h', w'].
+      if (b < static_cast<int64_t>(reference_images.size())) {
+        for (const auto& ref_img : reference_images[b]) {
+          auto ref = ref_img.to(device_, dtype_);
+          if (ref.dim() == 4) ref = ref.unsqueeze(2);  // [B,C,H,W]->[B,C,1,H,W]
+          auto lat = vae_->encode(ref.to(dtype_)).latent_dist.mode();
+          lat = lat.to(dtype_);
+          lat = (lat - latents_mean_tensor(lat)) / latents_std_tensor(lat);
+          items.push_back(lat.squeeze(0));  // [C, 1, h', w']
+        }
+      }
+
+      std::vector<torch::Tensor> sample_patches;
+      std::vector<torch::Tensor> sample_masks;
+      std::vector<std::array<int64_t, 3>> sample_shapes;
+      for (size_t j = 0; j < items.size(); ++j) {
+        auto pr = patchify(items[j]);
+        sample_shapes.push_back(pr.second);
+        sample_patches.push_back(pr.first);
+        auto n = pr.first.size(0);
+        sample_masks.push_back(torch::full(
+            {n},
+            /*value=*/(j == 0),
+            torch::TensorOptions().device(device_).dtype(torch::kBool)));
+      }
+      auto combined = torch::cat(sample_patches, 0);
+      auto combined_mask = torch::cat(sample_masks, 0);
+      all_patches.push_back(combined);
+      all_target_masks.push_back(combined_mask);
+      all_shapes.push_back(sample_shapes);
+      max_patches = std::max(max_patches, combined.size(0));
+    }
+
+    auto padded = torch::zeros({batch_size,
+                                max_patches,
+                                num_channels_latents,
+                                patch_t_,
+                                patch_h_,
+                                patch_w_},
+                               options_);
+    auto target_mask = torch::zeros(
+        {batch_size, max_patches},
+        torch::TensorOptions().device(device_).dtype(torch::kBool));
+    for (int64_t b = 0; b < batch_size; ++b) {
+      int64_t n = all_patches[b].size(0);
+      padded.index_put_({b, torch::indexing::Slice(0, n)}, all_patches[b]);
+      target_mask.index_put_({b, torch::indexing::Slice(0, n)},
+                             all_target_masks[b]);
+    }
+    return std::make_tuple(padded, target_mask, all_shapes);
+  }
+
+  DiTForwardOutput forward(const DiTForwardInput& input) {
+    torch::NoGradGuard no_grad;
+    const auto& gp = input.generation_params;
+    int64_t num_inference_steps = gp.num_inference_steps;
+    double guidance_scale =
+        gp.true_cfg_scale > 0 ? gp.true_cfg_scale : gp.guidance_scale;
+    int64_t seed = gp.seed >= 0 ? gp.seed : 42;
+
+    // Collect reference images (one sample per batch entry).
+    std::vector<torch::Tensor> raw_images;
+    if (!input.images_list.empty()) {
+      raw_images = input.images_list;
+    } else if (input.images.defined()) {
+      raw_images.push_back(input.images);
+    } else {
+      LOG(FATAL) << "JoyImageEditPlus requires reference images";
+    }
+
+    int64_t batch_size = input.batch_size;
+    if (batch_size <= 0) {
+      if (input.prompt_embeds.defined()) {
+        batch_size = input.prompt_embeds.size(0);
+      } else if (!input.prompts.empty()) {
+        batch_size = static_cast<int64_t>(input.prompts.size());
+      } else {
+        batch_size = raw_images[0].dim() == 4 ? raw_images[0].size(0) : 1;
+      }
+    }
+
+    // Determine output resolution from the last reference image if unset.
+    int64_t height = gp.height;
+    int64_t width = gp.width;
+    if (height <= 0 || width <= 0) {
+      const auto& last = raw_images.back();
+      int64_t ih = last.size(last.dim() - 2);
+      int64_t iw = last.size(last.dim() - 1);
+      auto hw = joyimage_bucket(ih, iw);
+      height = hw.first;
+      width = hw.second;
+    }
+    height = (height / vae_scale_factor_spatial_) * vae_scale_factor_spatial_;
+    width = (width / vae_scale_factor_spatial_) * vae_scale_factor_spatial_;
+
+    // Reference images per sample, in two forms:
+    //  - vae_refs: VAE-preprocessed to [-1,1] (for latent encoding), each
+    //    bucket-resized to its own aspect bucket.
+    std::vector<std::vector<torch::Tensor>> vae_refs(batch_size);
+    for (int64_t b = 0; b < batch_size; ++b) {
+      for (const auto& imgs : raw_images) {
+        auto img = imgs.dim() == 4 ? imgs[b] : imgs;  // [C,H,W]
+        int64_t ih = img.size(1), iw = img.size(2);
+        auto hw = joyimage_bucket(ih, iw);
+        auto img4 = img.unsqueeze(0).to(device_);
+        vae_refs[b].push_back(vae_image_processor_->preprocess(
+            img4, hw.first, hw.second, /*resize_mode=*/"lanczos"));
+      }
+    }
+    bool do_cfg = guidance_scale > 1.0;
+    torch::Tensor prompt_embeds;
+    torch::Tensor prompt_embeds_mask;
+    const bool encode_text =
+        !input.prompt_embeds.defined() ||
+        (do_cfg && !input.negative_prompt_embeds.defined());
+    if (encode_text) {
+      CHECK(!text_encoder_.is_empty()) << "Qwen3-VL text encoder is not loaded";
+    }
+    if (input.prompt_embeds.defined()) {
+      prompt_embeds = input.prompt_embeds.to(options_.device(), dtype_);
+      prompt_embeds_mask = torch::ones(
+          {prompt_embeds.size(0), prompt_embeds.size(1)},
+          torch::TensorOptions().device(device_).dtype(torch::kLong));
+    } else {
+      CHECK(!input.prompts.empty())
+          << "JoyImageEditPlus requires `prompts` or `prompt_embeds`";
+      std::tie(prompt_embeds, prompt_embeds_mask) =
+          encode_prompts(input.prompts, raw_images, batch_size);
+    }
+
+    torch::Tensor neg_embeds, neg_embeds_mask;
+    if (do_cfg) {
+      if (input.negative_prompt_embeds.defined()) {
+        neg_embeds = input.negative_prompt_embeds.to(options_.device(), dtype_);
+        neg_embeds_mask = torch::ones(
+            {neg_embeds.size(0), neg_embeds.size(1)},
+            torch::TensorOptions().device(device_).dtype(torch::kLong));
+      } else {
+        std::vector<std::string> negative_prompts = input.negative_prompts;
+        if (negative_prompts.empty()) {
+          negative_prompts.resize(static_cast<size_t>(batch_size));
+        }
+        std::tie(neg_embeds, neg_embeds_mask) =
+            encode_prompts(negative_prompts, raw_images, batch_size);
+      }
+      // Pad/concat [negative, positive] to equal sequence length.
+      int64_t max_l = std::max(prompt_embeds.size(1), neg_embeds.size(1));
+      prompt_embeds = pad_seq(prompt_embeds, max_l);
+      neg_embeds = pad_seq(neg_embeds, max_l);
+      prompt_embeds_mask = pad_seq(prompt_embeds_mask, max_l);
+      neg_embeds_mask = pad_seq(neg_embeds_mask, max_l);
+    }
+    // Latents.
+    int64_t num_channels_latents = in_channels_;
+    auto lp = prepare_latents(batch_size,
+                              num_channels_latents,
+                              height,
+                              width,
+                              seed,
+                              vae_refs,
+                              input.latents);
+    auto latents = std::get<0>(lp);
+    auto target_mask = std::get<1>(lp);
+    auto shape_list = std::get<2>(lp);
+    auto clean_backup = latents.clone();
+    const std::array<int64_t, 3>& target_shape = shape_list.front().front();
+    const int64_t target_patch_count =
+        target_shape[0] * target_shape[1] * target_shape[2];
+
+    // Timesteps (static shift; no dynamic shifting for Joy).
+    scheduler_->set_timesteps(num_inference_steps, device_);
+    scheduler_->set_begin_index(0);
+    auto timesteps = scheduler_->timesteps();
+    DiTCache::get_instance().set_infer_steps(num_inference_steps);
+    DiTCache::get_instance().set_num_blocks(num_layers_);
+
+    TransformerForwardContext prompt_transformer_context =
+        prepare_transformer_context(batch_size,
+                                    latents.size(1),
+                                    latents.device(),
+                                    prompt_embeds_mask,
+                                    shape_list);
+    TransformerForwardContext negative_transformer_context;
+    if (do_cfg) {
+      negative_transformer_context =
+          prepare_transformer_context(batch_size,
+                                      latents.size(1),
+                                      latents.device(),
+                                      neg_embeds_mask,
+                                      shape_list);
+    }
+
+    for (int64_t i = 0; i < timesteps.size(0); ++i) {
+      auto t = timesteps[i];
+
+      torch::Tensor noise_pred;
+      if (do_cfg) {
+        torch::Tensor t_expand = t.repeat({batch_size});
+        auto [cond, uncond] = exec_with_cfg([&](bool is_positive) {
+          const torch::Tensor& encoder_hidden_states =
+              is_positive ? prompt_embeds : neg_embeds;
+          const TransformerForwardContext& transformer_context =
+              is_positive ? prompt_transformer_context
+                          : negative_transformer_context;
+          return transformer_forward(latents,
+                                     t_expand,
+                                     encoder_hidden_states,
+                                     transformer_context,
+                                     /*use_cfg=*/!is_positive,
+                                     /*step_index=*/i + 1);
+        });
+        auto comb = uncond + guidance_scale * (cond - uncond);
+        // Norm-rescale (diffusers): comb * (||cond|| / ||comb||) over channel
+        // dim (2) of the 6D [B, N, C, pt, ph, pw] prediction.
+        auto cond_norm =
+            torch::norm(cond, 2, std::vector<int64_t>{2}, /*keepdim=*/true);
+        auto noise_norm =
+            torch::norm(comb, 2, std::vector<int64_t>{2}, /*keepdim=*/true);
+        noise_pred = comb * (cond_norm / noise_norm.clamp_min(1e-6));
+      } else {
+        torch::Tensor t_expand = t.repeat({batch_size});
+        noise_pred = transformer_forward(latents,
+                                         t_expand,
+                                         prompt_embeds,
+                                         prompt_transformer_context,
+                                         /*use_cfg=*/false,
+                                         /*step_index=*/i + 1);
+      }
+
+      latents = scheduler_->step(noise_pred, t, latents).to(latents.dtype());
+      latents.slice(/*dim=*/1, target_patch_count, latents.size(1))
+          .copy_(clean_backup.slice(
+              /*dim=*/1, target_patch_count, clean_backup.size(1)));
+    }
+
+    // Decode target patches per sample.
+    std::vector<torch::Tensor> images;
+    for (int64_t b = 0; b < batch_size; ++b) {
+      auto thw = shape_list[b][0];
+      int64_t lt = thw[0], lh = thw[1], lw = thw[2];
+      int64_t target_len = lt * lh * lw;
+      auto patches = latents[b].slice(0, 0, target_len);  // [len, C, pt,ph,pw]
+      int64_t c = patches.size(1);
+      auto vid = patches.reshape({lt, lh, lw, c, patch_t_, patch_h_, patch_w_});
+      vid = vid.permute({3, 0, 4, 1, 5, 2, 6})
+                .reshape({1, c, lt * patch_t_, lh * patch_h_, lw * patch_w_});
+      vid = vid * latents_std_tensor(vid) + latents_mean_tensor(vid);
+      auto img = vae_->decode(vid.to(dtype_)).sample;       // [1,C,1,H,W]
+      img = img.to(torch::kFloat32).squeeze(0).squeeze(1);  // [C,H,W]
+      images.push_back(img.unsqueeze(0));
+    }
+    auto image = torch::cat(images, 0);
+    image = vae_image_processor_->postprocess(image);
+
+    DiTForwardOutput out;
+    out.tensors = torch::chunk(image, batch_size, 0);
+    return out;
+  }
+
+  void load_model(std::unique_ptr<DiTModelLoader> loader) {
+    LOG(INFO) << "JoyImageEditPlusPipeline loading from "
+              << loader->model_root_path();
+    auto transformer_loader = loader->take_component_loader("transformer");
+    auto vae_loader = loader->take_component_loader("vae");
+    std::unique_ptr<DiTFolderLoader> text_encoder_loader;
+    if (!text_encoder_.is_empty()) {
+      text_encoder_loader = loader->take_component_loader("text_encoder");
+      std::unique_ptr<Tokenizer> tokenizer = text_encoder_loader->tokenizer();
+      CHECK(tokenizer != nullptr)
+          << "Failed to load JoyImageEditPlus Qwen3-VL tokenizer";
+      tokenizer_ = std::shared_ptr<Tokenizer>(std::move(tokenizer));
+      multimodal_processor_ =
+          create_multimodal_processor(text_encoder_model_args_, tokenizer_);
+    }
+
+    vae_->load_model(std::move(vae_loader));
+    vae_->to(options_.device(), dtype_);
+
+    transformer_->load_model(std::move(transformer_loader));
+    transformer_->to(options_.device(), dtype_);
+    transformer_->keep_fp32_modules();
+
+    if (text_encoder_loader != nullptr) {
+      text_encoder_->load_model(std::move(text_encoder_loader));
+    }
+
+    LOG(INFO) << "JoyImageEditPlusPipeline loaded.";
+  }
+
+ private:
+  static constexpr int64_t kPromptEmbeddingStartIndex = 34;
+
+  std::string build_qwen3_vl_prompt(const std::string& prompt,
+                                    size_t image_count) const {
+    std::string chat_prompt =
+        "<|im_start|>system\n \\nDescribe the image by detailing the "
+        "color, shape, size, texture, quantity, text, spatial relationships "
+        "of the objects and background:<|im_end|>\n<|im_start|>user\n";
+    for (size_t image_index = 0; image_index < image_count; ++image_index) {
+      chat_prompt += "<|vision_start|><|image_pad|><|vision_end|>";
+    }
+    chat_prompt += prompt;
+    chat_prompt += "<|im_end|>\n<|im_start|>assistant\n";
+    return chat_prompt;
+  }
+
+  std::vector<torch::Tensor> prepare_vl_images(
+      const std::vector<torch::Tensor>& raw_images,
+      int64_t batch_index,
+      int64_t batch_size) const {
+    std::vector<torch::Tensor> images;
+    images.reserve(raw_images.size());
+    for (const torch::Tensor& raw_image : raw_images) {
+      CHECK(raw_image.dim() == 3 || raw_image.dim() == 4)
+          << "JoyImageEditPlus image input must be [C,H,W] or [B,C,H,W]";
+      torch::Tensor image;
+      if (raw_image.dim() == 4) {
+        CHECK_EQ(raw_image.size(0), batch_size)
+            << "JoyImageEditPlus image batch must match prompt batch";
+        image = raw_image[batch_index];
+      } else {
+        CHECK_EQ(batch_size, 1)
+            << "Unbatched JoyImageEditPlus images require batch_size=1";
+        image = raw_image;
+      }
+      image = image.to(torch::kCPU).to(torch::kFloat32);
+      const float max_value = image.max().item<float>();
+      image = max_value <= 1.1f ? image.clamp(0.0f, 1.0f) * 255.0f
+                                : image.clamp(0.0f, 255.0f);
+      images.emplace_back(std::move(image));
+    }
+    return images;
+  }
+
+  ModelInputParams build_text_encoder_input(const torch::Tensor& tokens,
+                                            const MMData& mm_data) {
+    CHECK_LE(tokens.numel(), std::numeric_limits<int32_t>::max())
+        << "JoyImageEditPlus Qwen3-VL prompt is too long";
+    const int32_t sequence_length = static_cast<int32_t>(tokens.numel());
+    CHECK_GT(sequence_length, 0)
+        << "JoyImageEditPlus Qwen3-VL prompt must not be empty";
+
+    ModelInputParams params;
+    params.meta.num_sequences = 1;
+    params.meta.actual_num_sequences = 1;
+    params.meta.q_max_seq_len = sequence_length;
+    params.meta.kv_max_seq_len = sequence_length;
+    params.meta.batch_forward_type = BatchForwardType::PREFILL;
+    params.attention.host.q_seq_lens = {sequence_length};
+    params.attention.host.kv_seq_lens = {sequence_length};
+#if defined(USE_NPU)
+    params.attention.host.q_cu_seq_lens = {sequence_length};
+#else
+    params.attention.host.q_cu_seq_lens = {0, sequence_length};
+#endif
+    params.attention.device.q_seq_lens =
+        torch::tensor({sequence_length}, torch::kInt).to(tokens.device());
+    params.attention.device.kv_seq_lens =
+        torch::tensor({sequence_length}, torch::kInt).to(tokens.device());
+#if defined(USE_NPU)
+    params.attention.device.q_cu_seq_lens =
+        torch::tensor({sequence_length}, torch::kInt).to(tokens.device());
+#else
+    params.attention.device.q_cu_seq_lens =
+        torch::tensor({0, sequence_length}, torch::kInt).to(tokens.device());
+#endif
+
+    MMBatchData mm_batch(std::vector<MMData>{mm_data});
+    EncoderInputGatherVisitor input_gather;
+    mm_batch.foreach (input_gather);
+    CHECK(input_gather.finish(mm_batch))
+        << "JoyImageEditPlus failed to gather Qwen3-VL encoder inputs";
+    mm_batch.to(tokens.device());
+
+    ModelInputParams multimodal_params;
+    multimodal_params.multimodal.mm_data = mm_batch;
+    MMDict multimodal_embeddings =
+        text_encoder_->get_multimodal_embeddings(multimodal_params);
+    EncoderOutputScatterVisitor output_scatter(multimodal_embeddings);
+    CHECK(mm_batch.foreach (output_scatter));
+    CHECK(output_scatter.finish())
+        << "JoyImageEditPlus failed to scatter Qwen3-VL embeddings";
+
+    EncoderEmbeddingGatherVisitor embedding_gather(
+        tokens.device(),
+        mm_batch.type(),
+        params.attention.host.kv_seq_lens,
+        params.attention.host.q_seq_lens);
+    CHECK(mm_batch.foreach (embedding_gather));
+    CHECK(embedding_gather.finish(mm_batch))
+        << "JoyImageEditPlus failed to gather Qwen3-VL embeddings";
+    params.multimodal.mm_data = std::move(mm_batch);
+    params.embedding.input_embedding =
+        text_encoder_->get_input_embeddings(tokens, params);
+    return params;
+  }
+
+  std::pair<torch::Tensor, torch::Tensor> encode_single_prompt(
+      const std::string& prompt,
+      const std::vector<torch::Tensor>& raw_images,
+      int64_t batch_index,
+      int64_t batch_size) {
+    CHECK(!text_encoder_.is_empty()) << "Qwen3-VL text encoder is not loaded";
+    CHECK(tokenizer_ != nullptr) << "Qwen3-VL tokenizer is not loaded";
+    CHECK(multimodal_processor_ != nullptr)
+        << "Qwen3-VL multimodal processor is not loaded";
+    CHECK(mposition_generator_ != nullptr)
+        << "Qwen3-VL position generator is not loaded";
+
+    std::vector<torch::Tensor> images =
+        prepare_vl_images(raw_images, batch_index, batch_size);
+    std::vector<MMInputItem> input_items;
+    input_items.reserve(images.size());
+    for (torch::Tensor& image : images) {
+      MMInputItem item;
+      item.type = MMType::IMAGE;
+      item.decode_image = std::move(image);
+      input_items.emplace_back(std::move(item));
+    }
+    MMInput multimodal_input;
+    multimodal_input.insert(input_items);
+    MMData mm_data;
+    CHECK(multimodal_processor_->process_multimodal(multimodal_input, mm_data))
+        << "JoyImageEditPlus Qwen3-VL image preprocessing failed";
+
+    std::string chat_prompt = build_qwen3_vl_prompt(prompt, images.size());
+    std::vector<int32_t> token_ids;
+    CHECK(
+        multimodal_processor_->process_prompt(chat_prompt, mm_data, token_ids))
+        << "JoyImageEditPlus Qwen3-VL prompt processing failed";
+    UpdateMMItemScheduleStateVisitor schedule_visitor(
+        /*computed_token_num=*/0,
+        static_cast<int32_t>(token_ids.size()),
+        /*seq_idx=*/0);
+    CHECK(mm_data.foreach (schedule_visitor))
+        << "JoyImageEditPlus failed to schedule Qwen3-VL multimodal inputs";
+    auto [positions, mrope_position_delta] = mposition_generator_->generate(
+        token_ids, mm_data, text_encoder_model_args_);
+    static_cast<void>(mrope_position_delta);
+
+    torch::Tensor tokens =
+        torch::tensor(token_ids, torch::TensorOptions().dtype(torch::kInt32))
+            .to(device_);
+    positions = positions.to(device_);
+    ModelInputParams input_params = build_text_encoder_input(tokens, mm_data);
+    ModelOutput model_output = text_encoder_->forward(
+        tokens, positions, text_encoder_empty_kv_caches_, input_params);
+    CHECK(model_output.residual.defined())
+        << "JoyImageEditPlus requires Qwen3-VL pre-norm hidden states from "
+           "the TORCH backend.";
+    CHECK_GT(model_output.residual.size(0), kPromptEmbeddingStartIndex)
+        << "JoyImageEditPlus Qwen3-VL prompt is shorter than the embedding "
+           "prefix";
+
+    torch::Tensor prompt_embeddings =
+        model_output.residual
+            .slice(
+                /*dim=*/0, kPromptEmbeddingStartIndex)
+            .unsqueeze(0)
+            .to(options_);
+    torch::Tensor prompt_mask =
+        torch::ones({1, prompt_embeddings.size(1)},
+                    torch::TensorOptions().device(device_).dtype(torch::kLong));
+    return {prompt_embeddings, prompt_mask};
+  }
+
+  std::pair<torch::Tensor, torch::Tensor> encode_prompts(
+      const std::vector<std::string>& prompts,
+      const std::vector<torch::Tensor>& raw_images,
+      int64_t batch_size) {
+    CHECK_EQ(static_cast<int64_t>(prompts.size()), batch_size)
+        << "JoyImageEditPlus prompt batch size mismatch";
+    std::vector<torch::Tensor> embeddings;
+    std::vector<torch::Tensor> masks;
+    embeddings.reserve(static_cast<size_t>(batch_size));
+    masks.reserve(static_cast<size_t>(batch_size));
+    int64_t max_sequence_length = 0;
+    for (int64_t batch_index = 0; batch_index < batch_size; ++batch_index) {
+      auto [embedding, mask] =
+          encode_single_prompt(prompts[static_cast<size_t>(batch_index)],
+                               raw_images,
+                               batch_index,
+                               batch_size);
+      max_sequence_length = std::max(max_sequence_length, embedding.size(1));
+      embeddings.emplace_back(std::move(embedding));
+      masks.emplace_back(std::move(mask));
+    }
+    for (int64_t batch_index = 0; batch_index < batch_size; ++batch_index) {
+      embeddings[static_cast<size_t>(batch_index)] = pad_seq(
+          embeddings[static_cast<size_t>(batch_index)], max_sequence_length);
+      masks[static_cast<size_t>(batch_index)] =
+          pad_seq(masks[static_cast<size_t>(batch_index)], max_sequence_length);
+    }
+    return {torch::cat(embeddings, /*dim=*/0), torch::cat(masks, /*dim=*/0)};
+  }
+
+  torch::Tensor transformer_forward(
+      const torch::Tensor& hidden_states,
+      const torch::Tensor& timestep,
+      const torch::Tensor& encoder_hidden_states,
+      const TransformerForwardContext& forward_context,
+      bool use_cfg,
+      int64_t step_index) {
+    return transformer_->forward(hidden_states,
+                                 timestep,
+                                 encoder_hidden_states,
+                                 forward_context.rope_cos,
+                                 forward_context.rope_sin,
+                                 forward_context.attention_mask,
+                                 use_cfg,
+                                 step_index);
+  }
+
+  TransformerForwardContext prepare_transformer_context(
+      int64_t batch_size,
+      int64_t image_sequence_length,
+      const torch::Device& device,
+      const torch::Tensor& encoder_hidden_states_mask,
+      const JoyImageShapeList& shape_list) {
+    CHECK_EQ(static_cast<int64_t>(shape_list.size()), batch_size)
+        << "shape_list batch size must match transformer batch size";
+
+    std::vector<torch::Tensor> cos_list;
+    std::vector<torch::Tensor> sin_list;
+    cos_list.reserve(batch_size);
+    sin_list.reserve(batch_size);
+    for (int64_t batch_index = 0; batch_index < batch_size; ++batch_index) {
+      std::vector<torch::Tensor> cos_parts;
+      std::vector<torch::Tensor> sin_parts;
+      int64_t temporal_offset = 0;
+      for (const auto& shape : shape_list[batch_index]) {
+        std::array<int64_t, 3> start = {temporal_offset, 0, 0};
+        std::array<int64_t, 3> stop = {
+            temporal_offset + shape[0], shape[1], shape[2]};
+        auto rope = rope_for_range(start, stop, device);
+        cos_parts.emplace_back(rope.first);
+        sin_parts.emplace_back(rope.second);
+        temporal_offset += shape[0];
+      }
+
+      torch::Tensor cos = torch::cat(cos_parts, /*dim=*/0);
+      torch::Tensor sin = torch::cat(sin_parts, /*dim=*/0);
+      const int64_t actual_sequence_length = cos.size(0);
+      if (actual_sequence_length < image_sequence_length) {
+        cos = torch::constant_pad_nd(
+            cos,
+            {0, 0, 0, image_sequence_length - actual_sequence_length},
+            /*value=*/1.0);
+        sin = torch::constant_pad_nd(
+            sin,
+            {0, 0, 0, image_sequence_length - actual_sequence_length},
+            /*value=*/0.0);
+      }
+      cos_list.emplace_back(cos);
+      sin_list.emplace_back(sin);
+    }
+
+    torch::Tensor rope_cos = torch::stack(cos_list, /*dim=*/0);
+    torch::Tensor rope_sin = torch::stack(sin_list, /*dim=*/0);
+    torch::Tensor attention_mask;
+
+    if (encoder_hidden_states_mask.defined()) {
+      torch::Tensor image_mask = torch::zeros(
+          {batch_size, image_sequence_length},
+          torch::TensorOptions().device(device).dtype(torch::kBool));
+      for (int64_t batch_index = 0; batch_index < batch_size; ++batch_index) {
+        int64_t actual_sequence_length = 0;
+        for (const auto& shape : shape_list[batch_index]) {
+          actual_sequence_length += shape[0] * shape[1] * shape[2];
+        }
+        image_mask.index_put_(
+            {batch_index, torch::indexing::Slice(0, actual_sequence_length)},
+            true);
+      }
+      torch::Tensor full_mask =
+          torch::cat({image_mask, encoder_hidden_states_mask.to(torch::kBool)},
+                     /*dim=*/1);
+#if defined(USE_NPU)
+      const int64_t sequence_length = full_mask.size(1);
+      attention_mask = full_mask.logical_not()
+                           .unsqueeze(1)
+                           .unsqueeze(1)
+                           .expand({batch_size,
+                                    /*num_heads=*/1,
+                                    sequence_length,
+                                    sequence_length})
+                           .contiguous();
+#else
+      attention_mask = full_mask.unsqueeze(1).unsqueeze(1);
+#endif
+    }
+
+    return {rope_cos, rope_sin, attention_mask};
+  }
+
+  std::pair<torch::Tensor, torch::Tensor> rope_for_range(
+      const std::array<int64_t, 3>& start,
+      const std::array<int64_t, 3>& stop,
+      const torch::Device& device) {
+    std::vector<int64_t> dimensions = rope_dim_list_;
+    if (dimensions.empty()) {
+      const int64_t dimension = head_dim_ / 3;
+      dimensions = {dimension, dimension, dimension};
+    }
+
+    torch::TensorOptions float_options =
+        torch::TensorOptions().dtype(torch::kFloat32).device(device);
+    std::vector<torch::Tensor> grids;
+    grids.reserve(3);
+    for (int64_t dimension_index = 0; dimension_index < 3; ++dimension_index) {
+      grids.emplace_back(torch::arange(
+          start[dimension_index], stop[dimension_index], float_options));
+    }
+    auto mesh = torch::meshgrid({grids[0], grids[1], grids[2]}, "ij");
+
+    std::vector<torch::Tensor> cos_parts;
+    std::vector<torch::Tensor> sin_parts;
+    cos_parts.reserve(3);
+    sin_parts.reserve(3);
+    for (int64_t dimension_index = 0; dimension_index < 3; ++dimension_index) {
+      torch::Tensor position = mesh[dimension_index].reshape({-1});
+      const int64_t dimension = dimensions[dimension_index];
+      torch::Tensor index =
+          torch::arange(0, dimension, 2, float_options)
+              .slice(/*dim=*/0, /*start=*/0, /*end=*/dimension / 2);
+      torch::Tensor frequencies =
+          1.0 / torch::pow(static_cast<double>(theta_), index / dimension);
+      torch::Tensor angles = torch::outer(position, frequencies);
+      cos_parts.emplace_back(angles.cos().repeat_interleave(2, /*dim=*/1));
+      sin_parts.emplace_back(angles.sin().repeat_interleave(2, /*dim=*/1));
+    }
+    return {torch::cat(cos_parts, /*dim=*/1), torch::cat(sin_parts, /*dim=*/1)};
+  }
+
+  // Nearest 1024-base aspect bucket (h, w).
+  std::pair<int64_t, int64_t> joyimage_bucket(int64_t h, int64_t w) {
+    static const std::vector<std::pair<int64_t, int64_t>> kBuckets = {
+        {512, 1792},  {512, 1856}, {512, 1920}, {512, 1984}, {512, 2048},
+        {576, 1600},  {576, 1664}, {576, 1728}, {576, 1792}, {640, 1472},
+        {640, 1536},  {640, 1600}, {704, 1344}, {704, 1408}, {704, 1472},
+        {768, 1216},  {768, 1280}, {768, 1344}, {832, 1152}, {832, 1216},
+        {896, 1088},  {896, 1152}, {960, 1024}, {960, 1088}, {1024, 960},
+        {1024, 1024}, {1088, 896}, {1088, 960}, {1152, 832}, {1152, 896},
+        {1216, 768},  {1216, 832}, {1280, 768}, {1344, 704}, {1344, 768},
+        {1408, 704},  {1472, 640}, {1472, 704}, {1536, 640}, {1600, 576},
+        {1600, 640},  {1664, 576}, {1728, 576}, {1792, 512}, {1792, 576},
+        {1856, 512},  {1920, 512}, {1984, 512}, {2048, 512}};
+    double target = static_cast<double>(h) / static_cast<double>(w);
+    int64_t best_h = 1024, best_w = 1024;
+    double best_diff = std::numeric_limits<double>::max();
+    for (const auto& hw : kBuckets) {
+      double diff =
+          std::abs(static_cast<double>(hw.first) / hw.second - target);
+      if (diff < best_diff) {
+        best_diff = diff;
+        best_h = hw.first;
+        best_w = hw.second;
+      }
+    }
+    return {best_h, best_w};
+  }
+
+  torch::Tensor pad_seq(const torch::Tensor& x, int64_t target_len) {
+    int64_t cur = x.size(1);
+    if (cur >= target_len) return x.slice(1, cur - target_len, cur);
+    int64_t pad = target_len - cur;
+    std::vector<int64_t> shape = x.sizes().vec();
+    shape[1] = pad;
+    auto zeros = torch::zeros(shape, x.options());
+    return torch::cat({x, zeros}, 1);
+  }
+
+  DiTModelContext context_;
+  const ParallelArgs parallel_args_;
+  const ModelArgs& vae_model_args_;
+  torch::Device device_ = torch::kCPU;
+  torch::ScalarType dtype_;
+  torch::TensorOptions options_;
+
+  AutoencoderKLWan vae_{nullptr};
+  joyimage::JoyImageEditPlusTransformer3DModel transformer_{nullptr};
+  FlowMatchEulerDiscreteScheduler scheduler_{nullptr};
+  xllm::VAEImageProcessor vae_image_processor_{nullptr};
+  JoyImageEditTextEncoder text_encoder_{nullptr};
+  std::shared_ptr<Tokenizer> tokenizer_;
+  std::unique_ptr<MultimodalProcessorBase> multimodal_processor_;
+  std::unique_ptr<MPositionGenerator> mposition_generator_;
+  ModelArgs text_encoder_model_args_;
+  std::vector<KVCache> text_encoder_empty_kv_caches_;
+
+  int64_t in_channels_;
+  int64_t num_layers_;
+  int64_t head_dim_;
+  int64_t theta_;
+  int64_t patch_t_, patch_h_, patch_w_;
+  int64_t latent_channels_;
+  int64_t vae_scale_factor_spatial_;
+  std::vector<int64_t> rope_dim_list_;
+  std::vector<double> latents_mean_;
+  std::vector<double> latents_std_;
+};
+
+TORCH_MODULE(JoyImageEditPlusPipeline);
+
+REGISTER_DIT_MODEL(JoyImageEditPlusPipeline, JoyImageEditPlusPipeline);
+}  // namespace xllm
diff --git a/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h
new file mode 100644
index 0000000000..86388b2b18
--- /dev/null
+++ b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h
@@ -0,0 +1,985 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+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 <glog/logging.h>
+#include <torch/torch.h>
+#if defined(USE_NPU)
+#include <torch_npu/csrc/aten/CustomFunctions.h>
+#endif
+
+#include <cmath>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "core/framework/dit_cache/dit_cache.h"
+#include "core/framework/dit_model_loader.h"
+#include "core/framework/model_context.h"
+#include "core/framework/parallel_state/parallel_state.h"
+#include "core/framework/state_dict/state_dict.h"
+#include "core/framework/state_dict/utils.h"
+#include "core/layers/common/ada_layer_norm.h"
+#include "core/layers/common/add_matmul.h"
+#include "core/layers/common/rms_norm.h"
+#include "models/dit/transformers/transformer_qwen_image.h"
+#include "models/dit/utils/sequence_parallel_mixin.h"
+#include "models/model_registry.h"
+
+namespace xllm {
+namespace joyimage {
+
+// ---------------------------------------------------------------------------
+// Rotary position embedding (batched [B, S, D] cos/sin)
+// ---------------------------------------------------------------------------
+// Mirrors diffusers `_apply_rotary_emb_batched`:
+//   x_out = x * cos + rotate_half(x) * sin
+// where rotate_half interleaves pairs: [-x2, x1] over the last dim.
+// x:        [B, S, H, D]
+// cos/sin:  [B, S, D] or [1, S, D] (broadcast over heads)
+inline torch::Tensor apply_rotary_emb_batched(const torch::Tensor& x,
+                                              const torch::Tensor& cos,
+                                              const torch::Tensor& sin) {
+  auto cos_b = (cos.dim() == 3 ? cos.unsqueeze(2) : cos)
+                   .to(torch::kFloat32);  // [B or 1, S, 1, D]
+  auto sin_b = (sin.dim() == 3 ? sin.unsqueeze(2) : sin)
+                   .to(torch::kFloat32);  // [B or 1, S, 1, D]
+
+#if defined(USE_NPU)
+  auto x_float = x.to(torch::kFloat32);
+  if (cos_b.size(0) == 1) {
+    auto x_out = at_npu::native::custom_ops::npu_rotary_mul(
+        x_float, cos_b, sin_b, "interleave");
+    return x_out.to(x.dtype());
+  }
+
+  CHECK_EQ(cos_b.size(0), x.size(0))
+      << "RoPE batch mismatch: x batch=" << x.size(0)
+      << ", cos batch=" << cos_b.size(0);
+  std::vector<torch::Tensor> outputs;
+  outputs.reserve(x.size(0));
+  for (int64_t batch_index = 0; batch_index < x.size(0); ++batch_index) {
+    auto x_slice = x_float.slice(0, batch_index, batch_index + 1);
+    auto cos_slice = cos_b.slice(0, batch_index, batch_index + 1);
+    auto sin_slice = sin_b.slice(0, batch_index, batch_index + 1);
+    outputs.push_back(at_npu::native::custom_ops::npu_rotary_mul(
+        x_slice, cos_slice, sin_slice, "interleave"));
+  }
+  return torch::cat(outputs, 0).to(x.dtype());
+#else
+  auto x_float = x.to(torch::kFloat32);
+  // rotate_half: view last dim as pairs, produce [-x_imag, x_real]
+  auto x_pairs = x_float.reshape(
+      {x_float.size(0), x_float.size(1), x_float.size(2), -1, 2});
+  auto x_real = x_pairs.select(-1, 0);
+  auto x_imag = x_pairs.select(-1, 1);
+  auto rotated =
+      torch::stack({-x_imag, x_real}, -1).flatten(3);  // [B, S, H, D]
+
+  auto out = x_float * cos_b + rotated * sin_b;
+  return out.to(x.dtype());
+#endif
+}
+
+// Joint attention over [B, S, H, D] with an optional bool mask. NPU uses a
+// materialized block mask (True == masked out); the fallback uses a broadcast
+// keep mask (True == attend).
+inline torch::Tensor joyimage_joint_attention(
+    const torch::Tensor& query,  // [B, S, H, D]
+    const torch::Tensor& key,    // [B, S, H, D]
+    const torch::Tensor& value,  // [B, S, H, D]
+    int64_t num_heads,
+    const torch::Tensor& attn_mask = torch::Tensor()) {
+#if defined(USE_NPU)
+  auto results = at_npu::native::custom_ops::npu_fusion_attention(
+      query,
+      key,
+      value,
+      num_heads,
+      /*input_layout=*/"BSND",
+      /*pse=*/torch::nullopt,
+      /*padding_mask=*/torch::nullopt,
+      /*atten_mask=*/attn_mask,
+      /*scale=*/std::pow(static_cast<double>(query.size(3)), -0.5),
+      /*keep_prob=*/1.0,
+      /*pre_tockens=*/65535,
+      /*next_tockens=*/65535);
+  return std::get<0>(results);  // [B, S, H, D]
+#else
+  auto q = query.transpose(1, 2);  // [B, H, S, D]
+  auto k = key.transpose(1, 2);
+  auto v = value.transpose(1, 2);
+  auto out_dtype = q.dtype();
+  c10::optional<torch::Tensor> mask = c10::nullopt;
+  if (attn_mask.defined()) {
+    mask = attn_mask;  // [B, 1, 1, S] bool broadcasts over heads/query
+  }
+  auto output = torch::scaled_dot_product_attention(q,
+                                                    k,
+                                                    v,
+                                                    mask,
+                                                    /*dropout_p=*/0.0,
+                                                    /*is_causal=*/false);
+  if (output.dtype() != out_dtype) output = output.to(out_dtype);
+  return output.transpose(1, 2).contiguous();  // [B, S, H, D]
+#endif
+}
+
+// ---------------------------------------------------------------------------
+// Wan-style learnable modulation table
+// ---------------------------------------------------------------------------
+// modulate_table: [1, factor, hidden] learnable parameter.
+// forward(x[B, hidden]) -> factor tensors of [B, hidden], from
+// (modulate_table + x).chunk(factor).
+class ModulateImpl final : public torch::nn::Module {
+ public:
+  ModulateImpl(int64_t hidden_size, int64_t factor) : factor_(factor) {
+    modulate_table_ = register_parameter(
+        "modulate_table", torch::zeros({1, factor, hidden_size}));
+  }
+
+  std::vector<torch::Tensor> forward(const torch::Tensor& x) {
+    // x: [B, hidden] -> [B, 1, hidden]
+    auto xin = (x.dim() != 3) ? x.unsqueeze(1) : x;
+    auto summed = modulate_table_.to(xin.dtype()) + xin;  // [B, factor, hidden]
+    auto chunks = summed.chunk(factor_, /*dim=*/1);
+    std::vector<torch::Tensor> out;
+    out.reserve(factor_);
+    for (auto& c : chunks) {
+      out.emplace_back(c.squeeze(1));  // [B, hidden]
+    }
+    return out;
+  }
+
+  void load_state_dict(const StateDict& state_dict) {
+    weight::load_weight(state_dict, "modulate_table", modulate_table_, loaded_);
+  }
+  void verify_loaded_weights(const std::string& prefix) const {
+    CHECK(loaded_) << "weight not loaded for " << prefix + "modulate_table";
+  }
+
+ private:
+  int64_t factor_;
+  torch::Tensor modulate_table_;
+  bool loaded_{false};
+};
+TORCH_MODULE(Modulate);
+
+// PixArt-style text projection: Linear -> gelu_tanh -> Linear.
+class TextProjectionImpl final : public torch::nn::Module {
+ public:
+  TextProjectionImpl(const ModelContext& context,
+                     int64_t in_features,
+                     int64_t hidden_size)
+      : options_(context.get_tensor_options()) {
+    linear_1_ = register_module("linear_1",
+                                layer::AddMatmulWeightTransposed(
+                                    in_features, hidden_size, true, options_));
+    linear_2_ = register_module("linear_2",
+                                layer::AddMatmulWeightTransposed(
+                                    hidden_size, hidden_size, true, options_));
+  }
+
+  torch::Tensor forward(const torch::Tensor& caption) {
+    auto x = linear_1_->forward(caption);
+    x = torch::gelu(x, "tanh");
+    x = linear_2_->forward(x);
+    return x;
+  }
+
+  void load_state_dict(const StateDict& state_dict) {
+    linear_1_->load_state_dict(state_dict.get_dict_with_prefix("linear_1."));
+    linear_2_->load_state_dict(state_dict.get_dict_with_prefix("linear_2."));
+  }
+  void verify_loaded_weights(const std::string& prefix) {
+    linear_1_->verify_loaded_weights(prefix + "linear_1.");
+    linear_2_->verify_loaded_weights(prefix + "linear_2.");
+  }
+
+ private:
+  layer::AddMatmulWeightTransposed linear_1_{nullptr};
+  layer::AddMatmulWeightTransposed linear_2_{nullptr};
+  torch::TensorOptions options_;
+};
+TORCH_MODULE(TextProjection);
+
+// condition_embedder: timesteps -> time_embedder -> (temb, time_proj);
+// text_embedder(encoder_hidden_states).
+class TimeTextEmbeddingImpl final : public torch::nn::Module {
+ public:
+  TimeTextEmbeddingImpl(const ModelContext& context,
+                        int64_t hidden_size,
+                        int64_t time_freq_dim,
+                        int64_t time_proj_dim,
+                        int64_t text_embed_dim)
+      : options_(context.get_tensor_options()) {
+    timesteps_proj_ =
+        register_module("timesteps_proj",
+                        qwenimage::Timesteps(context,
+                                             time_freq_dim,
+                                             /*flip_sin_to_cos=*/true,
+                                             /*downscale_freq_shift=*/0.0,
+                                             /*scale=*/1.0));
+    time_embedder_ = register_module(
+        "time_embedder",
+        qwenimage::TimestepEmbedding(context, time_freq_dim, hidden_size));
+    time_embedder_->to(torch::kFloat32);
+    time_proj_ =
+        register_module("time_proj",
+                        layer::AddMatmulWeightTransposed(
+                            hidden_size, time_proj_dim, true, options_));
+    text_embedder_ = register_module(
+        "text_embedder", TextProjection(context, text_embed_dim, hidden_size));
+  }
+
+  // Returns {temb[B, hidden], timestep_proj[B, time_proj_dim],
+  //          text[B, L, hidden]}
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> forward(
+      const torch::Tensor& timestep,
+      const torch::Tensor& encoder_hidden_states) {
+    auto t = timesteps_proj_->forward(timestep);
+    auto temb = time_embedder_->forward(t);  // [B, hidden], computed in FP32
+    temb = temb.to(encoder_hidden_states.dtype());
+    auto act = torch::silu(temb);
+    auto timestep_proj = time_proj_->forward(act);  // [B, time_proj_dim]
+    auto text = text_embedder_->forward(encoder_hidden_states);
+    return std::make_tuple(temb, timestep_proj, text);
+  }
+
+  void load_state_dict(const StateDict& state_dict) {
+    time_embedder_->load_state_dict(
+        state_dict.get_dict_with_prefix("time_embedder."));
+    time_proj_->load_state_dict(state_dict.get_dict_with_prefix("time_proj."));
+    text_embedder_->load_state_dict(
+        state_dict.get_dict_with_prefix("text_embedder."));
+  }
+  void verify_loaded_weights(const std::string& prefix) {
+    time_embedder_->verify_loaded_weights(prefix + "time_embedder.");
+    time_proj_->verify_loaded_weights(prefix + "time_proj.");
+    text_embedder_->verify_loaded_weights(prefix + "text_embedder.");
+  }
+
+  void keep_fp32_modules() { time_embedder_->to(torch::kFloat32); }
+
+ private:
+  qwenimage::Timesteps timesteps_proj_{nullptr};
+  qwenimage::TimestepEmbedding time_embedder_{nullptr};
+  layer::AddMatmulWeightTransposed time_proj_{nullptr};
+  TextProjection text_embedder_{nullptr};
+  torch::TensorOptions options_;
+};
+TORCH_MODULE(TimeTextEmbedding);
+
+using JoyAttentionOutput = std::tuple<torch::Tensor, torch::Tensor>;
+
+class JoyAttentionImpl final : public torch::nn::Module {
+ public:
+  JoyAttentionImpl(const ModelContext& context,
+                   int64_t dim,
+                   int64_t num_heads,
+                   int64_t head_dim,
+                   ProcessGroup* sp_group,
+                   double eps = 1e-6)
+      : heads_(num_heads),
+        head_dim_(head_dim),
+        sp_group_(sp_group),
+        options_(context.get_tensor_options()) {
+    const int64_t inner = num_heads * head_dim;
+
+    img_attn_q_ = register_module(
+        "img_attn_q",
+        layer::AddMatmulWeightTransposed(dim, inner, true, options_));
+    img_attn_k_ = register_module(
+        "img_attn_k",
+        layer::AddMatmulWeightTransposed(dim, inner, true, options_));
+    img_attn_v_ = register_module(
+        "img_attn_v",
+        layer::AddMatmulWeightTransposed(dim, inner, true, options_));
+    img_attn_q_norm_ = register_module("img_attn_q_norm",
+                                       layer::RMSNorm(head_dim, eps, options_));
+    img_attn_k_norm_ = register_module("img_attn_k_norm",
+                                       layer::RMSNorm(head_dim, eps, options_));
+    img_attn_proj_ = register_module(
+        "img_attn_proj",
+        layer::AddMatmulWeightTransposed(inner, dim, true, options_));
+
+    txt_attn_q_ = register_module(
+        "txt_attn_q",
+        layer::AddMatmulWeightTransposed(dim, inner, true, options_));
+    txt_attn_k_ = register_module(
+        "txt_attn_k",
+        layer::AddMatmulWeightTransposed(dim, inner, true, options_));
+    txt_attn_v_ = register_module(
+        "txt_attn_v",
+        layer::AddMatmulWeightTransposed(dim, inner, true, options_));
+    txt_attn_q_norm_ = register_module("txt_attn_q_norm",
+                                       layer::RMSNorm(head_dim, eps, options_));
+    txt_attn_k_norm_ = register_module("txt_attn_k_norm",
+                                       layer::RMSNorm(head_dim, eps, options_));
+    txt_attn_proj_ = register_module(
+        "txt_attn_proj",
+        layer::AddMatmulWeightTransposed(inner, dim, true, options_));
+  }
+
+  JoyAttentionOutput forward(const torch::Tensor& hidden_states,
+                             const torch::Tensor& encoder_hidden_states,
+                             const torch::Tensor& rope_cos,
+                             const torch::Tensor& rope_sin,
+                             const torch::Tensor& attn_mask) {
+    std::vector<int64_t> reshape = {heads_, -1};
+    torch::Tensor img_q =
+        img_attn_q_->forward(hidden_states).unflatten(-1, reshape);
+    auto img_q_work = parallel_state::all_to_all_4D(img_q,
+                                                    /*scatter_idx=*/2,
+                                                    /*gather_idx=*/1,
+                                                    /*async_ops=*/true,
+                                                    sp_group_);
+    torch::Tensor img_k =
+        img_attn_k_->forward(hidden_states).unflatten(-1, reshape);
+    auto img_k_work = parallel_state::all_to_all_4D(img_k,
+                                                    /*scatter_idx=*/2,
+                                                    /*gather_idx=*/1,
+                                                    /*async_ops=*/true,
+                                                    sp_group_);
+    torch::Tensor img_v =
+        img_attn_v_->forward(hidden_states).unflatten(-1, reshape);
+    auto img_v_work = parallel_state::all_to_all_4D(img_v,
+                                                    /*scatter_idx=*/2,
+                                                    /*gather_idx=*/1,
+                                                    /*async_ops=*/true,
+                                                    sp_group_);
+
+    torch::Tensor txt_q =
+        txt_attn_q_->forward(encoder_hidden_states).unflatten(-1, reshape);
+    auto txt_q_work = parallel_state::all_to_all_4D(txt_q,
+                                                    /*scatter_idx=*/2,
+                                                    /*gather_idx=*/1,
+                                                    /*async_ops=*/true,
+                                                    sp_group_);
+    torch::Tensor txt_k =
+        txt_attn_k_->forward(encoder_hidden_states).unflatten(-1, reshape);
+    auto txt_k_work = parallel_state::all_to_all_4D(txt_k,
+                                                    /*scatter_idx=*/2,
+                                                    /*gather_idx=*/1,
+                                                    /*async_ops=*/true,
+                                                    sp_group_);
+    torch::Tensor txt_v =
+        txt_attn_v_->forward(encoder_hidden_states).unflatten(-1, reshape);
+    auto txt_v_work = parallel_state::all_to_all_4D(txt_v,
+                                                    /*scatter_idx=*/2,
+                                                    /*gather_idx=*/1,
+                                                    /*async_ops=*/true,
+                                                    sp_group_);
+
+    img_q = xllm::dit::SequenceParallelMixin::unpad_tensor(
+        img_q_work(), /*tensor_name=*/"hidden_states", /*dim=*/1);
+    img_k = xllm::dit::SequenceParallelMixin::unpad_tensor(
+        img_k_work(), /*tensor_name=*/"hidden_states", /*dim=*/1);
+    txt_q = xllm::dit::SequenceParallelMixin::unpad_tensor(
+        txt_q_work(), /*tensor_name=*/"encoder_hidden_states", /*dim=*/1);
+    txt_k = xllm::dit::SequenceParallelMixin::unpad_tensor(
+        txt_k_work(), /*tensor_name=*/"encoder_hidden_states", /*dim=*/1);
+
+    img_q = std::get<0>(img_attn_q_norm_->forward(img_q));
+    img_k = std::get<0>(img_attn_k_norm_->forward(img_k));
+    txt_q = std::get<0>(txt_attn_q_norm_->forward(txt_q));
+    txt_k = std::get<0>(txt_attn_k_norm_->forward(txt_k));
+
+    if (rope_cos.defined()) {
+      img_q = apply_rotary_emb_batched(img_q, rope_cos, rope_sin);
+      img_k = apply_rotary_emb_batched(img_k, rope_cos, rope_sin);
+    }
+
+    img_v = xllm::dit::SequenceParallelMixin::unpad_tensor(
+        img_v_work(), /*tensor_name=*/"hidden_states", /*dim=*/1);
+    txt_v = xllm::dit::SequenceParallelMixin::unpad_tensor(
+        txt_v_work(), /*tensor_name=*/"encoder_hidden_states", /*dim=*/1);
+
+    torch::Tensor joint_q = torch::cat({img_q, txt_q}, /*dim=*/1);
+    torch::Tensor joint_k = torch::cat({img_k, txt_k}, /*dim=*/1);
+    torch::Tensor joint_v = torch::cat({img_v, txt_v}, /*dim=*/1);
+    const int32_t sp_size = sp_group_ == nullptr ? 1 : sp_group_->world_size();
+    const int64_t local_heads = heads_ / sp_size;
+    torch::Tensor joint = joyimage_joint_attention(
+        joint_q, joint_k, joint_v, local_heads, attn_mask);
+    joint = joint.to(joint_q.dtype());
+
+    const int64_t image_sequence_length = img_q.size(1);
+    torch::Tensor img_out = joint.slice(
+        /*dim=*/1, /*start=*/0, /*end=*/image_sequence_length);
+    torch::Tensor txt_out = joint.slice(/*dim=*/1,
+                                        /*start=*/image_sequence_length,
+                                        /*end=*/joint.size(1));
+
+    img_out = xllm::dit::SequenceParallelMixin::pad_tensor(
+        img_out, /*tensor_name=*/"hidden_states", /*dim=*/1);
+    auto img_out_work = parallel_state::all_to_all_4D(img_out,
+                                                      /*scatter_idx=*/1,
+                                                      /*gather_idx=*/2,
+                                                      /*async_ops=*/true,
+                                                      sp_group_);
+    txt_out = xllm::dit::SequenceParallelMixin::pad_tensor(
+        txt_out, /*tensor_name=*/"encoder_hidden_states", /*dim=*/1);
+    auto txt_out_work = parallel_state::all_to_all_4D(txt_out,
+                                                      /*scatter_idx=*/1,
+                                                      /*gather_idx=*/2,
+                                                      /*async_ops=*/true,
+                                                      sp_group_);
+
+    img_out = img_out_work().flatten(/*start_dim=*/2, /*end_dim=*/3);
+    img_out = img_attn_proj_->forward(img_out);
+    txt_out = txt_out_work().flatten(/*start_dim=*/2, /*end_dim=*/3);
+    txt_out = txt_attn_proj_->forward(txt_out);
+    return std::make_tuple(img_out, txt_out);
+  }
+
+  void load_state_dict(const StateDict& state_dict) {
+    load_qkv_projections(state_dict,
+                         /*prefix=*/"img_attn_qkv",
+                         img_attn_q_,
+                         img_attn_k_,
+                         img_attn_v_);
+    img_attn_q_norm_->load_state_dict(
+        state_dict.get_dict_with_prefix("img_attn_q_norm."));
+    img_attn_k_norm_->load_state_dict(
+        state_dict.get_dict_with_prefix("img_attn_k_norm."));
+    img_attn_proj_->load_state_dict(
+        state_dict.get_dict_with_prefix("img_attn_proj."));
+    load_qkv_projections(state_dict,
+                         /*prefix=*/"txt_attn_qkv",
+                         txt_attn_q_,
+                         txt_attn_k_,
+                         txt_attn_v_);
+    txt_attn_q_norm_->load_state_dict(
+        state_dict.get_dict_with_prefix("txt_attn_q_norm."));
+    txt_attn_k_norm_->load_state_dict(
+        state_dict.get_dict_with_prefix("txt_attn_k_norm."));
+    txt_attn_proj_->load_state_dict(
+        state_dict.get_dict_with_prefix("txt_attn_proj."));
+  }
+  void verify_loaded_weights(const std::string& prefix) {
+    img_attn_q_->verify_loaded_weights(prefix + "img_attn_qkv.q.");
+    img_attn_k_->verify_loaded_weights(prefix + "img_attn_qkv.k.");
+    img_attn_v_->verify_loaded_weights(prefix + "img_attn_qkv.v.");
+    img_attn_q_norm_->verify_loaded_weights(prefix + "img_attn_q_norm.");
+    img_attn_k_norm_->verify_loaded_weights(prefix + "img_attn_k_norm.");
+    img_attn_proj_->verify_loaded_weights(prefix + "img_attn_proj.");
+    txt_attn_q_->verify_loaded_weights(prefix + "txt_attn_qkv.q.");
+    txt_attn_k_->verify_loaded_weights(prefix + "txt_attn_qkv.k.");
+    txt_attn_v_->verify_loaded_weights(prefix + "txt_attn_qkv.v.");
+    txt_attn_q_norm_->verify_loaded_weights(prefix + "txt_attn_q_norm.");
+    txt_attn_k_norm_->verify_loaded_weights(prefix + "txt_attn_k_norm.");
+    txt_attn_proj_->verify_loaded_weights(prefix + "txt_attn_proj.");
+  }
+
+ private:
+  int64_t heads_;
+  int64_t head_dim_;
+  ProcessGroup* sp_group_{nullptr};
+  layer::AddMatmulWeightTransposed img_attn_q_{nullptr};
+  layer::AddMatmulWeightTransposed img_attn_k_{nullptr};
+  layer::AddMatmulWeightTransposed img_attn_v_{nullptr};
+  layer::RMSNorm img_attn_q_norm_{nullptr};
+  layer::RMSNorm img_attn_k_norm_{nullptr};
+  layer::AddMatmulWeightTransposed img_attn_proj_{nullptr};
+  layer::AddMatmulWeightTransposed txt_attn_q_{nullptr};
+  layer::AddMatmulWeightTransposed txt_attn_k_{nullptr};
+  layer::AddMatmulWeightTransposed txt_attn_v_{nullptr};
+  layer::RMSNorm txt_attn_q_norm_{nullptr};
+  layer::RMSNorm txt_attn_k_norm_{nullptr};
+  layer::AddMatmulWeightTransposed txt_attn_proj_{nullptr};
+
+ private:
+  void load_qkv_projections(const StateDict& state_dict,
+                            const std::string& prefix,
+                            layer::AddMatmulWeightTransposed& q_projection,
+                            layer::AddMatmulWeightTransposed& k_projection,
+                            layer::AddMatmulWeightTransposed& v_projection) {
+    torch::Tensor fused_weight = state_dict.get_tensor(prefix + ".weight");
+    torch::Tensor fused_bias = state_dict.get_tensor(prefix + ".bias");
+    if (!fused_weight.defined() && !fused_bias.defined()) {
+      return;
+    }
+
+    const int64_t inner = heads_ * head_dim_;
+    CHECK(fused_weight.defined()) << prefix << ".weight is missing";
+    CHECK(fused_bias.defined()) << prefix << ".bias is missing";
+    CHECK_EQ(fused_weight.dim(), 2) << prefix << ".weight must be 2D";
+    CHECK_EQ(fused_weight.size(0), inner * 3)
+        << prefix << ".weight output size mismatch";
+    CHECK_EQ(fused_bias.dim(), 1) << prefix << ".bias must be 1D";
+    CHECK_EQ(fused_bias.size(0), inner * 3)
+        << prefix << ".bias output size mismatch";
+
+    auto load_projection = [&](layer::AddMatmulWeightTransposed& projection,
+                               int64_t projection_index) {
+      const int64_t start = projection_index * inner;
+      std::unordered_map<std::string, torch::Tensor> projection_tensors;
+      projection_tensors.reserve(2);
+      projection_tensors.emplace(
+          "weight", fused_weight.slice(/*dim=*/0, start, start + inner));
+      projection_tensors.emplace(
+          "bias", fused_bias.slice(/*dim=*/0, start, start + inner));
+      StateDict projection_state_dict(std::move(projection_tensors));
+      projection->load_state_dict(projection_state_dict);
+    };
+
+    load_projection(q_projection, /*projection_index=*/0);
+    load_projection(k_projection, /*projection_index=*/1);
+    load_projection(v_projection, /*projection_index=*/2);
+  }
+
+  torch::TensorOptions options_;
+};
+TORCH_MODULE(JoyAttention);
+
+// Double-stream transformer block.
+class TransformerBlockImpl final : public torch::nn::Module {
+ public:
+  TransformerBlockImpl(const ModelContext& context,
+                       int64_t dim,
+                       int64_t num_heads,
+                       int64_t head_dim,
+                       double mlp_width_ratio,
+                       ProcessGroup* sp_group,
+                       double eps = 1e-6)
+      : eps_(eps) {
+    int64_t mlp_hidden = static_cast<int64_t>(dim * mlp_width_ratio);
+
+    img_mod_ = register_module("img_mod", Modulate(dim, 6));
+#if defined(USE_NPU)
+    const torch::TensorOptions options = context.get_tensor_options();
+    img_norm1_ = register_module(
+        "img_norm1",
+        layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options));
+    img_norm2_ = register_module(
+        "img_norm2",
+        layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options));
+#endif
+    img_mlp_ = register_module(
+        "img_mlp", qwenimage::FeedForward(context, dim, dim, /*mult=*/4));
+    txt_mod_ = register_module("txt_mod", Modulate(dim, 6));
+#if defined(USE_NPU)
+    txt_norm1_ = register_module(
+        "txt_norm1",
+        layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options));
+    txt_norm2_ = register_module(
+        "txt_norm2",
+        layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options));
+#endif
+    txt_mlp_ = register_module(
+        "txt_mlp", qwenimage::FeedForward(context, dim, dim, /*mult=*/4));
+    attn_ = register_module(
+        "attn", JoyAttention(context, dim, num_heads, head_dim, sp_group, eps));
+    (void)mlp_hidden;  // FeedForward uses dim*4 internally (== dim*ratio)
+  }
+
+  // FP32 layernorm without affine.
+  static torch::Tensor fp32_norm(const torch::Tensor& x, double eps) {
+    auto xf = x.to(torch::kFloat32);
+    auto out =
+        torch::layer_norm(xf, {xf.size(-1)}, /*weight=*/{}, /*bias=*/{}, eps);
+    return out.to(x.dtype());
+  }
+
+  std::tuple<torch::Tensor, torch::Tensor> forward(
+      const torch::Tensor& hidden_states,
+      const torch::Tensor& encoder_hidden_states,
+      const torch::Tensor& temb,  // [B, hidden]
+      const torch::Tensor& rope_cos,
+      const torch::Tensor& rope_sin,
+      const torch::Tensor& attn_mask) {
+    auto img_mod = img_mod_->forward(temb);  // 6 x [B, hidden]
+    auto txt_mod = txt_mod_->forward(temb);
+
+    auto& img_s1 = img_mod[0];
+    auto& img_c1 = img_mod[1];
+    auto& img_g1 = img_mod[2];
+    auto& img_s2 = img_mod[3];
+    auto& img_c2 = img_mod[4];
+    auto& img_g2 = img_mod[5];
+    auto& txt_s1 = txt_mod[0];
+    auto& txt_c1 = txt_mod[1];
+    auto& txt_g1 = txt_mod[2];
+    auto& txt_s2 = txt_mod[3];
+    auto& txt_c2 = txt_mod[4];
+    auto& txt_g2 = txt_mod[5];
+
+    auto hs = hidden_states;
+    auto ehs = encoder_hidden_states;
+
+    // --- attention ---
+#if defined(USE_NPU)
+    auto img_modulated = img_norm1_->forward(hs, img_c1, img_s1);
+    auto txt_modulated = txt_norm1_->forward(ehs, txt_c1, txt_s1);
+#else
+    auto img_normed = fp32_norm(hs, eps_);
+    auto txt_normed = fp32_norm(ehs, eps_);
+    auto img_modulated =
+        img_normed * (1 + img_c1.unsqueeze(1)) + img_s1.unsqueeze(1);
+    auto txt_modulated =
+        txt_normed * (1 + txt_c1.unsqueeze(1)) + txt_s1.unsqueeze(1);
+#endif
+
+    JoyAttentionOutput attn_out = attn_->forward(
+        img_modulated, txt_modulated, rope_cos, rope_sin, attn_mask);
+    auto img_attn = std::get<0>(attn_out);
+    auto txt_attn = std::get<1>(attn_out);
+
+    img_attn.mul_(img_g1.unsqueeze(1));
+    img_attn.add_(hs);
+    hs = img_attn;
+    txt_attn.mul_(txt_g1.unsqueeze(1));
+    txt_attn.add_(ehs);
+    ehs = txt_attn;
+
+    // --- FFN ---
+#if defined(USE_NPU)
+    auto img_ffn_in = img_norm2_->forward(hs, img_c2, img_s2);
+    auto txt_ffn_in = txt_norm2_->forward(ehs, txt_c2, txt_s2);
+#else
+    auto img_ffn_normed = fp32_norm(hs, eps_);
+    auto txt_ffn_normed = fp32_norm(ehs, eps_);
+    auto img_ffn_in =
+        img_ffn_normed * (1 + img_c2.unsqueeze(1)) + img_s2.unsqueeze(1);
+    auto txt_ffn_in =
+        txt_ffn_normed * (1 + txt_c2.unsqueeze(1)) + txt_s2.unsqueeze(1);
+#endif
+    auto img_ffn = img_mlp_->forward(img_ffn_in);
+    auto txt_ffn = txt_mlp_->forward(txt_ffn_in);
+    img_ffn.mul_(img_g2.unsqueeze(1));
+    img_ffn.add_(hs);
+    hs = img_ffn;
+    txt_ffn.mul_(txt_g2.unsqueeze(1));
+    txt_ffn.add_(ehs);
+    ehs = txt_ffn;
+
+    return std::make_tuple(hs, ehs);
+  }
+
+  void load_state_dict(const StateDict& state_dict) {
+    img_mod_->load_state_dict(state_dict.get_dict_with_prefix("img_mod."));
+    txt_mod_->load_state_dict(state_dict.get_dict_with_prefix("txt_mod."));
+    img_mlp_->load_state_dict(state_dict.get_dict_with_prefix("img_mlp."));
+    txt_mlp_->load_state_dict(state_dict.get_dict_with_prefix("txt_mlp."));
+    attn_->load_state_dict(state_dict.get_dict_with_prefix("attn."));
+  }
+  void verify_loaded_weights(const std::string& prefix) {
+    img_mod_->verify_loaded_weights(prefix + "img_mod.");
+    txt_mod_->verify_loaded_weights(prefix + "txt_mod.");
+    img_mlp_->verify_loaded_weights(prefix + "img_mlp.");
+    txt_mlp_->verify_loaded_weights(prefix + "txt_mlp.");
+    attn_->verify_loaded_weights(prefix + "attn.");
+  }
+
+ private:
+  double eps_;
+  Modulate img_mod_{nullptr};
+  Modulate txt_mod_{nullptr};
+#if defined(USE_NPU)
+  layer::AdaLayerNorm img_norm1_{nullptr};
+  layer::AdaLayerNorm img_norm2_{nullptr};
+  layer::AdaLayerNorm txt_norm1_{nullptr};
+  layer::AdaLayerNorm txt_norm2_{nullptr};
+#endif
+  qwenimage::FeedForward img_mlp_{nullptr};
+  qwenimage::FeedForward txt_mlp_{nullptr};
+  JoyAttention attn_{nullptr};
+};
+TORCH_MODULE(TransformerBlock);
+
+class JoyImageEditPlusTransformer3DModelImpl final
+    : public torch::nn::Module,
+      public xllm::dit::SequenceParallelMixin {
+ public:
+  JoyImageEditPlusTransformer3DModelImpl(const ModelContext& context,
+                                         const ParallelArgs& parallel_args)
+      : xllm::dit::SequenceParallelMixin(
+            /*process_group=*/parallel_args.dit_sp_group_,
+            /*input_sequence_dims=*/
+            {{"hidden_states", 1}, {"encoder_hidden_states", 1}},
+            /*output_sequence_dims=*/{{"hidden_states", 1}}),
+        options_(context.get_tensor_options()) {
+    auto model_args = context.get_model_args();
+    hidden_size_ = model_args.hidden_size();
+    num_heads_ = model_args.num_attention_heads();
+    int64_t num_layers = model_args.num_layers();
+    in_channels_ = model_args.in_channels();
+    int64_t out_channels = model_args.out_channels();
+    out_channels_ = (out_channels > 0) ? out_channels : in_channels_;
+    patch_size_ = model_args.wan_patch_size();  // {pt, ph, pw}
+    double mlp_width_ratio = model_args.mlp_width_ratio();
+    int64_t text_dim = model_args.text_dim();
+
+    head_dim_ = hidden_size_ / num_heads_;
+    CHECK_EQ(hidden_size_ % num_heads_, 0)
+        << "hidden_size must be divisible by num_attention_heads";
+    const int32_t sp_size = parallel_args.dit_sp_group_ == nullptr
+                                ? 1
+                                : parallel_args.dit_sp_group_->world_size();
+    CHECK_EQ(num_heads_ % sp_size, 0)
+        << "JoyImageEditPlus attention heads must be divisible by sp_size";
+
+    // Conv3d patchifier: kernel == stride == patch_size.
+    img_in_ = register_module(
+        "img_in",
+        torch::nn::Conv3d(
+            torch::nn::Conv3dOptions(
+                in_channels_,
+                hidden_size_,
+                {patch_size_[0], patch_size_[1], patch_size_[2]})
+                .stride({patch_size_[0], patch_size_[1], patch_size_[2]})
+                .bias(true)));
+
+    condition_embedder_ =
+        register_module("condition_embedder",
+                        TimeTextEmbedding(context,
+                                          hidden_size_,
+                                          /*time_freq_dim=*/256,
+                                          /*time_proj_dim=*/hidden_size_ * 6,
+                                          text_dim));
+
+    double_blocks_ = register_module("double_blocks", torch::nn::ModuleList());
+    for (int64_t i = 0; i < num_layers; ++i) {
+      auto block = TransformerBlock(context,
+                                    hidden_size_,
+                                    num_heads_,
+                                    head_dim_,
+                                    mlp_width_ratio,
+                                    parallel_args.dit_sp_group_);
+      double_blocks_->push_back(block);
+      block_layers_.push_back(block);
+    }
+
+    int64_t patch_prod = patch_size_[0] * patch_size_[1] * patch_size_[2];
+    proj_out_ = register_module(
+        "proj_out",
+        layer::AddMatmulWeightTransposed(
+            hidden_size_, out_channels_ * patch_prod, true, options_));
+  }
+
+  // hidden_states: [B, N, C, pt, ph, pw]
+  // timestep:      [B]
+  // encoder_hidden_states: [B, L, text_dim]
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const torch::Tensor& timestep,
+                        const torch::Tensor& encoder_hidden_states,
+                        const torch::Tensor& rope_cos,
+                        const torch::Tensor& rope_sin,
+                        const torch::Tensor& attention_mask,
+                        bool use_cfg = false,
+                        int64_t step_index = 1) {
+    xllm::dit::SequenceParallelTensorMap model_outputs =
+        sequence_parallel_forward(
+            {{"hidden_states", hidden_states},
+             {"encoder_hidden_states", encoder_hidden_states}},
+            [this,
+             &timestep,
+             &rope_cos,
+             &rope_sin,
+             &attention_mask,
+             use_cfg,
+             step_index](
+                const xllm::dit::SequenceParallelTensorMap& model_inputs) {
+              torch::Tensor output =
+                  forward_impl(model_inputs.at("hidden_states"),
+                               timestep,
+                               model_inputs.at("encoder_hidden_states"),
+                               rope_cos,
+                               rope_sin,
+                               attention_mask,
+                               use_cfg,
+                               step_index);
+              return xllm::dit::SequenceParallelTensorMap{
+                  {"hidden_states", output}};
+            });
+    return model_outputs.at("hidden_states");
+  }
+
+ private:
+  torch::Tensor forward_impl(const torch::Tensor& local_hidden_states,
+                             const torch::Tensor& timestep,
+                             const torch::Tensor& local_encoder_hidden_states,
+                             const torch::Tensor& rope_cos,
+                             const torch::Tensor& rope_sin,
+                             const torch::Tensor& attention_mask,
+                             bool use_cfg,
+                             int64_t step_index) {
+    int64_t batch_size = local_hidden_states.size(0);
+    int64_t sequence_length = local_hidden_states.size(1);
+    int64_t channels = local_hidden_states.size(2);
+    int64_t pt = local_hidden_states.size(3);
+    int64_t ph = local_hidden_states.size(4);
+    int64_t pw = local_hidden_states.size(5);
+
+    // 1. Condition embeddings.
+    auto cond =
+        condition_embedder_->forward(timestep, local_encoder_hidden_states);
+    auto vec = std::get<1>(cond);  // [B, hidden*6] -> but we use per-block temb
+    auto txt = std::get<2>(cond);  // [B, L, hidden]
+    // vec is the projected timestep [B, 6*hidden]; blocks re-add modulate_table
+    // to a [B, hidden] signal. Diffusers passes vec.unflatten(1,(6,-1)) then
+    // the Modulate table adds to it. To match, feed each block the [B, 6,
+    // hidden] signal; our Modulate expects [B, hidden] or [B, 1, hidden]. We
+    // therefore pass timestep_proj reshaped to [B, 6, hidden] and let Modulate
+    // add table.
+    auto temb6 = vec.unflatten(1, std::vector<int64_t>{6, -1});  // [B,6,hidden]
+
+    // 2. Patchify via Conv3d.
+    auto x = local_hidden_states.reshape(
+        {batch_size * sequence_length, channels, pt, ph, pw});
+    x = img_in_->forward(x);  // [B*N, hidden, 1, 1, 1]
+    auto img = x.reshape({batch_size, sequence_length, hidden_size_});
+
+    // 3. Blocks with optional DiT cache.
+    torch::Tensor original_img = img;
+    torch::Tensor original_txt = txt;
+    TensorMap step_before_map = {
+        {"hidden_states", img},
+        {"encoder_hidden_states", txt},
+        {"original_hidden_states", original_img},
+        {"original_encoder_hidden_states", original_txt}};
+    CacheStepIn step_before(step_index, step_before_map);
+    const bool use_step_cache =
+        DiTCache::get_instance().on_before_step(step_before, use_cfg);
+
+    if (!use_step_cache) {
+      for (int64_t block_index = 0;
+           block_index < static_cast<int64_t>(block_layers_.size());
+           ++block_index) {
+        CacheBlockIn block_before(block_index);
+        const bool use_block_cache =
+            DiTCache::get_instance().on_before_block(block_before, use_cfg);
+        if (!use_block_cache) {
+          std::tie(img, txt) = block_layers_[block_index]->forward(
+              img, txt, temb6, rope_cos, rope_sin, attention_mask);
+        }
+
+        TensorMap block_after_map = {
+            {"hidden_states", img},
+            {"encoder_hidden_states", txt},
+            {"original_hidden_states", original_img},
+            {"original_encoder_hidden_states", original_txt}};
+        CacheBlockIn block_after(block_index, block_after_map);
+        CacheBlockOut block_output =
+            DiTCache::get_instance().on_after_block(block_after, use_cfg);
+        img = block_output.tensors.at("hidden_states");
+        txt = block_output.tensors.at("encoder_hidden_states");
+      }
+    }
+
+    TensorMap step_after_map = {
+        {"hidden_states", img},
+        {"encoder_hidden_states", txt},
+        {"original_hidden_states", original_img},
+        {"original_encoder_hidden_states", original_txt}};
+    CacheStepIn step_after(step_index, step_after_map);
+    CacheStepOut step_output =
+        DiTCache::get_instance().on_after_step(step_after, use_cfg);
+    img = step_output.tensors.at("hidden_states");
+
+    // 4. Output projection + reshape to [B, N, C_out, pt, ph, pw].
+    img = proj_out_->forward(fp32_norm_out(img));
+    img = img.reshape({batch_size, sequence_length, pt, ph, pw, out_channels_})
+              .permute({0, 1, 5, 2, 3, 4});
+    return img;
+  }
+
+ public:
+  static torch::Tensor fp32_norm_out(const torch::Tensor& x) {
+    auto xf = x.to(torch::kFloat32);
+    auto out =
+        torch::layer_norm(xf, {xf.size(-1)}, /*weight=*/{}, /*bias=*/{}, 1e-6);
+    return out.to(x.dtype());
+  }
+
+  void load_model(std::unique_ptr<DiTFolderLoader> loader) {
+    for (const auto& state_dict : loader->get_state_dicts()) {
+      // Conv3d img_in: load raw weight/bias.
+      auto w = state_dict->get_tensor("img_in.weight");
+      auto b = state_dict->get_tensor("img_in.bias");
+      if (w.defined()) {
+        img_in_->weight.data().copy_(w.to(options_));
+        img_in_weight_loaded_ = true;
+      }
+      if (b.defined()) {
+        img_in_->bias.data().copy_(b.to(options_));
+        img_in_bias_loaded_ = true;
+      }
+      condition_embedder_->load_state_dict(
+          state_dict->get_dict_with_prefix("condition_embedder."));
+      proj_out_->load_state_dict(state_dict->get_dict_with_prefix("proj_out."));
+      for (size_t i = 0; i < block_layers_.size(); ++i) {
+        auto prefix = "double_blocks." + std::to_string(i) + ".";
+        block_layers_[i]->load_state_dict(
+            state_dict->get_dict_with_prefix(prefix));
+      }
+    }
+    verify_loaded_weights();
+    LOG(INFO) << "JoyImageEditPlus transformer loaded successfully.";
+  }
+
+  void verify_loaded_weights() {
+    CHECK(img_in_weight_loaded_) << "img_in.weight not loaded";
+    CHECK(img_in_bias_loaded_) << "img_in.bias not loaded";
+    condition_embedder_->verify_loaded_weights("condition_embedder.");
+    proj_out_->verify_loaded_weights("proj_out.");
+    for (size_t i = 0; i < block_layers_.size(); ++i) {
+      block_layers_[i]->verify_loaded_weights("double_blocks." +
+                                              std::to_string(i) + ".");
+    }
+  }
+
+  void keep_fp32_modules() { condition_embedder_->keep_fp32_modules(); }
+
+ private:
+  torch::TensorOptions options_;
+  int64_t hidden_size_;
+  int64_t num_heads_;
+  int64_t head_dim_;
+  int64_t in_channels_;
+  int64_t out_channels_;
+  std::vector<int64_t> patch_size_;
+
+  torch::nn::Conv3d img_in_{nullptr};
+  TimeTextEmbedding condition_embedder_{nullptr};
+  torch::nn::ModuleList double_blocks_{nullptr};
+  std::vector<TransformerBlock> block_layers_;
+  layer::AddMatmulWeightTransposed proj_out_{nullptr};
+  bool img_in_weight_loaded_{false};
+  bool img_in_bias_loaded_{false};
+};
+TORCH_MODULE(JoyImageEditPlusTransformer3DModel);
+
+REGISTER_MODEL_ARGS(JoyImageEditPlusTransformer3DModel, [&] {
+  LOAD_ARG_OR(dtype, "dtype", "bfloat16");
+  LOAD_ARG_OR(hidden_size, "hidden_size", 4096);
+  LOAD_ARG_OR(num_attention_heads, "num_attention_heads", 32);
+  LOAD_ARG_OR(num_layers, "num_layers", 40);
+  LOAD_ARG_OR(in_channels, "in_channels", 16);
+  LOAD_ARG_OR(out_channels, "out_channels", 16);
+  LOAD_ARG_OR(wan_patch_size, "patch_size", (std::vector<int64_t>{1, 2, 2}));
+  LOAD_ARG_OR(mlp_width_ratio, "mlp_width_ratio", 4.0);
+  LOAD_ARG_OR(text_dim, "text_dim", 4096);
+  LOAD_ARG_OR(
+      rope_dim_list, "rope_dim_list", (std::vector<int64_t>{16, 56, 56}));
+  LOAD_ARG_OR(rope_theta_dit, "theta", 10000);
+});
+
+}  // namespace joyimage
+}  // namespace xllm
diff --git a/xllm/models/dit/utils/dit_parallel_mixin.h b/xllm/models/dit/utils/dit_parallel_mixin.h
index cbbb52b00e..6f68ed6891 100644
--- a/xllm/models/dit/utils/dit_parallel_mixin.h
+++ b/xllm/models/dit/utils/dit_parallel_mixin.h
@@ -70,8 +70,5 @@ class CFGParallelMixin {
 // Mixin for VAE parallelism (to be implemented).
 class VaeParallelMixin {};
 
-// Mixin for sequence parallelism (to be implemented).
-class SpParallelMixin {};
-
 }  // namespace dit
 }  // namespace xllm
diff --git a/xllm/models/dit/utils/sequence_parallel_mixin.h b/xllm/models/dit/utils/sequence_parallel_mixin.h
new file mode 100644
index 0000000000..8fef8490c0
--- /dev/null
+++ b/xllm/models/dit/utils/sequence_parallel_mixin.h
@@ -0,0 +1,180 @@
+/* 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 <glog/logging.h>
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "core/framework/parallel_state/parallel_state.h"
+
+namespace xllm::dit {
+
+using SequenceParallelTensorMap =
+    std::unordered_map<std::string, torch::Tensor>;
+using SequenceParallelTensorDims = std::unordered_map<std::string, int64_t>;
+
+class SequenceParallelMixin {
+ public:
+  static torch::Tensor pad_tensor(const torch::Tensor& input,
+                                  const std::string& tensor_name,
+                                  int64_t dim) {
+    if (!input.defined()) {
+      return input;
+    }
+
+    return pad_right(input, padding_length(tensor_name), dim);
+  }
+
+  static torch::Tensor unpad_tensor(const torch::Tensor& input,
+                                    const std::string& tensor_name,
+                                    int64_t dim) {
+    if (!input.defined()) {
+      return input;
+    }
+
+    return unpad_right(input, padding_length(tensor_name), dim);
+  }
+
+ protected:
+  SequenceParallelMixin(ProcessGroup* process_group,
+                        SequenceParallelTensorDims input_sequence_dims,
+                        SequenceParallelTensorDims output_sequence_dims)
+      : process_group_(process_group),
+        input_sequence_dims_(std::move(input_sequence_dims)),
+        output_sequence_dims_(std::move(output_sequence_dims)) {}
+
+  template <typename ForwardFn>
+  SequenceParallelTensorMap sequence_parallel_forward(
+      const SequenceParallelTensorMap& inputs,
+      ForwardFn&& forward_fn) {
+    padding_lengths_.clear();
+    SequenceParallelTensorMap local_inputs = inputs;
+    for (const auto& [tensor_name, sequence_dim] : input_sequence_dims_) {
+      auto tensor_it = local_inputs.find(tensor_name);
+      CHECK(tensor_it != local_inputs.end())
+          << "Missing registered sequence-parallel input: " << tensor_name;
+      tensor_it->second =
+          scatter_sequence(tensor_it->second, tensor_name, sequence_dim);
+    }
+
+    SequenceParallelTensorMap outputs =
+        std::forward<ForwardFn>(forward_fn)(local_inputs);
+    for (const auto& [tensor_name, sequence_dim] : output_sequence_dims_) {
+      auto tensor_it = outputs.find(tensor_name);
+      CHECK(tensor_it != outputs.end())
+          << "Missing registered sequence-parallel output: " << tensor_name;
+      tensor_it->second =
+          gather_sequence(tensor_it->second, tensor_name, sequence_dim);
+    }
+    return outputs;
+  }
+
+ private:
+  int32_t world_size() const {
+    return process_group_ == nullptr ? 1 : process_group_->world_size();
+  }
+
+  bool sequence_parallel_enabled() const { return world_size() > 1; }
+
+  torch::Tensor scatter_sequence(const torch::Tensor& input,
+                                 const std::string& tensor_name,
+                                 int64_t sequence_dim) {
+    if (!input.defined()) {
+      return input;
+    }
+
+    const int64_t sequence_length = input.size(sequence_dim);
+    const int64_t padding_length =
+        (world_size() - sequence_length % world_size()) % world_size();
+    padding_lengths_[tensor_name] = padding_length;
+    if (!sequence_parallel_enabled()) {
+      return input;
+    }
+
+    torch::Tensor padded_input = pad_tensor(input, tensor_name, sequence_dim);
+    return parallel_state::scatter(
+        padded_input, process_group_, static_cast<int32_t>(sequence_dim));
+  }
+
+  torch::Tensor gather_sequence(const torch::Tensor& input,
+                                const std::string& tensor_name,
+                                int64_t sequence_dim) const {
+    if (!sequence_parallel_enabled() || !input.defined()) {
+      return input;
+    }
+
+    torch::Tensor output = parallel_state::gather(
+        input.contiguous(), process_group_, static_cast<int32_t>(sequence_dim));
+    return unpad_tensor(output, tensor_name, sequence_dim);
+  }
+
+  static int64_t normalize_dim(const torch::Tensor& input, int64_t dim) {
+    const int64_t normalized_dim = dim < 0 ? input.dim() + dim : dim;
+    CHECK_GE(normalized_dim, 0) << "Invalid tensor dimension: " << dim;
+    CHECK_LT(normalized_dim, input.dim())
+        << "Invalid tensor dimension: " << dim;
+    return normalized_dim;
+  }
+
+  static torch::Tensor pad_right(const torch::Tensor& input,
+                                 int64_t padding_length,
+                                 int64_t dim) {
+    if (!input.defined() || padding_length == 0) {
+      return input;
+    }
+
+    const int64_t normalized_dim = normalize_dim(input, dim);
+    std::vector<int64_t> padding(static_cast<size_t>(input.dim() * 2), 0);
+    const int64_t padding_index = 2 * (input.dim() - normalized_dim - 1) + 1;
+    padding[static_cast<size_t>(padding_index)] = padding_length;
+    return torch::pad(input, padding, "constant", 0);
+  }
+
+  static torch::Tensor unpad_right(const torch::Tensor& input,
+                                   int64_t padding_length,
+                                   int64_t dim) {
+    if (!input.defined() || padding_length == 0) {
+      return input;
+    }
+
+    const int64_t normalized_dim = normalize_dim(input, dim);
+    CHECK_GE(input.size(normalized_dim), padding_length)
+        << "Padding length exceeds tensor size";
+    return input.narrow(normalized_dim,
+                        /*start=*/0,
+                        input.size(normalized_dim) - padding_length);
+  }
+
+  static int64_t padding_length(const std::string& tensor_name) {
+    auto padding_it = padding_lengths_.find(tensor_name);
+    CHECK(padding_it != padding_lengths_.end())
+        << "Missing sequence-parallel padding metadata: " << tensor_name;
+    return padding_it->second;
+  }
+
+  ProcessGroup* process_group_{nullptr};
+  const SequenceParallelTensorDims input_sequence_dims_;
+  const SequenceParallelTensorDims output_sequence_dims_;
+  inline static SequenceParallelTensorDims padding_lengths_;
+};
+
+}  // namespace xllm::dit
diff --git a/xllm/models/model_registry.cpp b/xllm/models/model_registry.cpp
index d5fe6e2b58..50bd2d9cd1 100644
--- a/xllm/models/model_registry.cpp
+++ b/xllm/models/model_registry.cpp
@@ -116,8 +116,9 @@ bool resolve_model_registration(const std::string& model_type,
   if (backend == kAutoBackend) {
     effective_backend =
         is_torch_only_model_type(model_type) ? kTorchBackend : kAtbBackend;
-  } else if (model_type == "qwen3" || model_type == "qwen3_moe") {
-    // qwen3/qwen3_moe support both backends.
+  } else if (model_type == "qwen3" || model_type == "qwen3_moe" ||
+             model_type == "qwen3_vl") {
+    // qwen3/qwen3_moe/qwen3_vl support both backends.
   } else if (is_torch_only_model_type(model_type)) {
     if (backend != kTorchBackend) {
       if (error_message != nullptr) {
@@ -141,6 +142,8 @@ bool resolve_model_registration(const std::string& model_type,
     *resolved_name = "qwen3_atb";
   } else if (model_type == "qwen3_moe" && effective_backend == kAtbBackend) {
     *resolved_name = "qwen3_moe_atb";
+  } else if (model_type == "qwen3_vl" && effective_backend == kAtbBackend) {
+    *resolved_name = "qwen3_vl_atb";
   } else {
     *resolved_name = model_type;
   }
diff --git a/xllm/models/models.h b/xllm/models/models.h
index 646e0343d8..6a99242a8f 100644
--- a/xllm/models/models.h
+++ b/xllm/models/models.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include "dit/pipelines/pipeline_flux2.h"                // IWYU pragma: keep
 #include "dit/pipelines/pipeline_flux_control.h"         // IWYU pragma: keep
 #include "dit/pipelines/pipeline_flux_fill.h"            // IWYU pragma: keep
+#include "dit/pipelines/pipeline_joyimage_edit_plus.h"   // IWYU pragma: keep
 #include "dit/pipelines/pipeline_qwenimage_edit_plus.h"  // IWYU pragma: keep
 #include "dit/pipelines/pipeline_wan_i2v.h"              // IWYU pragma: keep
 #include "llm/deepseek_v4.h"                             // IWYU pragma: keep
@@ -62,6 +63,7 @@ limitations under the License.
 #include "vlm/npu/qwen3_vl.h"                            // IWYU pragma: keep
 #include "vlm/npu/qwen3_vl_moe.h"                        // IWYU pragma: keep
 #include "vlm/qwen3_5.h"                                 // IWYU pragma: keep
+#include "vlm/qwen3_vl.h"                                // IWYU pragma: keep
 
 #elif defined(USE_MLU)
 #include "dit/pipelines/pipeline_flux.h"          // IWYU pragma: keep
diff --git a/xllm/models/vlm/npu/qwen3_vl.h b/xllm/models/vlm/npu/qwen3_vl.h
index 3c19ee7c51..a591f25bb8 100644
--- a/xllm/models/vlm/npu/qwen3_vl.h
+++ b/xllm/models/vlm/npu/qwen3_vl.h
@@ -841,11 +841,15 @@ TORCH_MODULE(Qwen3_VLForConditionalGeneration);
 using Qwen3VLMultimodalProcessor = MultimodalProcessor<Qwen3VLPromptProcessor,
                                                        Qwen2VLImageProcessor,
                                                        Qwen3VLVideoProcessor>;
-REGISTER_MULTIMODAL_PROCESSOR(qwen3_vl, Qwen3VLMultimodalProcessor);
-REGISTER_CAUSAL_VLM_MODEL(qwen3_vl, Qwen3_VLForConditionalGeneration);
+REGISTER_MULTIMODAL_PROCESSOR_WITH_VARNAME(qwen3_vl_atb,
+                                           qwen3_vl_atb,
+                                           Qwen3VLMultimodalProcessor);
+REGISTER_CAUSAL_VLM_MODEL_WITH_VARNAME(qwen3_vl_atb,
+                                       qwen3_vl_atb,
+                                       Qwen3_VLForConditionalGeneration);
 REGISTER_MPOSITION_GENERATOR(qwen3_vl, xllm::Qwen3VLMPositionGenerator);
 
-REGISTER_MODEL_ARGS(qwen3_vl, [&] {
+REGISTER_MODEL_ARGS_WITH_VARNAME(qwen3_vl_atb, qwen3_vl_atb, [&] {
   // text config
   // LOAD_ARG_OR(attention_dropout, "attention_dropout", 0.0);
   LOAD_ARG_OR(model_type, "model_type", "qwen3_vl");
diff --git a/xllm/models/vlm/qwen3_vl.h b/xllm/models/vlm/qwen3_vl.h
index 7350ae3cf9..bb86f84b20 100644
--- a/xllm/models/vlm/qwen3_vl.h
+++ b/xllm/models/vlm/qwen3_vl.h
@@ -461,12 +461,12 @@ class Qwen3_VisionTransformerImpl : public torch::nn::Module {
     m_sin = m_sin.repeat({1, 2});
 
     torch::Tensor cu_seqlens_cpu = cu_seqlens.cpu();
-    std::vector<int> cu_seqlens_vec(
+    std::vector<int32_t> cu_seqlens_vec(
         cu_seqlens_cpu.data_ptr<int>(),  // full seqlen vec
         cu_seqlens_cpu.data_ptr<int>() + cu_seqlens_cpu.numel());
     std::vector<torch::Tensor> deepstack_feature_lists;
     deepstack_feature_lists.reserve(deepstack_visual_indexes_.size());
-    for (int idx = 0; idx < layers_.size(); ++idx) {
+    for (int32_t idx = 0; idx < layers_.size(); ++idx) {
       hidden_states = layers_[idx](
           hidden_states, m_cos, m_sin, cu_seqlens, cu_seqlens_vec, idx);
       auto it = std::find(deepstack_visual_indexes_.begin(),
@@ -516,6 +516,23 @@ class Qwen3_VisionTransformerImpl : public torch::nn::Module {
     }
   }
 
+  void verify_loaded_weights(const std::string& prefix) const {
+    patch_embed_->verify_loaded_weights(prefix + "patch_embed.");
+    for (size_t idx = 0; idx < layers_.size(); ++idx) {
+      layers_[idx]->verify_loaded_weights(prefix + "blocks." +
+                                          std::to_string(idx) + ".");
+    }
+    merger_->verify_loaded_weights(prefix + "merger.");
+    for (size_t idx = 0; idx < deepstack_merger_layers_.size(); ++idx) {
+      deepstack_merger_layers_[idx]->verify_loaded_weights(
+          prefix + "deepstack_merger_list." + std::to_string(idx) + ".");
+    }
+    CHECK(is_emb_weight_loaded)
+        << "weight is not loaded for " << prefix + "pos_embed.weight";
+  }
+
+  void merge_loaded_weights() {}
+
  private:
   int hidden_size_ = 0;
   int num_heads_ = 0;
@@ -557,11 +574,13 @@ TORCH_MODULE(Qwen3_VLForConditionalGeneration);
 using Qwen3VLMultimodalProcessor = MultimodalProcessor<Qwen3VLPromptProcessor,
                                                        Qwen2VLImageProcessor,
                                                        Qwen3VLVideoProcessor>;
+
 REGISTER_MULTIMODAL_PROCESSOR(qwen3_vl, Qwen3VLMultimodalProcessor);
 REGISTER_CAUSAL_VLM_MODEL(qwen3_vl, Qwen3_VLForConditionalGeneration);
 REGISTER_MPOSITION_GENERATOR(qwen3_vl, Qwen3VLMPositionGenerator);
 
-REGISTER_MODEL_ARGS(qwen3_vl, [&] {
+const ModelArgsLoader kQwen3VLModelArgsLoader = [](const JsonReader& json,
+                                                   ModelArgs* args) {
   // text config
   // LOAD_ARG_OR(attention_dropout, "attention_dropout", 0.0);
   LOAD_ARG_OR(model_type, "model_type", "qwen3_vl");
@@ -617,7 +636,11 @@ REGISTER_MODEL_ARGS(qwen3_vl, [&] {
 
   LOAD_ARG_OR(
       rope_scaling_rope_type, "vision_config.rope_scaling.type", "mrope");
-
   LOAD_ARG_OR(vocab_size, "text_config.vocab_size", 151936);
-});
+  return true;
+};
+
+REGISTER_MODEL_ARGS_LOADER(qwen3_vl, kQwen3VLModelArgsLoader);
+REGISTER_MODEL_ARGS_LOADER(Qwen3VLForConditionalGeneration,
+                           kQwen3VLModelArgsLoader);
 }  // namespace xllm
diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp
index 223888524b..f8e017096f 100644
--- a/xllm/xllm.cpp
+++ b/xllm/xllm.cpp
@@ -182,6 +182,8 @@ Options create_options(const std::string& instance_name, bool is_local) {
       .sp_size(static_cast<int32_t>(parallel_config.sp_size()))
       .cfg_size(static_cast<int32_t>(parallel_config.cfg_size()))
       .vae_size(static_cast<int32_t>(parallel_config.vae_size()))
+      .text_encoder_tp_size(
+          static_cast<int32_t>(parallel_config.text_encoder_tp_size()))
       .instance_name(instance_name)
       .enable_disagg_pd(disagg_pd_config.enable_disagg_pd())
       .enable_pd_ooc(disagg_pd_config.enable_pd_ooc())