diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 45b2f9f1aa..9855b1c780 100755 --- a/xllm/core/common/global_flags.h +++ b/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 --- diff --git a/xllm/core/framework/config/kernel_config.cpp b/xllm/core/framework/config/kernel_config.cpp index bc6bf419fe..2f80319d87 100644 --- a/xllm/core/framework/config/kernel_config.cpp +++ b/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 } diff --git a/xllm/core/framework/config/kernel_config.h b/xllm/core/framework/config/kernel_config.h index 6be36a3456..1becc47a3d 100644 --- a/xllm/core/framework/config/kernel_config.h +++ b/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 }; 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, 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/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 4f1162859b..cfb8c6157e 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/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 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 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 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 all2all_work; process_group->all_to_all_single(output, @@ -337,13 +343,19 @@ std::function 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; }; } @@ -351,6 +363,12 @@ std::function all_to_all_4D(const torch::Tensor& input, // 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]; diff --git a/xllm/core/framework/parallel_state/parallel_state.h b/xllm/core/framework/parallel_state/parallel_state.h index 905f3c78d9..810d24bc3b 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 + #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 all_to_all_4D(const torch::Tensor& input_, - int32_t scatter_idx, - int32_t gather_idx, - bool is_sync, - ProcessGroup* pg); +std::function 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. 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/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 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 npu_fused_infer_attention( std::string layout = input_layout; char* input_layout_ptr = const_cast(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 c6da986340..d44196b554 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -47,6 +47,11 @@ 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); + std::tuple npu_fused_infer_attention( const torch::Tensor& query, const torch::Tensor& key, @@ -61,7 +66,8 @@ std::tuple 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..2baf629558 100644 --- a/xllm/core/kernels/npu/rope.cpp +++ b/xllm/core/kernels/npu/rope.cpp @@ -20,6 +20,40 @@ 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 query_shape = query.sizes().vec(); + const std::vector 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(torch::Tensor& q, torch::Tensor& k, const torch::Tensor& cos_sin_cache, 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 78ed058521..fba2d104bd 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/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); } diff --git a/xllm/core/layers/common/qwen2_attention.cpp b/xllm/core/layers/common/qwen2_attention.cpp index 1def5d5d36..ec351e4b38 100644 --- a/xllm/core/layers/common/qwen2_attention.cpp +++ b/xllm/core/layers/common/qwen2_attention.cpp @@ -19,6 +19,9 @@ limitations under the License. #include +#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_}); 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 #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 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& cu_seq_len_vec, + float scale) { + CHECK_GE(cu_seq_len_vec.size(), static_cast(2)); + CHECK_EQ(cu_seq_len_vec.front(), 0); + + std::vector 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..19b06c2cdb 100644 --- a/xllm/core/layers/common/rms_norm.cpp +++ b/xllm/core/layers/common/rms_norm.cpp @@ -46,6 +46,19 @@ std::tuple> RMSNormImpl::forward( std::optional residual, std::optional inplace_output) { auto org_shape = input.sizes().vec(); + + if (Platform::is_npu() && mode_ == kLayerNormMode) { + torch::Tensor norm_input = input; + std::optional 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 +161,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 + +#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 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& 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& 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(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 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 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 + +#include +#include +#include + +#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 forward( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& cos, + const torch::Tensor& sin, + const torch::Tensor& seq_len, + std::vector& 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& seq_len_host); + + atb_speed::Model::Node rope_node_; + std::vector 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/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 + +#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/embed_vlm_worker_impl.cpp b/xllm/core/runtime/embed_vlm_worker_impl.cpp index b56f4a5c4d..cd7fd381d4 100644 --- a/xllm/core/runtime/embed_vlm_worker_impl.cpp +++ b/xllm/core/runtime/embed_vlm_worker_impl.cpp @@ -26,6 +26,7 @@ limitations under the License. #include #include "common/metrics.h" +#include "core/framework/config/kernel_config.h" #include "core/framework/config/model_config.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_input_params.h" @@ -71,6 +72,11 @@ std::optional EmbedVLMWorkerImpl::step( auto model_output = model_executor_->forward( flatten_tokens, flatten_positions, kv_caches_, params); auto hidden_states = model_output.hidden_states; + if (::xllm::KernelConfig::get_instance() + .enable_return_prenorm_hidden_states() && + model_output.residual.defined()) { + hidden_states = model_output.residual; + } ret = device_.synchronize_default_stream(); COUNTER_ADD(execution_latency_seconds_model, timer.elapsed_seconds()); diff --git a/xllm/models/dit/attn_processor/attn_processor.h b/xllm/models/dit/attn_processor/attn_processor.h new file mode 100644 index 0000000000..736afe4d5d --- /dev/null +++ b/xllm/models/dit/attn_processor/attn_processor.h @@ -0,0 +1,60 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "core/framework/parallel_state/process_group.h" + +namespace xllm::dit { + +template +class AttnProcessor { + public: + explicit AttnProcessor(AttentionType& attention) : attention_(attention) {} + virtual ~AttnProcessor() = default; + + virtual OutputType forward(InputTypes... inputs) = 0; + + protected: + AttentionType& attention() { return attention_; } + + private: + AttentionType& attention_; +}; + +template +class SequenceParallelAttnProcessor + : public AttnProcessor { + public: + SequenceParallelAttnProcessor(AttentionType& attention, + ProcessGroup* process_group) + : AttnProcessor(attention), + process_group_(process_group) { + CHECK(process_group_ != nullptr) + << "Sequence-parallel attention requires a process group"; + CHECK_GT(process_group_->world_size(), 1) + << "Sequence-parallel attention requires world_size greater than one"; + } + + protected: + ProcessGroup* process_group() const { return process_group_; } + + private: + ProcessGroup* process_group_{nullptr}; +}; + +} // namespace xllm::dit diff --git a/xllm/models/dit/attn_processor/attn_processor_factory.h b/xllm/models/dit/attn_processor/attn_processor_factory.h new file mode 100644 index 0000000000..9bf45e3d53 --- /dev/null +++ b/xllm/models/dit/attn_processor/attn_processor_factory.h @@ -0,0 +1,79 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include "models/dit/attn_processor/attn_processor.h" +#include "models/dit/parallel_mode.h" + +namespace xllm::dit { + +template +class AttnProcessorFactory final { + public: + using Creator = std::function( + AttentionType& attention, + ProcessGroup* process_group)>; + + static AttnProcessorFactory& get_instance() { + static AttnProcessorFactory instance; + return instance; + } + + bool register_creator(const std::string& model_name, + ParallelMode parallel_mode, + Creator creator) { + return creators_[model_name] + .emplace(parallel_mode, std::move(creator)) + .second; + } + + std::unique_ptr create_attn_processor( + const std::string& model_name, + ParallelMode parallel_mode, + AttentionType& attention, + ProcessGroup* process_group) const { + auto model_it = creators_.find(model_name); + CHECK(model_it != creators_.end()) + << "No attention processors registered for model: " << model_name; + + auto processor_it = model_it->second.find(parallel_mode); + CHECK(processor_it != model_it->second.end()) + << "No attention processor registered for model: " << model_name + << ", parallel mode: " << static_cast(parallel_mode); + return processor_it->second(attention, process_group); + } + + AttnProcessorFactory(const AttnProcessorFactory&) = delete; + AttnProcessorFactory& operator=(const AttnProcessorFactory&) = delete; + + private: + using ModeCreators = std::unordered_map; + + AttnProcessorFactory() = default; + ~AttnProcessorFactory() = default; + + std::unordered_map creators_; +}; + +} // namespace xllm::dit 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 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/parallel_mode.h b/xllm/models/dit/parallel_mode.h new file mode 100644 index 0000000000..224eee6e4c --- /dev/null +++ b/xllm/models/dit/parallel_mode.h @@ -0,0 +1,47 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +#include "core/framework/parallel_state/process_group.h" +#include "models/dit/sequence_parallel/sequence_parallel_mixin.h" + +namespace xllm::dit { + +enum class ParallelMode : int8_t { + DEFAULT = 0, + SEQUENCE_PARALLEL = 1, +}; + +template +ParallelMode resolve_parallel_mode(ProcessGroup* process_group) { + if (process_group == nullptr || process_group->world_size() <= 1) { + return ParallelMode::DEFAULT; + } + + constexpr bool kSupportsSequenceParallel = + std::is_base_of_v; + CHECK(kSupportsSequenceParallel) + << "Sequence parallelism is enabled, but the model does not inherit " + "SequenceParallelMixin"; + return ParallelMode::SEQUENCE_PARALLEL; +} + +} // namespace xllm::dit 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..b8903c87e7 --- /dev/null +++ b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h @@ -0,0 +1,684 @@ +/* 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 + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/framework/config/parallel_config.h" +#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/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/util.h" +#include "models/model_registry.h" +#include "util/tensor_helper.h" + +namespace xllm { + +// NPU build: text encoding (Qwen3-VL) runs outside this pipeline; callers pass +// precomputed `prompt_embeds` (and `negative_prompt_embeds` for CFG). The +// in-pipeline VLM path is intentionally not compiled for NPU because the NPU +// Qwen3 language model returns only post-norm hidden states, whereas JoyImage +// requires the pre-norm last-layer output. +class JoyImageEditPlusPipelineImpl : public torch::nn::Module { + public: + using JoyImageShapeList = std::vector>>; + + struct TransformerForwardContext { + torch::Tensor rope_cos; + torch::Tensor rope_sin; + torch::Tensor attention_mask; + }; + + JoyImageEditPlusPipelineImpl(const DiTModelContext& 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")); + + 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_); + 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> 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 prepare_latents( + int64_t batch_size, + int64_t num_channels_latents, + int64_t height, + int64_t width, + int64_t seed, + const std::vector>& reference_images, + const torch::Tensor& provided_latents) { + std::vector all_patches; + std::vector all_target_masks; + std::vector>> 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 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(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 sample_patches; + std::vector sample_masks; + std::vector> 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 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(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> 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=*/"default")); + } + } + // Prompt embeddings are produced by a separate Qwen3-VL embedding service + // (two-stage serving) and passed in as `prompt_embeds` / + // `negative_prompt_embeds`, mirroring the Qwen-Image-Edit NPU path. A + // ones-mask is derived per sequence; CFG pads negatives/positives to equal + // length (padded positions get a zero mask). + bool do_cfg = guidance_scale > 1.0; + CHECK(input.prompt_embeds.defined()) + << "JoyImageEditPlus (NPU) requires precomputed `prompt_embeds` from " + "the Qwen3-VL embedding service."; + torch::Tensor prompt_embeds = + input.prompt_embeds.to(options_.device(), dtype_); + torch::Tensor prompt_embeds_mask = + torch::ones({prompt_embeds.size(0), prompt_embeds.size(1)}, + torch::TensorOptions().device(device_).dtype(torch::kLong)); + + torch::Tensor neg_embeds, neg_embeds_mask; + if (do_cfg) { + CHECK(input.negative_prompt_embeds.defined()) + << "JoyImageEditPlus CFG (guidance_scale > 1) requires precomputed " + "`negative_prompt_embeds`."; + 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)); + // 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(); + + // 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_); + + const int32_t cfg_size = ::xllm::ParallelConfig::get_instance().cfg_size(); + torch::Tensor transformer_encoder_hidden_states; + torch::Tensor transformer_encoder_hidden_states_mask; + JoyImageShapeList transformer_shape_list = shape_list; + bool transformer_use_cfg = false; + int64_t transformer_batch_size = batch_size; + if (!do_cfg) { + transformer_encoder_hidden_states = prompt_embeds; + transformer_encoder_hidden_states_mask = prompt_embeds_mask; + } else if (cfg_size == 2) { + CHECK(parallel_args_.dit_cfg_group_ != nullptr) + << "JoyImageEditPlus CFG parallel requires dit_cfg_group_"; + const int32_t rank = parallel_args_.dit_cfg_group_->rank(); + if (rank == 0) { + transformer_encoder_hidden_states = prompt_embeds; + transformer_encoder_hidden_states_mask = prompt_embeds_mask; + } else { + transformer_encoder_hidden_states = neg_embeds; + transformer_encoder_hidden_states_mask = neg_embeds_mask; + transformer_use_cfg = true; + } + } else { + transformer_encoder_hidden_states = + torch::cat({neg_embeds, prompt_embeds}, /*dim=*/0); + transformer_encoder_hidden_states_mask = + torch::cat({neg_embeds_mask, prompt_embeds_mask}, /*dim=*/0); + transformer_shape_list.insert( + transformer_shape_list.end(), shape_list.begin(), shape_list.end()); + transformer_batch_size *= 2; + } + + TransformerForwardContext transformer_context = + prepare_transformer_context(transformer_batch_size, + latents.size(1), + latents.device(), + transformer_encoder_hidden_states_mask, + transformer_shape_list); + + for (int64_t i = 0; i < timesteps.size(0); ++i) { + auto t = timesteps[i]; + // Restore reference patches. + latents.index_put_({~target_mask}, clean_backup.index({~target_mask})); + + torch::Tensor noise_pred; + if (do_cfg) { + torch::Tensor cond; + torch::Tensor uncond; + if (cfg_size == 2) { + auto t_expand = t.repeat({batch_size}); + torch::Tensor local_pred = + sequence_parallel_forward(latents, + t_expand, + transformer_encoder_hidden_states, + transformer_context, + transformer_use_cfg, + /*step_index=*/i + 1); + torch::Tensor gathered_pred = + xllm::parallel_state::gather(local_pred, + parallel_args_.dit_cfg_group_, + /*dim=*/0); + auto chunks = gathered_pred.chunk(2, 0); + cond = chunks[0]; + uncond = chunks[1]; + } else { + auto model_in = torch::cat({latents, latents}, 0); + auto t_expand = t.repeat({model_in.size(0)}); + auto pred = + sequence_parallel_forward(model_in, + t_expand, + transformer_encoder_hidden_states, + transformer_context, + /*use_cfg=*/false, + /*step_index=*/i + 1); + auto chunks = pred.chunk(2, 0); + uncond = chunks[0]; + cond = chunks[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{2}, /*keepdim=*/true); + auto noise_norm = + torch::norm(comb, 2, std::vector{2}, /*keepdim=*/true); + noise_pred = comb * (cond_norm / noise_norm.clamp_min(1e-6)); + } else { + auto t_expand = t.repeat({batch_size}); + noise_pred = + sequence_parallel_forward(latents, + t_expand, + transformer_encoder_hidden_states, + transformer_context, + /*use_cfg=*/false, + /*step_index=*/i + 1); + } + + latents = scheduler_->step(noise_pred, t, latents).to(latents.dtype()); + } + + // Restore refs and decode target patches per sample. + latents.index_put_({~target_mask}, clean_backup.index({~target_mask})); + std::vector 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 loader) { + LOG(INFO) << "JoyImageEditPlusPipeline loading from " + << loader->model_root_path(); + // Only the DiT (transformer) and VAE run in this service. The text encoder + // (Qwen3-VL) is a separate embedding service; its weights are not loaded + // here. + auto transformer_loader = loader->take_component_loader("transformer"); + auto vae_loader = loader->take_component_loader("vae"); + + 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(); + + LOG(INFO) << "JoyImageEditPlusPipeline loaded."; + } + + private: + torch::Tensor sequence_parallel_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) { + xllm::dit::SequenceParallelTensorMap model_outputs = + transformer_->sequence_parallel_forward( + {{"hidden_states", hidden_states}, + {"encoder_hidden_states", encoder_hidden_states}}, + [this, ×tep, &forward_context, use_cfg, step_index]( + const xllm::dit::SequenceParallelTensorMap& model_inputs) { + torch::Tensor output = transformer_->forward( + model_inputs.at("hidden_states"), + timestep, + model_inputs.at("encoder_hidden_states"), + forward_context.rope_cos, + forward_context.rope_sin, + forward_context.attention_mask, + use_cfg, + step_index); + return xllm::dit::SequenceParallelTensorMap{ + {"hidden_states", output}}; + }); + return model_outputs.at("hidden_states"); + } + + 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(shape_list.size()), batch_size) + << "shape_list batch size must match transformer batch size"; + + std::vector cos_list; + std::vector 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 cos_parts; + std::vector sin_parts; + int64_t temporal_offset = 0; + for (const auto& shape : shape_list[batch_index]) { + std::array start = {temporal_offset, 0, 0}; + std::array 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 !defined(USE_NPU) + 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); + attention_mask = full_mask.unsqueeze(1).unsqueeze(1); + } +#endif + + return {rope_cos, rope_sin, attention_mask}; + } + + std::pair rope_for_range( + const std::array& start, + const std::array& stop, + const torch::Device& device) { + std::vector 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 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 cos_parts; + std::vector 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(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 joyimage_bucket(int64_t h, int64_t w) { + static const std::vector> 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(h) / static_cast(w); + int64_t best_h = 1024, best_w = 1024; + double best_diff = std::numeric_limits::max(); + for (const auto& hw : kBuckets) { + double diff = + std::abs(static_cast(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 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}; + + 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 rope_dim_list_; + std::vector latents_mean_; + std::vector latents_std_; +}; + +TORCH_MODULE(JoyImageEditPlusPipeline); + +// Only the transformer and VAE components are instantiated by this pipeline; +// their args loaders live in transformer_joyimage_edit_plus.h and +// autoencoder_kl_wan.h. The text_encoder / processor / tokenizer components in +// model_index.json belong to the separate Qwen3-VL embedding service and are +// not loaded here (the DiT loader simply skips components without a factory). + +REGISTER_DIT_MODEL(JoyImageEditPlusPipeline, JoyImageEditPlusPipeline); +} // namespace xllm diff --git a/xllm/models/dit/pipelines/pipeline_wan_i2v.h b/xllm/models/dit/pipelines/pipeline_wan_i2v.h index 3bf0830083..9088803bea 100644 --- a/xllm/models/dit/pipelines/pipeline_wan_i2v.h +++ b/xllm/models/dit/pipelines/pipeline_wan_i2v.h @@ -231,7 +231,7 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module, image.options()); video_condition = torch::cat({image, zeros, last_img}, 2); } - video_condition = video_condition.to(options_.device()).to(torch::kFloat32); + video_condition = video_condition.to(options_.device(), options_.dtype()); torch::Tensor latents_mean = torch::tensor(latents_mean_, torch::dtype(torch::kFloat32)) @@ -624,7 +624,7 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module, torch::Tensor latents_std = 1.0 / latents_std_raw; prepared_latents = prepared_latents / latents_std; prepared_latents = prepared_latents + latents_mean; - video = vae_->decode(prepared_latents.to(torch::kFloat32)).sample; + video = vae_->decode(prepared_latents.to(options_.dtype())).sample; video = video_processor_->postprocess_video(video); return video; } diff --git a/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h b/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h new file mode 100644 index 0000000000..13dd09ad77 --- /dev/null +++ b/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h @@ -0,0 +1,114 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include "core/framework/parallel_state/parallel_state.h" +#include "models/dit/utils/sequence_parallel_pad_manager.h" + +namespace xllm::dit { + +using SequenceParallelTensorMap = + std::unordered_map; +using SequenceParallelTensorDims = std::unordered_map; + +class SequenceParallelMixin { + public: + 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 + SequenceParallelTensorMap sequence_parallel_forward( + const SequenceParallelTensorMap& inputs, + ForwardFn&& forward_fn) const { + SequenceParallelTensorMap split_inputs = + split_sequence_parallel_inputs(inputs); + SequenceParallelTensorMap outputs = + std::forward(forward_fn)(split_inputs); + return gather_sequence_parallel_outputs(outputs); + } + + private: + SequenceParallelTensorMap split_sequence_parallel_inputs( + const SequenceParallelTensorMap& inputs) const { + SequenceParallelTensorMap split_inputs = inputs; + if (!sequence_parallel_enabled()) { + return split_inputs; + } + + for (const auto& [tensor_name, sequence_dim] : input_sequence_dims_) { + auto tensor_it = split_inputs.find(tensor_name); + CHECK(tensor_it != split_inputs.end()) + << "Missing registered sequence-parallel input: " << tensor_name; + if (!tensor_it->second.defined()) { + continue; + } + + torch::Tensor padded_tensor = + SequenceParallelPadManager::get_instance().pad_tensor( + tensor_it->second, tensor_name, sequence_dim); + tensor_it->second = parallel_state::scatter( + padded_tensor, process_group_, static_cast(sequence_dim)); + } + return split_inputs; + } + + SequenceParallelTensorMap gather_sequence_parallel_outputs( + const SequenceParallelTensorMap& outputs) const { + SequenceParallelTensorMap gathered_outputs = outputs; + if (!sequence_parallel_enabled()) { + return gathered_outputs; + } + + for (const auto& [tensor_name, sequence_dim] : output_sequence_dims_) { + auto tensor_it = gathered_outputs.find(tensor_name); + CHECK(tensor_it != gathered_outputs.end()) + << "Missing registered sequence-parallel output: " << tensor_name; + if (!tensor_it->second.defined()) { + continue; + } + + tensor_it->second = + parallel_state::gather(tensor_it->second.contiguous(), + process_group_, + static_cast(sequence_dim)); + SequenceParallelPadManager::get_instance().unpad_tensor( + tensor_it->second, tensor_name, sequence_dim); + } + return gathered_outputs; + } + + bool sequence_parallel_enabled() const { + return process_group_ != nullptr && process_group_->world_size() > 1; + } + + ProcessGroup* process_group_{nullptr}; + SequenceParallelTensorDims input_sequence_dims_; + SequenceParallelTensorDims output_sequence_dims_; +}; + +} // namespace xllm::dit 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..626de0df73 --- /dev/null +++ b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h @@ -0,0 +1,969 @@ +/* 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 +#include +#if defined(USE_NPU) +#include +#endif + +#include +#include +#include +#include +#include + +#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/add_matmul.h" +#include "core/layers/common/rms_norm.h" +#include "models/dit/attn_processor/attn_processor.h" +#include "models/dit/attn_processor/attn_processor_factory.h" +#include "models/dit/transformers/transformer_qwen_image.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 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 +// (True == attend). NPU uses npu_fusion_attention (BSND layout); the fallback +// uses scaled_dot_product_attention. +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) + // JoyImage uses full bidirectional attention. Do not pass the pipeline-side + // padding mask to npu_fusion_attention: that mask is [B,1,1,Skv], while the + // NPU kernel requires an explicit Sq dimension ([B,1,Sq,Skv], [Sq,Skv], ...). + // This mirrors the QwenImageEditPlus/Wan DiT NPU paths, which run full + // attention with atten_mask=nullopt. + 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=*/torch::nullopt, + /*scale=*/std::pow(static_cast(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 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 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 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 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); + +// Internal module that only registers parameters used by Joy attention. +// Attention forward is implemented by processor classes. +class JoyAttentionImpl final : public torch::nn::Module { + public: + JoyAttentionImpl(const ModelContext& context, + int64_t dim, + int64_t num_heads, + int64_t head_dim, + double eps = 1e-6) + : options_(context.get_tensor_options()), heads_(num_heads) { + int64_t inner = num_heads * head_dim; + + img_attn_qkv_ = register_module( + "img_attn_qkv", + layer::AddMatmulWeightTransposed(dim, inner * 3, 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_qkv_ = register_module( + "txt_attn_qkv", + layer::AddMatmulWeightTransposed(dim, inner * 3, 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_)); + } + + void load_state_dict(const StateDict& state_dict) { + img_attn_qkv_->load_state_dict( + state_dict.get_dict_with_prefix("img_attn_qkv.")); + 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.")); + txt_attn_qkv_->load_state_dict( + state_dict.get_dict_with_prefix("txt_attn_qkv.")); + 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_qkv_->verify_loaded_weights(prefix + "img_attn_qkv."); + 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_qkv_->verify_loaded_weights(prefix + "txt_attn_qkv."); + 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."); + } + + public: + int64_t heads_; + layer::AddMatmulWeightTransposed img_attn_qkv_{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_qkv_{nullptr}; + layer::RMSNorm txt_attn_q_norm_{nullptr}; + layer::RMSNorm txt_attn_k_norm_{nullptr}; + layer::AddMatmulWeightTransposed txt_attn_proj_{nullptr}; + + private: + torch::TensorOptions options_; +}; +TORCH_MODULE(JoyAttention); + +using JoyAttnProcessorOutput = std::tuple; +using JoyAttnProcessorBase = xllm::dit::AttnProcessor; +using JoySequenceParallelAttnProcessorBase = + xllm::dit::SequenceParallelAttnProcessor; + +class JoyAttnProcessor final : public JoyAttnProcessorBase { + public: + explicit JoyAttnProcessor(JoyAttention& attention) + : JoyAttnProcessorBase(attention) {} + + JoyAttnProcessorOutput 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) override { + JoyAttention& attn = attention(); + auto img_qkv = attn->img_attn_qkv_->forward(hidden_states); + auto img_chunks = img_qkv.chunk(3, -1); + auto txt_qkv = attn->txt_attn_qkv_->forward(encoder_hidden_states); + auto txt_chunks = txt_qkv.chunk(3, -1); + + std::vector reshape = {attn->heads_, -1}; + auto img_q = img_chunks[0].unflatten(-1, reshape); + auto img_k = img_chunks[1].unflatten(-1, reshape); + auto img_v = img_chunks[2].unflatten(-1, reshape); + auto txt_q = txt_chunks[0].unflatten(-1, reshape); + auto txt_k = txt_chunks[1].unflatten(-1, reshape); + auto txt_v = txt_chunks[2].unflatten(-1, reshape); + + img_q = std::get<0>(attn->img_attn_q_norm_->forward(img_q)); + img_k = std::get<0>(attn->img_attn_k_norm_->forward(img_k)); + txt_q = std::get<0>(attn->txt_attn_q_norm_->forward(txt_q)); + txt_k = std::get<0>(attn->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); + } + + auto joint_q = torch::cat({img_q, txt_q}, /*dim=*/1); + auto joint_k = torch::cat({img_k, txt_k}, /*dim=*/1); + auto joint_v = torch::cat({img_v, txt_v}, /*dim=*/1); + auto joint = joyimage_joint_attention( + joint_q, joint_k, joint_v, attn->heads_, attn_mask); + joint = joint.to(joint_q.dtype()); + + const int64_t image_seq_len = img_q.size(1); + auto img_out = joint.slice( + /*dim=*/1, /*start=*/0, /*end=*/image_seq_len); + auto txt_out = joint.slice( + /*dim=*/1, /*start=*/image_seq_len, /*end=*/joint.size(1)); + img_out = img_out.flatten(/*start_dim=*/2, /*end_dim=*/3); + txt_out = txt_out.flatten(/*start_dim=*/2, /*end_dim=*/3); + + img_out = attn->img_attn_proj_->forward(img_out); + txt_out = attn->txt_attn_proj_->forward(txt_out); + return std::make_tuple(img_out, txt_out); + } +}; + +class JoySequenceParallelAttnProcessor final + : public JoySequenceParallelAttnProcessorBase { + public: + JoySequenceParallelAttnProcessor(JoyAttention& attention, + ProcessGroup* process_group) + : JoySequenceParallelAttnProcessorBase(attention, process_group) { + CHECK_EQ(attention->heads_ % process_group->world_size(), 0) + << "JoyImageEditPlus attention heads must be divisible by sp_size"; + } + + JoyAttnProcessorOutput 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) override { + JoyAttention& attn = attention(); + auto img_qkv = attn->img_attn_qkv_->forward(hidden_states); + auto img_chunks = img_qkv.chunk(3, -1); + + std::vector reshape = {attn->heads_, -1}; + auto img_q = img_chunks[0].unflatten(-1, reshape); + auto img_k = img_chunks[1].unflatten(-1, reshape); + auto img_v = img_chunks[2].unflatten(-1, reshape); + + auto img_q_handler = + xllm::parallel_state::all_to_all_4D(img_q, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"hidden_states"); + auto img_k_handler = + xllm::parallel_state::all_to_all_4D(img_k, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"hidden_states"); + auto img_v_handler = + xllm::parallel_state::all_to_all_4D(img_v, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"hidden_states"); + + auto txt_qkv = attn->txt_attn_qkv_->forward(encoder_hidden_states); + auto txt_chunks = txt_qkv.chunk(3, -1); + auto txt_q = txt_chunks[0].unflatten(-1, reshape); + auto txt_k = txt_chunks[1].unflatten(-1, reshape); + auto txt_v = txt_chunks[2].unflatten(-1, reshape); + + auto txt_q_handler = xllm::parallel_state::all_to_all_4D( + txt_q, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"encoder_hidden_states"); + auto txt_k_handler = xllm::parallel_state::all_to_all_4D( + txt_k, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"encoder_hidden_states"); + auto txt_v_handler = xllm::parallel_state::all_to_all_4D( + txt_v, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"encoder_hidden_states"); + + img_q = img_q_handler(); + img_k = img_k_handler(); + txt_q = txt_q_handler(); + txt_k = txt_k_handler(); + + img_q = std::get<0>(attn->img_attn_q_norm_->forward(img_q)); + img_k = std::get<0>(attn->img_attn_k_norm_->forward(img_k)); + txt_q = std::get<0>(attn->txt_attn_q_norm_->forward(txt_q)); + txt_k = std::get<0>(attn->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 = img_v_handler(); + txt_v = txt_v_handler(); + + auto joint_q = torch::cat({img_q, txt_q}, /*dim=*/1); + auto joint_k = torch::cat({img_k, txt_k}, /*dim=*/1); + auto joint_v = torch::cat({img_v, txt_v}, /*dim=*/1); + const int64_t local_heads = attn->heads_ / process_group()->world_size(); + auto joint = joyimage_joint_attention( + joint_q, joint_k, joint_v, local_heads, attn_mask); + joint = joint.to(joint_q.dtype()); + + const int64_t image_seq_len = img_q.size(1); + auto img_out = joint.slice( + /*dim=*/1, /*start=*/0, /*end=*/image_seq_len); + auto txt_out = joint.slice( + /*dim=*/1, /*start=*/image_seq_len, /*end=*/joint.size(1)); + + auto img_out_handler = + xllm::parallel_state::all_to_all_4D(img_out, + /*scatter_idx=*/1, + /*gather_idx=*/2, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"hidden_states"); + auto txt_out_handler = xllm::parallel_state::all_to_all_4D( + txt_out, + /*scatter_idx=*/1, + /*gather_idx=*/2, + /*async_ops=*/true, + process_group(), + /*enable_sp_pad=*/true, + /*tensor_name=*/"encoder_hidden_states"); + + img_out = img_out_handler().flatten(/*start_dim=*/2, /*end_dim=*/3); + img_out = attn->img_attn_proj_->forward(img_out); + txt_out = txt_out_handler().flatten(/*start_dim=*/2, /*end_dim=*/3); + txt_out = attn->txt_attn_proj_->forward(txt_out); + return std::make_tuple(img_out, txt_out); + } +}; + +inline constexpr char kJoyImageEditPlusModelName[] = + "JoyImageEditPlusTransformer3DModel"; + +using JoyAttnProcessorFactory = + xllm::dit::AttnProcessorFactory; + +inline void register_joy_attn_processors() { + static const bool kRegistered = []() { + JoyAttnProcessorFactory& factory = JoyAttnProcessorFactory::get_instance(); + const bool default_registered = factory.register_creator( + kJoyImageEditPlusModelName, + xllm::dit::ParallelMode::DEFAULT, + [](JoyAttention& attention, ProcessGroup* /*process_group*/) { + return std::make_unique(attention); + }); + const bool sequence_parallel_registered = factory.register_creator( + kJoyImageEditPlusModelName, + xllm::dit::ParallelMode::SEQUENCE_PARALLEL, + [](JoyAttention& attention, ProcessGroup* process_group) { + return std::make_unique( + attention, process_group); + }); + return default_registered && sequence_parallel_registered; + }(); + CHECK(kRegistered) << "Failed to register Joy attention processors"; +} + +// 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, + const std::string& model_name, + xllm::dit::ParallelMode parallel_mode, + ProcessGroup* sp_group, + double eps = 1e-6) + : eps_(eps) { + int64_t mlp_hidden = static_cast(dim * mlp_width_ratio); + + img_mod_ = register_module("img_mod", Modulate(dim, 6)); + img_mlp_ = register_module( + "img_mlp", qwenimage::FeedForward(context, dim, dim, /*mult=*/4)); + txt_mod_ = register_module("txt_mod", Modulate(dim, 6)); + txt_mlp_ = register_module( + "txt_mlp", qwenimage::FeedForward(context, dim, dim, /*mult=*/4)); + attn_ = register_module( + "attn", JoyAttention(context, dim, num_heads, head_dim, eps)); + attn_processor_ = + JoyAttnProcessorFactory::get_instance().create_attn_processor( + model_name, parallel_mode, attn_, sp_group); + (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 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 --- + 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); + + auto attn_out = attn_processor_->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); + + hs = hs + img_attn * img_g1.unsqueeze(1); + ehs = ehs + txt_attn * txt_g1.unsqueeze(1); + + // --- FFN --- + 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); + auto img_ffn = img_mlp_->forward(img_ffn_in); + auto txt_ffn = txt_mlp_->forward(txt_ffn_in); + hs = hs + img_ffn * img_g2.unsqueeze(1); + ehs = ehs + txt_ffn * txt_g2.unsqueeze(1); + + 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}; + qwenimage::FeedForward img_mlp_{nullptr}; + qwenimage::FeedForward txt_mlp_{nullptr}; + JoyAttention attn_{nullptr}; + std::unique_ptr attn_processor_; +}; +TORCH_MODULE(TransformerBlock); + +class JoyImageEditPlusTransformer3DModelImpl + : 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()) { + register_joy_attn_processors(); + parallel_mode_ = xllm::dit::resolve_parallel_mode< + JoyImageEditPlusTransformer3DModelImpl>(parallel_args.dit_sp_group_); + + 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"; + + // 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, + kJoyImageEditPlusModelName, + parallel_mode_, + 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) { + int64_t B = hidden_states.size(0); + int64_t N = hidden_states.size(1); + int64_t C = hidden_states.size(2); + int64_t pt = hidden_states.size(3); + int64_t ph = hidden_states.size(4); + int64_t pw = hidden_states.size(5); + + // 1. Condition embeddings. + auto cond = condition_embedder_->forward(timestep, 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{6, -1}); // [B,6,hidden] + + // 2. Patchify via Conv3d. + auto x = hidden_states.reshape({B * N, C, pt, ph, pw}); + x = img_in_->forward(x); // [B*N, hidden, 1, 1, 1] + auto img = x.reshape({B, N, 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(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({B, N, pt, ph, pw, out_channels_}) + .permute({0, 1, 5, 2, 3, 4}); + return img; + } + + 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 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 patch_size_; + xllm::dit::ParallelMode parallel_mode_{xllm::dit::ParallelMode::DEFAULT}; + + torch::nn::Conv3d img_in_{nullptr}; + TimeTextEmbedding condition_embedder_{nullptr}; + torch::nn::ModuleList double_blocks_{nullptr}; + std::vector 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{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{16, 56, 56})); + LOAD_ARG_OR(rope_theta_dit, "theta", 10000); +}); + +} // namespace joyimage +} // namespace xllm diff --git a/xllm/models/llm/npu/qwen3.h b/xllm/models/llm/npu/qwen3.h index 074a74a157..91e16e4db7 100644 --- a/xllm/models/llm/npu/qwen3.h +++ b/xllm/models/llm/npu/qwen3.h @@ -264,15 +264,27 @@ class QWen3ModelImpl : public LlmModelImplBase { } } } + // Capture the pre-norm output of the last decoder layer before the final + // RMSNorm. Diffusers' JoyImage-Edit-Plus text encoding consumes these + // pre-norm hidden states; the embedding service enables this flag so they + // are surfaced in ModelOutput.residual. `h` here already includes the + // residual stream (the last layer folds it in via intra-layer addnorm), so + // it equals the pre-norm hidden states regardless of interlayer addnorm. + torch::Tensor prenorm_hidden_states; + if (::xllm::KernelConfig::get_instance() + .enable_return_prenorm_hidden_states()) { + prenorm_hidden_states = h.clone(); + } auto hidden_states = norm_(h, 0); if (capture_aux_hidden_states_) { CHECK_EQ(capture_idx, static_cast(layers_to_capture_set_.size())) << "Captured aux hidden layer count mismatch."; torch::Tensor aux_hidden_states = aux_output_buffer_.slice(0, 0, num_tokens); - return ModelOutput(hidden_states, torch::Tensor(), aux_hidden_states); + return ModelOutput( + hidden_states, prenorm_hidden_states, aux_hidden_states); } - return ModelOutput(hidden_states); + return ModelOutput(hidden_states, prenorm_hidden_states); } protected: diff --git a/xllm/models/model_registry.cpp b/xllm/models/model_registry.cpp index 6305f59281..edfb2623a4 100644 --- a/xllm/models/model_registry.cpp +++ b/xllm/models/model_registry.cpp @@ -115,8 +115,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) { @@ -140,6 +141,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 5e583143b7..392f01f2bb 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..3ab5b21a0a 100644 --- a/xllm/models/vlm/npu/qwen3_vl.h +++ b/xllm/models/vlm/npu/qwen3_vl.h @@ -789,6 +789,10 @@ class Qwen3_VLForConditionalGenerationImpl : public torch::nn::Module { torch::Tensor pooler(const torch::Tensor& hidden_states, const torch::Tensor& seleted_idxes) { + if (::xllm::ModelConfig::get_instance() + .enable_return_mm_full_embeddings()) { + return hidden_states; + } return language_model_->pooler(hidden_states, seleted_idxes); } @@ -841,11 +845,15 @@ TORCH_MODULE(Qwen3_VLForConditionalGeneration); using Qwen3VLMultimodalProcessor = MultimodalProcessor; -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..37566bed7e 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 cu_seqlens_vec( + std::vector cu_seqlens_vec( cu_seqlens_cpu.data_ptr(), // full seqlen vec cu_seqlens_cpu.data_ptr() + cu_seqlens_cpu.numel()); std::vector 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,6 +574,7 @@ TORCH_MODULE(Qwen3_VLForConditionalGeneration); using Qwen3VLMultimodalProcessor = MultimodalProcessor; + REGISTER_MULTIMODAL_PROCESSOR(qwen3_vl, Qwen3VLMultimodalProcessor); REGISTER_CAUSAL_VLM_MODEL(qwen3_vl, Qwen3_VLForConditionalGeneration); REGISTER_MPOSITION_GENERATOR(qwen3_vl, Qwen3VLMPositionGenerator); diff --git a/xllm/models/vlm/qwen3_vl_base.h b/xllm/models/vlm/qwen3_vl_base.h index f624c511b9..c8edcd9df6 100644 --- a/xllm/models/vlm/qwen3_vl_base.h +++ b/xllm/models/vlm/qwen3_vl_base.h @@ -15,6 +15,7 @@ limitations under the License. #pragma once +#include "core/framework/config/model_config.h" #include "core/framework/model/model_output.h" #include "core/framework/multimodal/mm_data_item.h" #include "core/layers/common/lm_head.h" @@ -183,6 +184,10 @@ class Qwen3VLForConditionalGenerationBase : public torch::nn::Module { torch::Tensor pooler(const torch::Tensor& hidden_states, const torch::Tensor& seleted_idxes) { + if (::xllm::ModelConfig::get_instance() + .enable_return_mm_full_embeddings()) { + return hidden_states; + } return language_model_->pooler(hidden_states, seleted_idxes); }