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

xllm/core/common/global_flags.h

@@ -396,6 +396,8 @@ DECLARE_bool(enable_split_rmsnorm_rope);
 DECLARE_bool(enable_aclnn_matmul);
 
 DECLARE_bool(enable_aclnn_swiglu);
+
+DECLARE_bool(enable_return_prenorm_hidden_states);
 #endif
 
 // --- chat template config ---

xllm/core/framework/config/kernel_config.cpp

@@ -46,6 +46,12 @@ DEFINE_bool(enable_split_rmsnorm_rope,
             false,
             "enable fused split rmsnorm rope ops.");
 
+DEFINE_bool(enable_return_prenorm_hidden_states,
+            false,
+            "return the pre-norm output of the last decoder layer in "
+            "ModelOutput.residual (used by the Qwen3-VL embedding service for "
+            "JoyImage-Edit-Plus, which consumes pre-norm hidden states).");
+
 DEFINE_bool(enable_aclnn_matmul,
             false,
             "enable ACLNN matmul backend for supported NPU ATB layers.");
@@ -85,6 +91,7 @@ void KernelConfig::from_flags() {
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_split_rmsnorm_rope);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_aclnn_matmul);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_aclnn_swiglu);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_return_prenorm_hidden_states);
 #endif
 }
 
@@ -98,6 +105,7 @@ void KernelConfig::from_json(const JsonReader& json) {
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_split_rmsnorm_rope);
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_aclnn_matmul);
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_aclnn_swiglu);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(enable_return_prenorm_hidden_states);
 #endif
 }
 
@@ -121,6 +129,8 @@ void KernelConfig::append_config_json(
       config_json, default_config, enable_aclnn_matmul);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, enable_aclnn_swiglu);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, enable_return_prenorm_hidden_states);
 #endif
 }
 

xllm/core/framework/config/kernel_config.h

@@ -47,7 +47,8 @@ class KernelConfig final {
          "enable_interlayer_addnorm",
          "enable_split_rmsnorm_rope",
          "enable_aclnn_matmul",
-         "enable_aclnn_swiglu"}};
+         "enable_aclnn_swiglu",
+         "enable_return_prenorm_hidden_states"}};
     return kOptionCategory;
   }
 
@@ -67,6 +68,8 @@ class KernelConfig final {
   PROPERTY(bool, enable_aclnn_matmul) = false;
 
   PROPERTY(bool, enable_aclnn_swiglu) = false;
+
+  PROPERTY(bool, enable_return_prenorm_hidden_states) = false;
 #endif
 };
 

xllm/core/framework/model/model_args.h

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

xllm/core/framework/parallel_state/parallel_state.cpp

@@ -16,6 +16,7 @@ limitations under the License.
 #include "parallel_state.h"
 
 #include "core/util/utils.h"
+#include "models/dit/utils/sequence_parallel_pad_manager.h"
 #include "runtime/options.h"
 #include "util/net.h"
 
@@ -262,7 +263,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];
@@ -272,7 +272,9 @@ std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input,
                                              int32_t scatter_idx,
                                              int32_t gather_idx,
                                              bool async_ops,
-                                             ProcessGroup* process_group) {
+                                             ProcessGroup* process_group,
+                                             bool enable_sp_pad,
+                                             const std::string& tensor_name) {
   if (!process_group) {
     return [input]() { return input; };
   }
@@ -282,12 +284,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"
@@ -323,7 +323,13 @@ std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input,
                    .transpose(0, 1)
                    .contiguous()
                    .reshape({bs, seqlen, shard_head_num, head_size});
-      return [output]() { return output; };
+      return [output, enable_sp_pad, tensor_name]() mutable {
+        if (enable_sp_pad) {
+          xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor(
+              output, tensor_name, /*dim=*/1);
+        }
+        return output;
+      };
     } else {
       c10::intrusive_ptr<c10d::Work> all2all_work;
       process_group->all_to_all_single(output,
@@ -337,20 +343,32 @@ std::function<torch::Tensor()> all_to_all_4D(const torch::Tensor& input,
               bs,
               seqlen,
               shard_head_num,
-              head_size]() mutable -> torch::Tensor {
+              head_size,
+              enable_sp_pad,
+              tensor_name]() 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});
+        if (enable_sp_pad) {
+          xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor(
+              comm_output, tensor_name, /*dim=*/1);
+        }
         return comm_output;
       };
     }
   } else if (scatter_idx == 1 && gather_idx == 2) {
     // branch B : from "head shard" -> "sequence shard"
     // input: (bs, seqlen, head_num / group_size, head_size)
     // output (bs, seqlen / group_size, head_num, haed_size)
+    if (enable_sp_pad) {
+      send_input =
+          xllm::dit::SequenceParallelPadManager::get_instance().pad_tensor(
+              send_input, tensor_name, /*dim=*/1);
+    }
+
     auto sizes = send_input.sizes().vec();
     const int64_t bs = sizes[0];
     const int64_t seqlen = sizes[1];

xllm/core/framework/parallel_state/parallel_state.h

@@ -15,6 +15,8 @@ limitations under the License.
 
 #pragma once
 
+#include <string>
+
 #include "parallel_args.h"
 #include "process_group.h"
 
@@ -72,11 +74,14 @@ torch::Tensor scatter(torch::Tensor input,
                       ProcessGroup* process_group,
                       int dim = -1);
 
-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);
+std::function<torch::Tensor()> all_to_all_4D(
+    const torch::Tensor& input,
+    int32_t scatter_idx,
+    int32_t gather_idx,
+    bool async_ops,
+    ProcessGroup* process_group,
+    bool enable_sp_pad = false,
+    const std::string& tensor_name = "");
 
 // Create a process group where each process has a single device
 // devices: list of devices to create process groups on.

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
+}  // namespace xllm::kernel::npu

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;

xllm/core/kernels/npu/npu_ops_api.h

@@ -47,6 +47,16 @@ void batch_decode(const torch::Tensor& query,
                   const torch::Tensor& seq_lens,
                   torch::Tensor& output);
 
+void apply_rotary_pos_emb(torch::Tensor& query,
+                          torch::Tensor& key,
+                          const torch::Tensor& cos,
+                          const torch::Tensor& sin);
+
+void apply_rotary_pos_emb_atb(torch::Tensor& query,
+                              torch::Tensor& key,
+                              const torch::Tensor& cos,
+                              const torch::Tensor& sin);
+
 std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
     const torch::Tensor& query,
     const torch::Tensor& key,
@@ -61,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,

xllm/core/kernels/npu/rope.cpp

@@ -20,6 +20,58 @@ limitations under the License.
 
 namespace xllm::kernel::npu {
 
+void apply_rotary_pos_emb(torch::Tensor& query,
+                          torch::Tensor& key,
+                          const torch::Tensor& cos,
+                          const torch::Tensor& sin) {
+  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, "BSND");
+  query = std::get<0>(rotary_result).view(query_shape);
+  key = std::get<1>(rotary_result).view(key_shape);
+}
+
+void apply_rotary_pos_emb_atb(torch::Tensor& query,
+                              torch::Tensor& key,
+                              const torch::Tensor& cos,
+                              const torch::Tensor& sin) {
+  const int64_t rotary_dim = query.size(-1) / 2;
+  auto cos_input = cos.size(-1) == query.size(-1)
+                       ? cos.slice(/*dim=*/-1, /*start=*/0, /*end=*/rotary_dim)
+                       : cos;
+  auto sin_input = sin.size(-1) == query.size(-1)
+                       ? sin.slice(/*dim=*/-1, /*start=*/0, /*end=*/rotary_dim)
+                       : sin;
+  atb::npu_rotary_embedding_with_cos_sin(cos_input.contiguous(),
+                                         sin_input.contiguous(),
+                                         query,
+                                         key,
+                                         /*is_neox_style=*/true);
+}
+
 void apply_rotary(torch::Tensor& q,
                   torch::Tensor& k,
                   const torch::Tensor& cos_sin_cache,
@@ -40,7 +92,7 @@ void apply_rotary(torch::Tensor& q,
   q = q.view({1, q.size(0), -1, rotary_dim});
   k = k.view({1, k.size(0), -1, rotary_dim});
 
-  at_npu::native::custom_ops::npu_apply_rotary_pos_emb(q, k, cos, sin);
+  apply_rotary_pos_emb(q, k, cos, sin);
 }
 
 std::pair<torch::Tensor, torch::Tensor> apply_npu_partial_rotary_embedding(

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

xllm/core/layers/common/dense_mlp.cpp

@@ -109,7 +109,6 @@ torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
           {batch_size, intermediate_size_ / process_group_->world_size()},
           gate_up.options());
     }
-
     act_->forward(gate_up, output);
     return down_proj_->forward(output);
   }

xllm/core/layers/common/qwen2_attention.cpp

@@ -19,6 +19,9 @@ limitations under the License.
 
 #include <tuple>
 
+#if defined(USE_NPU)
+#include "kernels/npu/npu_ops_api.h"
+#endif
 #if defined(USE_CUDA) || defined(USE_DCU)
 #include "kernels/cuda/cuda_ops_api.h"
 #endif
@@ -180,7 +183,17 @@ torch::Tensor Qwen2AttentionImpl::forward(
 
   // 4. rope
   if (!fused_qk_norm_rope_applied) {
+#if defined(USE_NPU)
+    if (attn_metadata.mrope_cos.defined() &&
+        attn_metadata.mrope_sin.defined()) {
+      xllm::kernel::npu::apply_rotary_pos_emb(
+          q, k, attn_metadata.mrope_cos, attn_metadata.mrope_sin);
+    } else {
+      rotary_emb_->forward(q, k, positions, attn_metadata);
+    }
+#else
     rotary_emb_->forward(q, k, positions, attn_metadata);
+#endif
   }
   q = q.view({T, q_size_});
   k = k.view({T, kv_size_});