From cfb8e8e224a7bf263e6e9ba10f5dd92d070155ab Mon Sep 17 00:00:00 2001 From: shan-chen-feng Date: Tue, 28 Jul 2026 10:28:54 +0800 Subject: [PATCH 1/4] feat: support joyai-image_edit_plus inference on NPU device. --- xllm/core/common/global_flags.h | 2 + xllm/core/common/options.h | 2 + .../core/distributed_runtime/dist_manager.cpp | 4 +- xllm/core/distributed_runtime/master.cpp | 4 +- .../spawn_worker_protocol.h | 3 +- .../spawn_worker_server.cpp | 5 + .../spawn_worker_server/spawn_worker_server.h | 1 + .../spawn_worker_server_process.cpp | 11 +- .../distributed_runtime/worker_server.cpp | 21 +- .../core/framework/config/parallel_config.cpp | 8 + xllm/core/framework/config/parallel_config.h | 3 + xllm/core/framework/model/model_args.h | 6 + .../dit_collective_communicator.cpp | 24 +- .../dit_collective_communicator.h | 4 +- .../framework/parallel_state/dit_mapping.cpp | 32 + .../framework/parallel_state/dit_mapping.h | 3 + .../framework/parallel_state/parallel_args.h | 8 + .../parallel_state/parallel_state.cpp | 32 +- .../framework/parallel_state/parallel_state.h | 15 +- xllm/core/kernels/npu/active.cpp | 12 +- .../kernels/npu/npu_fused_infer_attention.cpp | 5 +- xllm/core/kernels/npu/npu_ops_api.h | 8 +- xllm/core/kernels/npu/rope.cpp | 34 + xllm/core/layers/common/CMakeLists.txt | 2 + xllm/core/layers/common/dense_mlp.cpp | 1 + xllm/core/layers/common/qwen2_attention.cpp | 13 + .../layers/common/qwen2_vision_attention.cpp | 57 ++ .../layers/common/qwen2_vision_attention.h | 6 + xllm/core/layers/common/rms_norm.cpp | 20 + xllm/core/layers/common/rms_norm.h | 2 + xllm/core/layers/npu/npu_rope_layer_impl.cpp | 169 +++ xllm/core/layers/npu/npu_rope_layer_impl.h | 72 ++ xllm/core/layers/npu_torch/attention.cpp | 34 +- xllm/core/layers/qwen3_vision_layer.cpp | 59 +- xllm/core/layers/qwen3_vision_layer.h | 27 +- xllm/core/runtime/options.h | 4 + .../dit/attn_processor/attn_processor.h | 60 ++ .../attn_processor/attn_processor_factory.h | 79 ++ .../dit/autoencoders/autoencoder_kl_wan.h | 41 +- xllm/models/dit/parallel_mode.h | 47 + .../pipelines/pipeline_joyimage_edit_plus.h | 955 +++++++++++++++++ xllm/models/dit/pipelines/pipeline_wan_i2v.h | 4 +- .../sequence_parallel_mixin.h | 114 +++ .../transformer_joyimage_edit_plus.h | 969 ++++++++++++++++++ xllm/models/model_registry.cpp | 7 +- xllm/models/models.h | 2 + xllm/models/vlm/npu/qwen3_vl.h | 10 +- xllm/models/vlm/qwen3_vl.h | 33 +- xllm/xllm.cpp | 2 + 49 files changed, 2945 insertions(+), 91 deletions(-) create mode 100644 xllm/core/layers/npu/npu_rope_layer_impl.cpp create mode 100644 xllm/core/layers/npu/npu_rope_layer_impl.h create mode 100644 xllm/models/dit/attn_processor/attn_processor.h create mode 100644 xllm/models/dit/attn_processor/attn_processor_factory.h create mode 100644 xllm/models/dit/parallel_mode.h create mode 100644 xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h create mode 100644 xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h create mode 100644 xllm/models/dit/transformers/transformer_joyimage_edit_plus.h diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 50cf1c6fc9..df93aa5b2e 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -113,6 +113,8 @@ DECLARE_int64(cfg_size); DECLARE_int64(vae_size); +DECLARE_int64(text_encoder_tp_size); + DECLARE_bool(enable_mm_encoder_dp); DECLARE_bool(enable_multi_stream_parallel); diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index c9b3ac67e3..e0839c9863 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -147,6 +147,8 @@ class Options { PROPERTY(int32_t, vae_size) = 1; + PROPERTY(int32_t, text_encoder_tp_size) = 1; + PROPERTY(std::optional, instance_name); PROPERTY(bool, enable_disagg_pd) = false; diff --git a/xllm/core/distributed_runtime/dist_manager.cpp b/xllm/core/distributed_runtime/dist_manager.cpp index eb66ba45d5..4ca3878761 100644 --- a/xllm/core/distributed_runtime/dist_manager.cpp +++ b/xllm/core/distributed_runtime/dist_manager.cpp @@ -189,12 +189,14 @@ void DistManager::setup_multi_node_workers( const int32_t sp_size = options.sp_size(); const int32_t cfg_size = options.cfg_size(); const int32_t vae_size = options.vae_size(); + const int32_t text_encoder_tp_size = options.text_encoder_tp_size(); LOG(INFO) << "Multi-node serving world_size = " << world_size << ", each_node_ranks = " << each_node_ranks << ", current node rank = " << options.node_rank() << ", nnodes = " << options.nnodes() << ", dp_size = " << dp_size << ", tp_size = " << tp_size << ", sp_size = " << sp_size - << ", cfg_size = " << cfg_size << ", vae_size = " << vae_size; + << ", cfg_size = " << cfg_size << ", vae_size = " << vae_size + << ", text_encoder_tp_size = " << text_encoder_tp_size; } else { LOG(INFO) << "Multi-node serving world_size = " << world_size << ", each_node_ranks = " << each_node_ranks diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index dd48dc3894..73603ba87d 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -615,6 +615,7 @@ Master::Master(const Options& options, EngineType type) .model_id(options.model_id()) .devices(devices) .backend(options.backend()) + .npu_kernel_backend(options_.npu_kernel_backend()) .enable_prefix_cache(options_.enable_prefix_cache()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .enable_offline_inference(options_.enable_offline_inference()) @@ -634,7 +635,8 @@ Master::Master(const Options& options, EngineType type) .tp_size(options_.tp_size()) .sp_size(options_.sp_size()) .cfg_size(options_.cfg_size()) - .vae_size(options_.vae_size()); + .vae_size(options_.vae_size()) + .text_encoder_tp_size(options_.text_encoder_tp_size()); auto dit_engine = std::make_unique(eng_options); engine_ = std::move(dit_engine); diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h index b51c5d84b4..f8aafd95e7 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_protocol.h @@ -21,10 +21,11 @@ limitations under the License. namespace xllm::spawn_worker_protocol { -inline constexpr int32_t kArgumentCount = 36; +inline constexpr int32_t kArgumentCount = 37; inline constexpr int32_t kMinimumArgumentCount = 34; inline constexpr int32_t kIndexerCacheDtypeArgumentIndex = 34; inline constexpr int32_t kEnableMtpDraftBodyTp1ArgumentIndex = 35; +inline constexpr int32_t kTextEncoderTpSizeArgumentIndex = 36; inline constexpr char kDefaultIndexerCacheDtype[] = "auto"; inline std::optional parse_indexer_cache_dtype( diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp index 5d621351fb..b875fc0be2 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.cpp @@ -88,6 +88,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, int32_t tp_size, int32_t sp_size, int32_t cfg_size, + int32_t text_encoder_tp_size, int32_t cp_size, int32_t ep_size, const InstanceRole& instance_role, @@ -100,6 +101,8 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, const bool is_dit_backend = backend == "dit"; const int32_t effective_sp_size = is_dit_backend ? sp_size : 1; const int32_t effective_cfg_size = is_dit_backend ? cfg_size : 1; + const int32_t effective_text_encoder_tp_size = + is_dit_backend ? text_encoder_tp_size : 1; runner_options.block_size(block_size) .backend(backend) .world_size(world_size) @@ -119,6 +122,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, .tp_size(tp_size) .sp_size(effective_sp_size) .cfg_size(effective_cfg_size) + .text_encoder_tp_size(effective_text_encoder_tp_size) .enable_shm(enable_shm) .input_shm_size(input_shm_size) .output_shm_size(output_shm_size) @@ -142,6 +146,7 @@ SpawnWorkerServer::SpawnWorkerServer(const std::string& master_node_addr, .tp_size(tp_size) .sp_size(effective_sp_size) .cfg_size(effective_cfg_size) + .text_encoder_tp_size(effective_text_encoder_tp_size) .communication_backend(communication_backend); DistributedConfig::get_instance().master_node_addr(master_node_addr); KVCacheConfig::get_instance() diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h index 3b4935f2a2..3ce9240d55 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server.h @@ -59,6 +59,7 @@ class SpawnWorkerServer final { int32_t tp_size, int32_t sp_size, int32_t cfg_size, + int32_t text_encoder_tp_size, int32_t cp_size, int32_t ep_size, const InstanceRole& instance_role, diff --git a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp index a2cdbc2c18..e676810f66 100644 --- a/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp +++ b/xllm/core/distributed_runtime/spawn_worker_server/spawn_worker_server_process.cpp @@ -60,6 +60,7 @@ limitations under the License. // @instance_role // @indexer_cache_dtype // @enable_mtp_draft_body_tp1 +// @text_encoder_tp_size int main(int argc, char* argv[]) { const std::optional parsed_indexer_cache_dtype = xllm::spawn_worker_protocol::parse_indexer_cache_dtype(argc, argv); @@ -117,8 +118,14 @@ int main(int argc, char* argv[]) { static_cast( atoi(argv[xllm::spawn_worker_protocol:: kEnableMtpDraftBodyTp1ArgumentIndex])) > 0; + const int32_t text_encoder_tp_size = + argc > xllm::spawn_worker_protocol::kTextEncoderTpSizeArgumentIndex + ? static_cast( + atoi(argv[xllm::spawn_worker_protocol:: + kTextEncoderTpSizeArgumentIndex])) + : 1; if (world_size < 1 || global_rank < 0 || global_rank >= world_size || - cp_size < 1 || ep_size < 1 || + cp_size < 1 || ep_size < 1 || text_encoder_tp_size < 1 || (instance_role_str != "DEFAULT" && instance_role_str != "PREFILL" && instance_role_str != "DECODE")) { LOG(ERROR) << "Invalid spawn worker topology: global_rank=" << global_rank @@ -159,6 +166,7 @@ int main(int argc, char* argv[]) { << ", max_encoder_cache_size = " << max_encoder_cache_size << ", dp_size = " << dp_size << ", tp_size = " << tp_size << ", sp_size = " << sp_size << ", cfg_size = " << cfg_size + << ", text_encoder_tp_size = " << text_encoder_tp_size << ", indexer_cache_dtype = " << indexer_cache_dtype << ", enable_mtp_draft_body_tp1 = " << enable_mtp_draft_body_tp1 << "\n"; @@ -193,6 +201,7 @@ int main(int argc, char* argv[]) { tp_size, sp_size, cfg_size, + text_encoder_tp_size, cp_size, ep_size, instance_role, diff --git a/xllm/core/distributed_runtime/worker_server.cpp b/xllm/core/distributed_runtime/worker_server.cpp index 42598464c5..9b25a3851c 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -166,14 +166,15 @@ void WorkerServer::create_server(const runtime::Options& options, const ParallelArgs* parallel_args = nullptr; std::unique_ptr comm; if (worker_type == WorkerType::DIT) { - auto dit_comm = - std::make_unique(worker_global_rank, - world_size, - options.dp_size(), - options.tp_size(), - options.sp_size(), - options.cfg_size(), - options.vae_size()); + auto dit_comm = std::make_unique( + worker_global_rank, + world_size, + options.dp_size(), + options.tp_size(), + options.sp_size(), + options.cfg_size(), + options.vae_size(), + options.text_encoder_tp_size()); comm = std::move(dit_comm); } else { auto common_comm = std::make_unique( @@ -315,6 +316,9 @@ void WorkerServer::create_spawn_server(int32_t local_rank, const char* sp_size_ptr = sp_size_str.c_str(); std::string cfg_size_str = std::to_string(options.cfg_size()); const char* cfg_size_ptr = cfg_size_str.c_str(); + std::string text_encoder_tp_size_str = + std::to_string(options.text_encoder_tp_size()); + const char* text_encoder_tp_size_ptr = text_encoder_tp_size_str.c_str(); const std::string& indexer_cache_dtype = KVCacheConfig::get_instance().indexer_cache_dtype(); const char* indexer_cache_dtype_ptr = indexer_cache_dtype.c_str(); @@ -372,6 +376,7 @@ void WorkerServer::create_spawn_server(int32_t local_rank, instance_role_ptr, indexer_cache_dtype_ptr, enable_mtp_draft_body_tp1_ptr, + text_encoder_tp_size_ptr, nullptr}; static_assert(std::size(argv) == spawn_worker_protocol::kArgumentCount + 1); pid_t pid; diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index d29e1e2786..f333d53993 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -44,6 +44,10 @@ DEFINE_int64(cfg_size, DEFINE_int64(vae_size, 1, "Vae patch parallelism size"); +DEFINE_int64(text_encoder_tp_size, + 1, + "Text encoder tensor parallelism size for DiT model."); + DEFINE_string( communication_backend, "hccl", @@ -80,6 +84,7 @@ void ParallelConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(cfg_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(vae_size); + XLLM_CONFIG_ASSIGN_FROM_FLAG(text_encoder_tp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(communication_backend); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_mm_encoder_dp); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_multi_stream_parallel); @@ -95,6 +100,7 @@ void ParallelConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size); XLLM_CONFIG_ASSIGN_FROM_JSON(vae_size); + XLLM_CONFIG_ASSIGN_FROM_JSON(text_encoder_tp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(communication_backend); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_mm_encoder_dp); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_multi_stream_parallel); @@ -114,6 +120,8 @@ void ParallelConfig::append_config_json( config_json, default_config, cfg_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, vae_size); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, text_encoder_tp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, communication_backend); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..bae219a705 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -48,6 +48,7 @@ class ParallelConfig final { "sp_size", "cfg_size", "vae_size", + "text_encoder_tp_size", "communication_backend", "enable_mm_encoder_dp", "enable_multi_stream_parallel", @@ -73,6 +74,8 @@ class ParallelConfig final { PROPERTY(int64_t, vae_size) = 1; + PROPERTY(int64_t, text_encoder_tp_size) = 1; + PROPERTY(std::string, communication_backend) = "hccl"; PROPERTY(bool, enable_mm_encoder_dp) = false; diff --git a/xllm/core/framework/model/model_args.h b/xllm/core/framework/model/model_args.h index 4f6b7cc570..4a0a371473 100644 --- a/xllm/core/framework/model/model_args.h +++ b/xllm/core/framework/model/model_args.h @@ -573,6 +573,12 @@ struct ModelArgs { PROPERTY(bool, zero_cond_t) = false; PROPERTY(bool, use_additional_t_cond) = false; PROPERTY(bool, use_layer3d_rope) = false; + + // JoyImage-Edit-Plus dit related args + PROPERTY(double, mlp_width_ratio) = 4.0; + PROPERTY(int64_t, text_dim) = 4096; + PROPERTY(std::vector, rope_dim_list) = { 16, 56, 56 }; + PROPERTY(int64_t, rope_theta_dit) = 10000; }; // Qwen hybrid models may describe full-attention layers explicitly via diff --git a/xllm/core/framework/parallel_state/dit_collective_communicator.cpp b/xllm/core/framework/parallel_state/dit_collective_communicator.cpp index cfbaa13248..428d8efb0f 100644 --- a/xllm/core/framework/parallel_state/dit_collective_communicator.cpp +++ b/xllm/core/framework/parallel_state/dit_collective_communicator.cpp @@ -35,13 +35,15 @@ limitations under the License. #include "util/net.h" namespace xllm { -DiTCollectiveCommunicator::DiTCollectiveCommunicator(int32_t global_rank, - int32_t world_size, - int32_t dit_dp_size, - int32_t dit_tp_size, - int32_t dit_sp_size, - int32_t dit_cfg_size, - int32_t dit_vae_size) +DiTCollectiveCommunicator::DiTCollectiveCommunicator( + int32_t global_rank, + int32_t world_size, + int32_t dit_dp_size, + int32_t dit_tp_size, + int32_t dit_sp_size, + int32_t dit_cfg_size, + int32_t dit_vae_size, + int32_t dit_text_encoder_tp_size) : CollectiveCommunicatorBase(global_rank, world_size) { parallel_args_ = std::make_unique(global_rank, world_size, @@ -50,13 +52,15 @@ DiTCollectiveCommunicator::DiTCollectiveCommunicator(int32_t global_rank, dit_sp_size, dit_cfg_size, dit_vae_size, + dit_text_encoder_tp_size, /*process_group=*/nullptr); DiTMapping::Options dit_mapping_options; dit_mapping_options.dit_tp_size(dit_tp_size) .dit_sp_size(dit_sp_size) .dit_cfg_size(dit_cfg_size) .dit_dp_size(dit_dp_size) - .dit_vae_size(dit_vae_size); + .dit_vae_size(dit_vae_size) + .dit_text_encoder_tp_size(dit_text_encoder_tp_size); dit_mapping_ = std::make_unique( world_size, global_rank, dit_mapping_options); } @@ -88,6 +92,8 @@ void DiTCollectiveCommunicator::create_process_groups( create_process_group_by_type("dp", dit_dp_group_, device); parallel_args_->dit_vae_group_ = create_process_group_by_type("vae", dit_vae_group_, device); + parallel_args_->dit_text_encoder_tp_group_ = create_process_group_by_type( + "text_encoder_tp", dit_text_encoder_tp_group_, device); } ProcessGroup* DiTCollectiveCommunicator::create_process_group_by_type( @@ -95,7 +101,7 @@ ProcessGroup* DiTCollectiveCommunicator::create_process_group_by_type( std::unique_ptr& member_group, const torch::Device& device) { int32_t group_size = parallel_args_->get_group_size_by_type(group_type); - if (group_size > 1 && dit_mapping_) { + if (group_size >= 1 && dit_mapping_) { auto parallel_info = dit_mapping_->get_parallel_info(group_type); auto group_id = parallel_info.current_group_id(); auto num_group = parallel_info.num_group(); diff --git a/xllm/core/framework/parallel_state/dit_collective_communicator.h b/xllm/core/framework/parallel_state/dit_collective_communicator.h index 83672c7b2d..2a84721f2a 100644 --- a/xllm/core/framework/parallel_state/dit_collective_communicator.h +++ b/xllm/core/framework/parallel_state/dit_collective_communicator.h @@ -28,7 +28,8 @@ class DiTCollectiveCommunicator : public CollectiveCommunicatorBase { int32_t dit_tp_size, int32_t dit_sp_size, int32_t dit_cfg_size, - int32_t dit_vae_size); + int32_t dit_vae_size, + int32_t dit_text_encoder_tp_size); ~DiTCollectiveCommunicator() = default; @@ -54,6 +55,7 @@ class DiTCollectiveCommunicator : public CollectiveCommunicatorBase { std::unique_ptr dit_dp_group_; std::unique_ptr dit_cfg_group_; std::unique_ptr dit_vae_group_; + std::unique_ptr dit_text_encoder_tp_group_; }; } // namespace xllm diff --git a/xllm/core/framework/parallel_state/dit_mapping.cpp b/xllm/core/framework/parallel_state/dit_mapping.cpp index eec1b0b97a..f81571a879 100644 --- a/xllm/core/framework/parallel_state/dit_mapping.cpp +++ b/xllm/core/framework/parallel_state/dit_mapping.cpp @@ -28,6 +28,7 @@ DiTMapping::DiTMapping(const int32_t world_size, cfg_.backend("hccl"); dp_.backend("hccl"); vae_.backend("hccl"); + text_encoder_tp_.backend("hccl"); parse_parallel_info(); validate(); RankGenerator rank_generator(world_size); @@ -47,6 +48,15 @@ DiTMapping::DiTMapping(const int32_t world_size, auto ranks_mapping_vae = rank_generator.get_ranks_mapping(vae_group_ranks, vae_group_order); set_group_by_type(vae_, "vae", ranks_mapping_vae.at("vae")); + + std::vector text_encoder_tp_group_ranks = { + text_encoder_tp_.group_size()}; + std::vector text_encoder_tp_group_order = {"text_encoder_tp"}; + auto ranks_mapping_text_encoder_tp = rank_generator.get_ranks_mapping( + text_encoder_tp_group_ranks, text_encoder_tp_group_order); + set_group_by_type(text_encoder_tp_, + "text_encoder_tp", + ranks_mapping_text_encoder_tp.at("text_encoder_tp")); } void DiTMapping::parse_parallel_info() { @@ -65,6 +75,9 @@ void DiTMapping::parse_parallel_info() { if (options_.dit_vae_size() != -1) { vae_.group_size(options_.dit_vae_size()); } + if (options_.dit_text_encoder_tp_size() != -1) { + text_encoder_tp_.group_size(options_.dit_text_encoder_tp_size()); + } } void DiTMapping::validate() { @@ -107,6 +120,21 @@ void DiTMapping::validate() { std::to_string(vae_.group_size()) + ", world_size is " + std::to_string(world_size_) + ". Please check `vae` and 'world_size'."; + + CHECK(text_encoder_tp_.group_size() >= 1 && + text_encoder_tp_.group_size() <= world_size_) + << "text_encoder_tp_size must be between 1 and world_size. " + "text_encoder_tp_size is " + + std::to_string(text_encoder_tp_.group_size()) + + ", world_size is " + std::to_string(world_size_) + + ". Please check `text_encoder_tp_size` and 'world_size'."; + + CHECK(world_size_ % text_encoder_tp_.group_size() == 0) + << "world_size could not be divided by text_encoder_tp_size. " + "text_encoder_tp_size is " + + std::to_string(text_encoder_tp_.group_size()) + + ", world_size is " + std::to_string(world_size_) + + ". Please check `text_encoder_tp_size` and 'world_size'."; } void DiTMapping::set_group_by_type( @@ -150,6 +178,8 @@ const ParallelInfo& DiTMapping::get_parallel_info( return dp_; } else if (group_type == "vae") { return vae_; + } else if (group_type == "text_encoder_tp") { + return text_encoder_tp_; } else { LOG(FATAL) << "get unexpected group_type: " << group_type; } @@ -161,6 +191,7 @@ nlohmann::json DiTMapping::to_json() { data["SpSize"] = options_.dit_sp_size(); data["TpSize"] = options_.dit_tp_size(); data["CfgSize"] = options_.dit_cfg_size(); + data["TextEncoderTpSize"] = options_.dit_text_encoder_tp_size(); data["worldSize"] = world_size_; data["rank"] = rank_; data["sp"] = sp_.to_json(); @@ -168,6 +199,7 @@ nlohmann::json DiTMapping::to_json() { data["cfg"] = cfg_.to_json(); data["dp"] = dp_.to_json(); data["vae"] = vae_.to_json(); + data["text_encoder_tp"] = text_encoder_tp_.to_json(); return data; } diff --git a/xllm/core/framework/parallel_state/dit_mapping.h b/xllm/core/framework/parallel_state/dit_mapping.h index 536e6cd469..b227e14574 100644 --- a/xllm/core/framework/parallel_state/dit_mapping.h +++ b/xllm/core/framework/parallel_state/dit_mapping.h @@ -36,6 +36,8 @@ class DiTMapping final { PROPERTY(int32_t, dit_dp_size) = -1; // vae size PROPERTY(int32_t, dit_vae_size) = -1; + // text encoder tensor parallel size + PROPERTY(int32_t, dit_text_encoder_tp_size) = -1; }; DiTMapping(const int32_t world_size, @@ -72,5 +74,6 @@ class DiTMapping final { ParallelInfo cfg_ = ParallelInfo(); ParallelInfo dp_ = ParallelInfo(); ParallelInfo vae_ = ParallelInfo(); + ParallelInfo text_encoder_tp_ = ParallelInfo(); }; } // namespace xllm diff --git a/xllm/core/framework/parallel_state/parallel_args.h b/xllm/core/framework/parallel_state/parallel_args.h index e2e3b0edc4..845929065d 100644 --- a/xllm/core/framework/parallel_state/parallel_args.h +++ b/xllm/core/framework/parallel_state/parallel_args.h @@ -101,6 +101,7 @@ struct ParallelArgs { int32_t sp_size, int32_t cfg_size, int32_t vae_size, + int32_t text_encoder_tp_size, ProcessGroup* process_group) : rank_(rank), world_size_(world_size), @@ -109,6 +110,7 @@ struct ParallelArgs { sp_size_(sp_size), cfg_size_(cfg_size), vae_size_(vae_size), + text_encoder_tp_size_(text_encoder_tp_size), process_group_(process_group) {} int32_t get_group_size_by_type(const std::string& group_type) const { @@ -124,6 +126,8 @@ struct ParallelArgs { return ep_size(); } else if (group_type == "vae") { return vae_size(); + } else if (group_type == "text_encoder_tp") { + return text_encoder_tp_size(); } else if (group_type == "cp") { return cp_size(); } else { @@ -185,6 +189,9 @@ struct ParallelArgs { // cfg size PROPERTY(int32_t, vae_size) = 1; + // text encoder tensor parallel size + PROPERTY(int32_t, text_encoder_tp_size) = 1; + // atb hccl mapping json data PROPERTY(nlohmann::json, mapping_data); @@ -225,6 +232,7 @@ struct ParallelArgs { ProcessGroup* dit_cfg_group_ = nullptr; ProcessGroup* dit_dp_group_ = nullptr; ProcessGroup* dit_vae_group_ = nullptr; + ProcessGroup* dit_text_encoder_tp_group_ = nullptr; }; } // namespace xllm diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 313ff262b6..ea6bb2f460 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" @@ -292,7 +293,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]; @@ -302,7 +302,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; }; } @@ -312,12 +314,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" @@ -353,7 +353,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, @@ -367,13 +373,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; }; } @@ -381,6 +393,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 b40ddfa34d..9953e30b47 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" @@ -78,11 +80,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 b9a40f2c4f..032d10536d 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -51,6 +51,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, @@ -65,7 +70,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 59ddaa0c37..a6b993f71d 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -113,6 +113,7 @@ torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) { return down_proj_->forward(gate_up, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } + return down_proj_->forward(gate_up); } diff --git a/xllm/core/layers/common/qwen2_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/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index df42b89b44..438ad4aad2 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -49,21 +49,29 @@ std::tuple> AttentionImpl::forward( return std::make_tuple(output, output_lse); } - bool only_prefill = + const bool only_prefill = attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; - torch::Tensor k_cache = kv_cache.get_k_cache(); - torch::Tensor v = value.view({-1, num_kv_heads_, head_size_}); - std::optional v_cache = kv_cache.get_v_cache(); - - // Reshape and cache key/value - xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; - reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_}); - reshape_paged_cache_params.value = v; - reshape_paged_cache_params.k_cache = k_cache; - reshape_paged_cache_params.v_cache = v_cache; - reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; - xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); + torch::Tensor k_cache; + std::optional v_cache; + if (kv_cache.empty()) { + CHECK(attn_metadata.is_prefill) + << "KV-cache-free NPU attention only supports full prefill."; + CHECK(!attn_metadata.use_expanded_decode_for_spec_verify_attention) + << "KV-cache-free NPU attention does not support expanded decode."; + } else { + k_cache = kv_cache.get_k_cache(); + v_cache = kv_cache.get_v_cache(); + + xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; + reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_}); + reshape_paged_cache_params.value = + value.view({-1, num_kv_heads_, head_size_}); + reshape_paged_cache_params.k_cache = k_cache; + reshape_paged_cache_params.v_cache = v_cache; + reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; + xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); + } if (attn_metadata.use_expanded_decode_for_spec_verify_attention) { decoder_forward(query, output, k_cache, v_cache, attn_metadata); 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/options.h b/xllm/core/runtime/options.h index 0d89b738b2..70ee0deb50 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -135,6 +135,10 @@ struct Options { // Default set as 1 PROPERTY(int32_t, vae_size) = 1; + // text encoder tensor parallelism size + // Default set as 1 + PROPERTY(int32_t, text_encoder_tp_size) = 1; + // enable enable_schedule_overlap to improve runtime execution efficiency. PROPERTY(bool, enable_schedule_overlap) = true; diff --git a/xllm/models/dit/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..04d5cdabf6 --- /dev/null +++ b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h @@ -0,0 +1,955 @@ +/* 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 + +#include "core/framework/config/kernel_config.h" +#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/kv_cache/kv_cache.h" +#include "core/framework/model/model_input_params.h" +#include "core/framework/model_context.h" +#include "core/framework/multimodal/mm_input.h" +#include "core/framework/multimodal/mm_visitor.h" +#include "core/framework/parallel_state/parallel_state.h" +#include "core/framework/parallel_state/process_group.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/framework/tokenizer/tokenizer.h" +#include "core/runtime/dit_forward_params.h" +#include "models/dit/autoencoders/autoencoder_kl_wan.h" +#include "models/dit/processors/vae_image_processor.h" +#include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h" +#include "models/dit/transformers/transformer_joyimage_edit_plus.h" +#include "models/dit/utils/util.h" +#include "models/model_registry.h" +#include "models/vlm/mposition/mposition.h" +#include "models/vlm/qwen3_vl.h" +#include "processors/multimodal_processor.h" +#include "util/tensor_helper.h" + +namespace xllm { + +using JoyImageEditTextEncoder = Qwen3_VLForConditionalGeneration; + +class JoyImageEditPlusPipelineImpl : public torch::nn::Module { + public: + 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")); + + if (context.has_component("text_encoder")) { + text_encoder_model_args_ = context.get_model_args("text_encoder"); +#if defined(USE_NPU) + CHECK_EQ(KernelConfig::get_instance().npu_kernel_backend(), "TORCH") + << "JoyImageEditPlus in-process Qwen3-VL text encoding requires " + "--npu_kernel_backend=TORCH."; +#endif + ProcessGroup* tp_group = parallel_args_.dit_text_encoder_tp_group_; + CHECK(tp_group != nullptr) + << "JoyImageEditPlus requires an injected text encoder TP group."; + ParallelArgs vlm_parallel_args( + tp_group->rank(), tp_group->world_size(), tp_group); + vlm_parallel_args.tp_size(tp_group->world_size()); + vlm_parallel_args.tp_group_ = tp_group; + text_encoder_empty_kv_caches_.resize( + static_cast(text_encoder_model_args_.n_layers())); + const ModelContext& source_vlm_context = + context.get_model_context("text_encoder"); + ModelContext vlm_context(vlm_parallel_args, + source_vlm_context.get_model_args(), + source_vlm_context.get_quant_args(), + source_vlm_context.get_tensor_options()); + text_encoder_ = JoyImageEditTextEncoder(vlm_context); + CHECK(!text_encoder_.is_empty()) + << "Failed to create JoyImageEditPlus Qwen3-VL text encoder"; + mposition_generator_ = + MPositionGeneratorFactory::get_instance().create_mposition_generator( + "qwen3_vl"); + } + + vae_image_processor_ = + xllm::VAEImageProcessor(context.get_model_context("vae"), + /*do_resize=*/true, + /*do_normalize=*/true, + /*do_binarize=*/false, + /*do_convert_rgb=*/false, + /*do_convert_grayscale=*/false, + latent_channels_, + /*scale_factor=*/vae_scale_factor_spatial_); + + register_module("vae", vae_); + register_module("scheduler", scheduler_); + register_module("transformer", transformer_); + if (!text_encoder_.is_empty()) { + register_module("text_encoder", text_encoder_); + } + register_module("vae_image_processor", vae_image_processor_); + + latents_mean_ = vae_model_args_.latents_mean(); + latents_std_ = vae_model_args_.latents_std(); + } + + torch::Tensor latents_mean_tensor(const torch::Tensor& ref) const { + return torch::tensor(latents_mean_, torch::kFloat32) + .view({1, latent_channels_, 1, 1, 1}) + .to(ref.device(), ref.dtype()); + } + torch::Tensor latents_std_tensor(const torch::Tensor& ref) const { + return torch::tensor(latents_std_, torch::kFloat32) + .view({1, latent_channels_, 1, 1, 1}) + .to(ref.device(), ref.dtype()); + } + + // Patchify a [C, T, H, W] latent into [num_patches, C, pt, ph, pw] and + // return {patches, (lt, lh, lw)}. + std::pair> 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")); + } + } + bool do_cfg = guidance_scale > 1.0; + torch::Tensor prompt_embeds; + torch::Tensor prompt_embeds_mask; + const bool encode_text = + !input.prompt_embeds.defined() || + (do_cfg && !input.negative_prompt_embeds.defined()); + if (encode_text) { + CHECK(!text_encoder_.is_empty()) << "Qwen3-VL text encoder is not loaded"; + } + if (input.prompt_embeds.defined()) { + prompt_embeds = input.prompt_embeds.to(options_.device(), dtype_); + prompt_embeds_mask = torch::ones( + {prompt_embeds.size(0), prompt_embeds.size(1)}, + torch::TensorOptions().device(device_).dtype(torch::kLong)); + } else { + CHECK(!input.prompts.empty()) + << "JoyImageEditPlus requires `prompts` or `prompt_embeds`"; + std::tie(prompt_embeds, prompt_embeds_mask) = + encode_prompts(input.prompts, raw_images, batch_size); + } + + torch::Tensor neg_embeds, neg_embeds_mask; + if (do_cfg) { + if (input.negative_prompt_embeds.defined()) { + neg_embeds = input.negative_prompt_embeds.to(options_.device(), dtype_); + neg_embeds_mask = torch::ones( + {neg_embeds.size(0), neg_embeds.size(1)}, + torch::TensorOptions().device(device_).dtype(torch::kLong)); + } else { + std::vector negative_prompts = input.negative_prompts; + if (negative_prompts.empty()) { + negative_prompts.resize(static_cast(batch_size)); + } + std::tie(neg_embeds, neg_embeds_mask) = + encode_prompts(negative_prompts, raw_images, batch_size); + } + // Pad/concat [negative, positive] to equal sequence length. + int64_t max_l = std::max(prompt_embeds.size(1), neg_embeds.size(1)); + prompt_embeds = pad_seq(prompt_embeds, max_l); + neg_embeds = pad_seq(neg_embeds, max_l); + prompt_embeds_mask = pad_seq(prompt_embeds_mask, max_l); + neg_embeds_mask = pad_seq(neg_embeds_mask, max_l); + } + // Latents. + int64_t num_channels_latents = in_channels_; + auto lp = prepare_latents(batch_size, + num_channels_latents, + height, + width, + seed, + vae_refs, + input.latents); + auto latents = std::get<0>(lp); + auto target_mask = std::get<1>(lp); + auto shape_list = std::get<2>(lp); + auto clean_backup = latents.clone(); + + // 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(); + auto transformer_loader = loader->take_component_loader("transformer"); + auto vae_loader = loader->take_component_loader("vae"); + std::unique_ptr text_encoder_loader; + if (!text_encoder_.is_empty()) { + text_encoder_loader = loader->take_component_loader("text_encoder"); + std::unique_ptr tokenizer = text_encoder_loader->tokenizer(); + CHECK(tokenizer != nullptr) + << "Failed to load JoyImageEditPlus Qwen3-VL tokenizer"; + tokenizer_ = std::shared_ptr(std::move(tokenizer)); + multimodal_processor_ = + create_multimodal_processor(text_encoder_model_args_, tokenizer_); + } + + vae_->load_model(std::move(vae_loader)); + vae_->to(options_.device(), dtype_); + + transformer_->load_model(std::move(transformer_loader)); + transformer_->to(options_.device(), dtype_); + transformer_->keep_fp32_modules(); + + if (text_encoder_loader != nullptr) { + text_encoder_->load_model(std::move(text_encoder_loader)); + } + + LOG(INFO) << "JoyImageEditPlusPipeline loaded."; + } + + private: + static constexpr int64_t kPromptEmbeddingStartIndex = 34; + + std::string build_qwen3_vl_prompt(const std::string& prompt, + size_t image_count) const { + std::string chat_prompt = + "<|im_start|>system\n \\nDescribe the image by detailing the " + "color, shape, size, texture, quantity, text, spatial relationships " + "of the objects and background:<|im_end|>\n<|im_start|>user\n"; + for (size_t image_index = 0; image_index < image_count; ++image_index) { + chat_prompt += "<|vision_start|><|image_pad|><|vision_end|>"; + } + chat_prompt += prompt; + chat_prompt += "<|im_end|>\n<|im_start|>assistant\n"; + return chat_prompt; + } + + std::vector prepare_vl_images( + const std::vector& raw_images, + int64_t batch_index, + int64_t batch_size) const { + std::vector images; + images.reserve(raw_images.size()); + for (const torch::Tensor& raw_image : raw_images) { + CHECK(raw_image.dim() == 3 || raw_image.dim() == 4) + << "JoyImageEditPlus image input must be [C,H,W] or [B,C,H,W]"; + torch::Tensor image; + if (raw_image.dim() == 4) { + CHECK_EQ(raw_image.size(0), batch_size) + << "JoyImageEditPlus image batch must match prompt batch"; + image = raw_image[batch_index]; + } else { + CHECK_EQ(batch_size, 1) + << "Unbatched JoyImageEditPlus images require batch_size=1"; + image = raw_image; + } + image = image.to(torch::kCPU).to(torch::kFloat32); + const float max_value = image.max().item(); + image = max_value <= 1.1f ? image.clamp(0.0f, 1.0f) * 255.0f + : image.clamp(0.0f, 255.0f); + images.emplace_back(std::move(image)); + } + return images; + } + + ModelInputParams build_text_encoder_input(const torch::Tensor& tokens, + const MMData& mm_data) { + CHECK_LE(tokens.numel(), std::numeric_limits::max()) + << "JoyImageEditPlus Qwen3-VL prompt is too long"; + const int32_t sequence_length = static_cast(tokens.numel()); + CHECK_GT(sequence_length, 0) + << "JoyImageEditPlus Qwen3-VL prompt must not be empty"; + + ModelInputParams params; + params.meta.num_sequences = 1; + params.meta.actual_num_sequences = 1; + params.meta.q_max_seq_len = sequence_length; + params.meta.kv_max_seq_len = sequence_length; + params.meta.batch_forward_type = BatchForwardType::PREFILL; + params.attention.host.q_seq_lens = {sequence_length}; + params.attention.host.kv_seq_lens = {sequence_length}; +#if defined(USE_NPU) + params.attention.host.q_cu_seq_lens = {sequence_length}; +#else + params.attention.host.q_cu_seq_lens = {0, sequence_length}; +#endif + params.attention.device.q_seq_lens = + torch::tensor({sequence_length}, torch::kInt).to(tokens.device()); + params.attention.device.kv_seq_lens = + torch::tensor({sequence_length}, torch::kInt).to(tokens.device()); +#if defined(USE_NPU) + params.attention.device.q_cu_seq_lens = + torch::tensor({sequence_length}, torch::kInt).to(tokens.device()); +#else + params.attention.device.q_cu_seq_lens = + torch::tensor({0, sequence_length}, torch::kInt).to(tokens.device()); +#endif + + MMBatchData mm_batch(std::vector{mm_data}); + EncoderInputGatherVisitor input_gather; + mm_batch.foreach (input_gather); + CHECK(input_gather.finish(mm_batch)) + << "JoyImageEditPlus failed to gather Qwen3-VL encoder inputs"; + mm_batch.to(tokens.device()); + + ModelInputParams multimodal_params; + multimodal_params.multimodal.mm_data = mm_batch; + MMDict multimodal_embeddings = + text_encoder_->get_multimodal_embeddings(multimodal_params); + EncoderOutputScatterVisitor output_scatter(multimodal_embeddings); + CHECK(mm_batch.foreach (output_scatter)); + CHECK(output_scatter.finish()) + << "JoyImageEditPlus failed to scatter Qwen3-VL embeddings"; + + EncoderEmbeddingGatherVisitor embedding_gather( + tokens.device(), + mm_batch.type(), + params.attention.host.kv_seq_lens, + params.attention.host.q_seq_lens); + CHECK(mm_batch.foreach (embedding_gather)); + CHECK(embedding_gather.finish(mm_batch)) + << "JoyImageEditPlus failed to gather Qwen3-VL embeddings"; + params.multimodal.mm_data = std::move(mm_batch); + params.embedding.input_embedding = + text_encoder_->get_input_embeddings(tokens, params); + return params; + } + + std::pair encode_single_prompt( + const std::string& prompt, + const std::vector& raw_images, + int64_t batch_index, + int64_t batch_size) { + CHECK(!text_encoder_.is_empty()) << "Qwen3-VL text encoder is not loaded"; + CHECK(tokenizer_ != nullptr) << "Qwen3-VL tokenizer is not loaded"; + CHECK(multimodal_processor_ != nullptr) + << "Qwen3-VL multimodal processor is not loaded"; + CHECK(mposition_generator_ != nullptr) + << "Qwen3-VL position generator is not loaded"; + + std::vector images = + prepare_vl_images(raw_images, batch_index, batch_size); + std::vector input_items; + input_items.reserve(images.size()); + for (torch::Tensor& image : images) { + MMInputItem item; + item.type = MMType::IMAGE; + item.decode_image = std::move(image); + input_items.emplace_back(std::move(item)); + } + MMInput multimodal_input; + multimodal_input.insert(input_items); + MMData mm_data; + CHECK(multimodal_processor_->process_multimodal(multimodal_input, mm_data)) + << "JoyImageEditPlus Qwen3-VL image preprocessing failed"; + + std::string chat_prompt = build_qwen3_vl_prompt(prompt, images.size()); + std::vector token_ids; + CHECK( + multimodal_processor_->process_prompt(chat_prompt, mm_data, token_ids)) + << "JoyImageEditPlus Qwen3-VL prompt processing failed"; + UpdateMMItemScheduleStateVisitor schedule_visitor( + /*computed_token_num=*/0, + static_cast(token_ids.size()), + /*seq_idx=*/0); + CHECK(mm_data.foreach (schedule_visitor)) + << "JoyImageEditPlus failed to schedule Qwen3-VL multimodal inputs"; + auto [positions, mrope_position_delta] = mposition_generator_->generate( + token_ids, mm_data, text_encoder_model_args_); + static_cast(mrope_position_delta); + + torch::Tensor tokens = + torch::tensor(token_ids, torch::TensorOptions().dtype(torch::kInt32)) + .to(device_); + positions = positions.to(device_); + ModelInputParams input_params = build_text_encoder_input(tokens, mm_data); + ModelOutput model_output = text_encoder_->forward( + tokens, positions, text_encoder_empty_kv_caches_, input_params); + CHECK(model_output.residual.defined()) + << "JoyImageEditPlus requires Qwen3-VL pre-norm hidden states from " + "the TORCH backend."; + CHECK_GT(model_output.residual.size(0), kPromptEmbeddingStartIndex) + << "JoyImageEditPlus Qwen3-VL prompt is shorter than the embedding " + "prefix"; + + torch::Tensor prompt_embeddings = + model_output.residual + .slice( + /*dim=*/0, kPromptEmbeddingStartIndex) + .unsqueeze(0) + .to(options_); + torch::Tensor prompt_mask = + torch::ones({1, prompt_embeddings.size(1)}, + torch::TensorOptions().device(device_).dtype(torch::kLong)); + return {prompt_embeddings, prompt_mask}; + } + + std::pair encode_prompts( + const std::vector& prompts, + const std::vector& raw_images, + int64_t batch_size) { + CHECK_EQ(static_cast(prompts.size()), batch_size) + << "JoyImageEditPlus prompt batch size mismatch"; + std::vector embeddings; + std::vector masks; + embeddings.reserve(static_cast(batch_size)); + masks.reserve(static_cast(batch_size)); + int64_t max_sequence_length = 0; + for (int64_t batch_index = 0; batch_index < batch_size; ++batch_index) { + auto [embedding, mask] = + encode_single_prompt(prompts[static_cast(batch_index)], + raw_images, + batch_index, + batch_size); + max_sequence_length = std::max(max_sequence_length, embedding.size(1)); + embeddings.emplace_back(std::move(embedding)); + masks.emplace_back(std::move(mask)); + } + for (int64_t batch_index = 0; batch_index < batch_size; ++batch_index) { + embeddings[static_cast(batch_index)] = pad_seq( + embeddings[static_cast(batch_index)], max_sequence_length); + masks[static_cast(batch_index)] = + pad_seq(masks[static_cast(batch_index)], max_sequence_length); + } + return {torch::cat(embeddings, /*dim=*/0), torch::cat(masks, /*dim=*/0)}; + } + + 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}; + JoyImageEditTextEncoder text_encoder_{nullptr}; + std::shared_ptr tokenizer_; + std::unique_ptr multimodal_processor_; + std::unique_ptr mposition_generator_; + ModelArgs text_encoder_model_args_; + std::vector text_encoder_empty_kv_caches_; + + int64_t in_channels_; + int64_t num_layers_; + int64_t head_dim_; + int64_t theta_; + int64_t patch_t_, patch_h_, patch_w_; + int64_t latent_channels_; + int64_t vae_scale_factor_spatial_; + std::vector rope_dim_list_; + std::vector latents_mean_; + std::vector latents_std_; +}; + +TORCH_MODULE(JoyImageEditPlusPipeline); + +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 b14e9c5817..e25eaf9dc0 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/model_registry.cpp b/xllm/models/model_registry.cpp index d5fe6e2b58..50bd2d9cd1 100644 --- a/xllm/models/model_registry.cpp +++ b/xllm/models/model_registry.cpp @@ -116,8 +116,9 @@ bool resolve_model_registration(const std::string& model_type, if (backend == kAutoBackend) { effective_backend = is_torch_only_model_type(model_type) ? kTorchBackend : kAtbBackend; - } else if (model_type == "qwen3" || model_type == "qwen3_moe") { - // qwen3/qwen3_moe support both backends. + } else if (model_type == "qwen3" || model_type == "qwen3_moe" || + model_type == "qwen3_vl") { + // qwen3/qwen3_moe/qwen3_vl support both backends. } else if (is_torch_only_model_type(model_type)) { if (backend != kTorchBackend) { if (error_message != nullptr) { @@ -141,6 +142,8 @@ bool resolve_model_registration(const std::string& model_type, *resolved_name = "qwen3_atb"; } else if (model_type == "qwen3_moe" && effective_backend == kAtbBackend) { *resolved_name = "qwen3_moe_atb"; + } else if (model_type == "qwen3_vl" && effective_backend == kAtbBackend) { + *resolved_name = "qwen3_vl_atb"; } else { *resolved_name = model_type; } diff --git a/xllm/models/models.h b/xllm/models/models.h index 646e0343d8..6a99242a8f 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -20,6 +20,7 @@ limitations under the License. #include "dit/pipelines/pipeline_flux2.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_flux_control.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_flux_fill.h" // IWYU pragma: keep +#include "dit/pipelines/pipeline_joyimage_edit_plus.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_qwenimage_edit_plus.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_wan_i2v.h" // IWYU pragma: keep #include "llm/deepseek_v4.h" // IWYU pragma: keep @@ -62,6 +63,7 @@ limitations under the License. #include "vlm/npu/qwen3_vl.h" // IWYU pragma: keep #include "vlm/npu/qwen3_vl_moe.h" // IWYU pragma: keep #include "vlm/qwen3_5.h" // IWYU pragma: keep +#include "vlm/qwen3_vl.h" // IWYU pragma: keep #elif defined(USE_MLU) #include "dit/pipelines/pipeline_flux.h" // IWYU pragma: keep diff --git a/xllm/models/vlm/npu/qwen3_vl.h b/xllm/models/vlm/npu/qwen3_vl.h index 3c19ee7c51..a591f25bb8 100644 --- a/xllm/models/vlm/npu/qwen3_vl.h +++ b/xllm/models/vlm/npu/qwen3_vl.h @@ -841,11 +841,15 @@ TORCH_MODULE(Qwen3_VLForConditionalGeneration); using Qwen3VLMultimodalProcessor = MultimodalProcessor; -REGISTER_MULTIMODAL_PROCESSOR(qwen3_vl, Qwen3VLMultimodalProcessor); -REGISTER_CAUSAL_VLM_MODEL(qwen3_vl, Qwen3_VLForConditionalGeneration); +REGISTER_MULTIMODAL_PROCESSOR_WITH_VARNAME(qwen3_vl_atb, + qwen3_vl_atb, + Qwen3VLMultimodalProcessor); +REGISTER_CAUSAL_VLM_MODEL_WITH_VARNAME(qwen3_vl_atb, + qwen3_vl_atb, + Qwen3_VLForConditionalGeneration); REGISTER_MPOSITION_GENERATOR(qwen3_vl, xllm::Qwen3VLMPositionGenerator); -REGISTER_MODEL_ARGS(qwen3_vl, [&] { +REGISTER_MODEL_ARGS_WITH_VARNAME(qwen3_vl_atb, qwen3_vl_atb, [&] { // text config // LOAD_ARG_OR(attention_dropout, "attention_dropout", 0.0); LOAD_ARG_OR(model_type, "model_type", "qwen3_vl"); diff --git a/xllm/models/vlm/qwen3_vl.h b/xllm/models/vlm/qwen3_vl.h index 7350ae3cf9..bb86f84b20 100644 --- a/xllm/models/vlm/qwen3_vl.h +++ b/xllm/models/vlm/qwen3_vl.h @@ -461,12 +461,12 @@ class Qwen3_VisionTransformerImpl : public torch::nn::Module { m_sin = m_sin.repeat({1, 2}); torch::Tensor cu_seqlens_cpu = cu_seqlens.cpu(); - std::vector 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,11 +574,13 @@ 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); -REGISTER_MODEL_ARGS(qwen3_vl, [&] { +const ModelArgsLoader kQwen3VLModelArgsLoader = [](const JsonReader& json, + ModelArgs* args) { // text config // LOAD_ARG_OR(attention_dropout, "attention_dropout", 0.0); LOAD_ARG_OR(model_type, "model_type", "qwen3_vl"); @@ -617,7 +636,11 @@ REGISTER_MODEL_ARGS(qwen3_vl, [&] { LOAD_ARG_OR( rope_scaling_rope_type, "vision_config.rope_scaling.type", "mrope"); - LOAD_ARG_OR(vocab_size, "text_config.vocab_size", 151936); -}); + return true; +}; + +REGISTER_MODEL_ARGS_LOADER(qwen3_vl, kQwen3VLModelArgsLoader); +REGISTER_MODEL_ARGS_LOADER(Qwen3VLForConditionalGeneration, + kQwen3VLModelArgsLoader); } // namespace xllm diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 223888524b..f8e017096f 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -182,6 +182,8 @@ Options create_options(const std::string& instance_name, bool is_local) { .sp_size(static_cast(parallel_config.sp_size())) .cfg_size(static_cast(parallel_config.cfg_size())) .vae_size(static_cast(parallel_config.vae_size())) + .text_encoder_tp_size( + static_cast(parallel_config.text_encoder_tp_size())) .instance_name(instance_name) .enable_disagg_pd(disagg_pd_config.enable_disagg_pd()) .enable_pd_ooc(disagg_pd_config.enable_pd_ooc()) From 0bfc4d770e05e6a06381701d459c97319b3e2fe3 Mon Sep 17 00:00:00 2001 From: shan-chen-feng Date: Tue, 28 Jul 2026 14:14:50 +0800 Subject: [PATCH 2/4] feat: support cfg mixin for joy-image-edit-plus. --- xllm/core/layers/common/dense_mlp.cpp | 2 +- .../pipelines/pipeline_joyimage_edit_plus.h | 145 +++++------- .../transformer_joyimage_edit_plus.h | 213 +++++++++++++----- 3 files changed, 221 insertions(+), 139 deletions(-) diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index a6b993f71d..2c394a65c1 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -113,7 +113,7 @@ torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) { return down_proj_->forward(gate_up, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } - + return down_proj_->forward(gate_up); } diff --git a/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h index 04d5cdabf6..92c7b9a215 100644 --- a/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h +++ b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h @@ -27,7 +27,6 @@ limitations under the License. #include #include "core/framework/config/kernel_config.h" -#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/kv_cache/kv_cache.h" @@ -35,7 +34,6 @@ limitations under the License. #include "core/framework/model_context.h" #include "core/framework/multimodal/mm_input.h" #include "core/framework/multimodal/mm_visitor.h" -#include "core/framework/parallel_state/parallel_state.h" #include "core/framework/parallel_state/process_group.h" #include "core/framework/state_dict/state_dict.h" #include "core/framework/tokenizer/tokenizer.h" @@ -44,6 +42,7 @@ limitations under the License. #include "models/dit/processors/vae_image_processor.h" #include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h" #include "models/dit/transformers/transformer_joyimage_edit_plus.h" +#include "models/dit/utils/dit_parallel_mixin.h" #include "models/dit/utils/util.h" #include "models/model_registry.h" #include "models/vlm/mposition/mposition.h" @@ -55,7 +54,8 @@ namespace xllm { using JoyImageEditTextEncoder = Qwen3_VLForConditionalGeneration; -class JoyImageEditPlusPipelineImpl : public torch::nn::Module { +class JoyImageEditPlusPipelineImpl : public torch::nn::Module, + public dit::CFGParallelMixin { public: using JoyImageShapeList = std::vector>>; @@ -66,7 +66,8 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { }; JoyImageEditPlusPipelineImpl(const DiTModelContext& context) - : context_(context), + : dit::CFGParallelMixin(context), + context_(context), parallel_args_(context.get_parallel_args()), vae_model_args_(context.get_model_args("vae")) { options_ = context.get_tensor_options(); @@ -314,7 +315,7 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { 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")); + img4, hw.first, hw.second, /*resize_mode=*/"lanczos")); } } bool do_cfg = guidance_scale > 1.0; @@ -373,6 +374,9 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { auto target_mask = std::get<1>(lp); auto shape_list = std::get<2>(lp); auto clean_backup = latents.clone(); + const std::array& target_shape = shape_list.front().front(); + const int64_t target_patch_count = + target_shape[0] * target_shape[1] * target_shape[2]; // Timesteps (static shift; no dynamic shifting for Joy). scheduler_->set_timesteps(num_inference_steps, device_); @@ -381,83 +385,41 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { 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, + TransformerForwardContext prompt_transformer_context = + prepare_transformer_context(batch_size, latents.size(1), latents.device(), - transformer_encoder_hidden_states_mask, - transformer_shape_list); + prompt_embeds_mask, + shape_list); + TransformerForwardContext negative_transformer_context; + if (do_cfg) { + negative_transformer_context = + prepare_transformer_context(batch_size, + latents.size(1), + latents.device(), + neg_embeds_mask, + shape_list); + } for (int64_t i = 0; i < timesteps.size(0); ++i) { auto t = timesteps[i]; - // 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]; - } + torch::Tensor t_expand = t.repeat({batch_size}); + auto [cond, uncond] = exec_with_cfg([&](bool is_positive) { + const torch::Tensor& encoder_hidden_states = + is_positive ? prompt_embeds : neg_embeds; + const TransformerForwardContext& transformer_context = + is_positive ? prompt_transformer_context + : negative_transformer_context; + return sequence_parallel_forward(latents, + t_expand, + encoder_hidden_states, + transformer_context, + /*use_cfg=*/!is_positive, + /*step_index=*/i + 1); + }); auto comb = uncond + guidance_scale * (cond - uncond); // Norm-rescale (diffusers): comb * (||cond|| / ||comb||) over channel // dim (2) of the 6D [B, N, C, pt, ph, pw] prediction. @@ -467,21 +429,22 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { 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); + torch::Tensor t_expand = t.repeat({batch_size}); + noise_pred = sequence_parallel_forward(latents, + t_expand, + prompt_embeds, + prompt_transformer_context, + /*use_cfg=*/false, + /*step_index=*/i + 1); } latents = scheduler_->step(noise_pred, t, latents).to(latents.dtype()); + latents.slice(/*dim=*/1, target_patch_count, latents.size(1)) + .copy_(clean_backup.slice( + /*dim=*/1, target_patch_count, clean_backup.size(1))); } - // Restore refs and decode target patches per sample. - latents.index_put_({~target_mask}, clean_backup.index({~target_mask})); + // Decode target patches per sample. std::vector images; for (int64_t b = 0; b < batch_size; ++b) { auto thw = shape_list[b][0]; @@ -818,7 +781,6 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { 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}, @@ -835,9 +797,20 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module { torch::Tensor full_mask = torch::cat({image_mask, encoder_hidden_states_mask.to(torch::kBool)}, /*dim=*/1); +#if defined(USE_NPU) + const int64_t sequence_length = full_mask.size(1); + attention_mask = full_mask.logical_not() + .unsqueeze(1) + .unsqueeze(1) + .expand({batch_size, + /*num_heads=*/1, + sequence_length, + sequence_length}) + .contiguous(); +#else attention_mask = full_mask.unsqueeze(1).unsqueeze(1); - } #endif + } return {rope_cos, rope_sin, attention_mask}; } diff --git a/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h index 626de0df73..834855c535 100644 --- a/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h +++ b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h @@ -24,6 +24,8 @@ limitations under the License. #include #include #include +#include +#include #include #include "core/framework/dit_cache/dit_cache.h" @@ -32,6 +34,7 @@ limitations under the License. #include "core/framework/parallel_state/parallel_state.h" #include "core/framework/state_dict/state_dict.h" #include "core/framework/state_dict/utils.h" +#include "core/layers/common/ada_layer_norm.h" #include "core/layers/common/add_matmul.h" #include "core/layers/common/rms_norm.h" #include "models/dit/attn_processor/attn_processor.h" @@ -94,9 +97,9 @@ inline torch::Tensor apply_rotary_emb_batched(const torch::Tensor& x, #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. +// Joint attention over [B, S, H, D] with an optional bool mask. NPU uses a +// materialized block mask (True == masked out); the fallback uses a broadcast +// keep mask (True == attend). inline torch::Tensor joyimage_joint_attention( const torch::Tensor& query, // [B, S, H, D] const torch::Tensor& key, // [B, S, H, D] @@ -104,11 +107,6 @@ inline torch::Tensor joyimage_joint_attention( 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, @@ -117,7 +115,7 @@ inline torch::Tensor joyimage_joint_attention( /*input_layout=*/"BSND", /*pse=*/torch::nullopt, /*padding_mask=*/torch::nullopt, - /*atten_mask=*/torch::nullopt, + /*atten_mask=*/attn_mask, /*scale=*/std::pow(static_cast(query.size(3)), -0.5), /*keep_prob=*/1.0, /*pre_tockens=*/65535, @@ -297,12 +295,20 @@ class JoyAttentionImpl final : public torch::nn::Module { 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_)); + : heads_(num_heads), + head_dim_(head_dim), + options_(context.get_tensor_options()) { + const int64_t inner = num_heads * head_dim; + + img_attn_q_ = register_module( + "img_attn_q", + layer::AddMatmulWeightTransposed(dim, inner, true, options_)); + img_attn_k_ = register_module( + "img_attn_k", + layer::AddMatmulWeightTransposed(dim, inner, true, options_)); + img_attn_v_ = register_module( + "img_attn_v", + layer::AddMatmulWeightTransposed(dim, inner, true, options_)); img_attn_q_norm_ = register_module("img_attn_q_norm", layer::RMSNorm(head_dim, eps, options_)); img_attn_k_norm_ = register_module("img_attn_k_norm", @@ -311,9 +317,15 @@ class JoyAttentionImpl final : public torch::nn::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_ = register_module( + "txt_attn_q", + layer::AddMatmulWeightTransposed(dim, inner, true, options_)); + txt_attn_k_ = register_module( + "txt_attn_k", + layer::AddMatmulWeightTransposed(dim, inner, true, options_)); + txt_attn_v_ = register_module( + "txt_attn_v", + layer::AddMatmulWeightTransposed(dim, inner, true, options_)); txt_attn_q_norm_ = register_module("txt_attn_q_norm", layer::RMSNorm(head_dim, eps, options_)); txt_attn_k_norm_ = register_module("txt_attn_k_norm", @@ -324,16 +336,22 @@ class JoyAttentionImpl final : public torch::nn::Module { } void load_state_dict(const StateDict& state_dict) { - img_attn_qkv_->load_state_dict( - state_dict.get_dict_with_prefix("img_attn_qkv.")); + load_qkv_projections(state_dict, + /*prefix=*/"img_attn_qkv", + img_attn_q_, + img_attn_k_, + img_attn_v_); img_attn_q_norm_->load_state_dict( state_dict.get_dict_with_prefix("img_attn_q_norm.")); img_attn_k_norm_->load_state_dict( state_dict.get_dict_with_prefix("img_attn_k_norm.")); img_attn_proj_->load_state_dict( state_dict.get_dict_with_prefix("img_attn_proj.")); - txt_attn_qkv_->load_state_dict( - state_dict.get_dict_with_prefix("txt_attn_qkv.")); + load_qkv_projections(state_dict, + /*prefix=*/"txt_attn_qkv", + txt_attn_q_, + txt_attn_k_, + txt_attn_v_); txt_attn_q_norm_->load_state_dict( state_dict.get_dict_with_prefix("txt_attn_q_norm.")); txt_attn_k_norm_->load_state_dict( @@ -342,11 +360,15 @@ class JoyAttentionImpl final : public torch::nn::Module { 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_->verify_loaded_weights(prefix + "img_attn_qkv.q."); + img_attn_k_->verify_loaded_weights(prefix + "img_attn_qkv.k."); + img_attn_v_->verify_loaded_weights(prefix + "img_attn_qkv.v."); img_attn_q_norm_->verify_loaded_weights(prefix + "img_attn_q_norm."); img_attn_k_norm_->verify_loaded_weights(prefix + "img_attn_k_norm."); img_attn_proj_->verify_loaded_weights(prefix + "img_attn_proj."); - txt_attn_qkv_->verify_loaded_weights(prefix + "txt_attn_qkv."); + txt_attn_q_->verify_loaded_weights(prefix + "txt_attn_qkv.q."); + txt_attn_k_->verify_loaded_weights(prefix + "txt_attn_qkv.k."); + txt_attn_v_->verify_loaded_weights(prefix + "txt_attn_qkv.v."); txt_attn_q_norm_->verify_loaded_weights(prefix + "txt_attn_q_norm."); txt_attn_k_norm_->verify_loaded_weights(prefix + "txt_attn_k_norm."); txt_attn_proj_->verify_loaded_weights(prefix + "txt_attn_proj."); @@ -354,16 +376,60 @@ class JoyAttentionImpl final : public torch::nn::Module { public: int64_t heads_; - layer::AddMatmulWeightTransposed img_attn_qkv_{nullptr}; + int64_t head_dim_; + layer::AddMatmulWeightTransposed img_attn_q_{nullptr}; + layer::AddMatmulWeightTransposed img_attn_k_{nullptr}; + layer::AddMatmulWeightTransposed img_attn_v_{nullptr}; layer::RMSNorm img_attn_q_norm_{nullptr}; layer::RMSNorm img_attn_k_norm_{nullptr}; layer::AddMatmulWeightTransposed img_attn_proj_{nullptr}; - layer::AddMatmulWeightTransposed txt_attn_qkv_{nullptr}; + layer::AddMatmulWeightTransposed txt_attn_q_{nullptr}; + layer::AddMatmulWeightTransposed txt_attn_k_{nullptr}; + layer::AddMatmulWeightTransposed txt_attn_v_{nullptr}; layer::RMSNorm txt_attn_q_norm_{nullptr}; layer::RMSNorm txt_attn_k_norm_{nullptr}; layer::AddMatmulWeightTransposed txt_attn_proj_{nullptr}; private: + void load_qkv_projections(const StateDict& state_dict, + const std::string& prefix, + layer::AddMatmulWeightTransposed& q_projection, + layer::AddMatmulWeightTransposed& k_projection, + layer::AddMatmulWeightTransposed& v_projection) { + torch::Tensor fused_weight = state_dict.get_tensor(prefix + ".weight"); + torch::Tensor fused_bias = state_dict.get_tensor(prefix + ".bias"); + if (!fused_weight.defined() && !fused_bias.defined()) { + return; + } + + const int64_t inner = heads_ * head_dim_; + CHECK(fused_weight.defined()) << prefix << ".weight is missing"; + CHECK(fused_bias.defined()) << prefix << ".bias is missing"; + CHECK_EQ(fused_weight.dim(), 2) << prefix << ".weight must be 2D"; + CHECK_EQ(fused_weight.size(0), inner * 3) + << prefix << ".weight output size mismatch"; + CHECK_EQ(fused_bias.dim(), 1) << prefix << ".bias must be 1D"; + CHECK_EQ(fused_bias.size(0), inner * 3) + << prefix << ".bias output size mismatch"; + + auto load_projection = [&](layer::AddMatmulWeightTransposed& projection, + int64_t projection_index) { + const int64_t start = projection_index * inner; + std::unordered_map projection_tensors; + projection_tensors.reserve(2); + projection_tensors.emplace( + "weight", fused_weight.slice(/*dim=*/0, start, start + inner)); + projection_tensors.emplace( + "bias", fused_bias.slice(/*dim=*/0, start, start + inner)); + StateDict projection_state_dict(std::move(projection_tensors)); + projection->load_state_dict(projection_state_dict); + }; + + load_projection(q_projection, /*projection_index=*/0); + load_projection(k_projection, /*projection_index=*/1); + load_projection(v_projection, /*projection_index=*/2); + } + torch::TensorOptions options_; }; TORCH_MODULE(JoyAttention); @@ -396,18 +462,19 @@ class JoyAttnProcessor final : public JoyAttnProcessorBase { 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); + auto img_q = + attn->img_attn_q_->forward(hidden_states).unflatten(-1, reshape); + auto img_k = + attn->img_attn_k_->forward(hidden_states).unflatten(-1, reshape); + auto img_v = + attn->img_attn_v_->forward(hidden_states).unflatten(-1, reshape); + auto txt_q = attn->txt_attn_q_->forward(encoder_hidden_states) + .unflatten(-1, reshape); + auto txt_k = attn->txt_attn_k_->forward(encoder_hidden_states) + .unflatten(-1, reshape); + auto txt_v = attn->txt_attn_v_->forward(encoder_hidden_states) + .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)); @@ -456,13 +523,9 @@ class JoySequenceParallelAttnProcessor final 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 = + attn->img_attn_q_->forward(hidden_states).unflatten(-1, reshape); auto img_q_handler = xllm::parallel_state::all_to_all_4D(img_q, @@ -472,6 +535,8 @@ class JoySequenceParallelAttnProcessor final process_group(), /*enable_sp_pad=*/true, /*tensor_name=*/"hidden_states"); + auto img_k = + attn->img_attn_k_->forward(hidden_states).unflatten(-1, reshape); auto img_k_handler = xllm::parallel_state::all_to_all_4D(img_k, /*scatter_idx=*/2, @@ -480,6 +545,8 @@ class JoySequenceParallelAttnProcessor final process_group(), /*enable_sp_pad=*/true, /*tensor_name=*/"hidden_states"); + auto img_v = + attn->img_attn_v_->forward(hidden_states).unflatten(-1, reshape); auto img_v_handler = xllm::parallel_state::all_to_all_4D(img_v, /*scatter_idx=*/2, @@ -489,11 +556,8 @@ class JoySequenceParallelAttnProcessor final /*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 = attn->txt_attn_q_->forward(encoder_hidden_states) + .unflatten(-1, reshape); auto txt_q_handler = xllm::parallel_state::all_to_all_4D( txt_q, @@ -503,6 +567,8 @@ class JoySequenceParallelAttnProcessor final process_group(), /*enable_sp_pad=*/true, /*tensor_name=*/"encoder_hidden_states"); + auto txt_k = attn->txt_attn_k_->forward(encoder_hidden_states) + .unflatten(-1, reshape); auto txt_k_handler = xllm::parallel_state::all_to_all_4D( txt_k, /*scatter_idx=*/2, @@ -511,6 +577,8 @@ class JoySequenceParallelAttnProcessor final process_group(), /*enable_sp_pad=*/true, /*tensor_name=*/"encoder_hidden_states"); + auto txt_v = attn->txt_attn_v_->forward(encoder_hidden_states) + .unflatten(-1, reshape); auto txt_v_handler = xllm::parallel_state::all_to_all_4D( txt_v, /*scatter_idx=*/2, @@ -620,9 +688,26 @@ class TransformerBlockImpl final : public torch::nn::Module { int64_t mlp_hidden = static_cast(dim * mlp_width_ratio); img_mod_ = register_module("img_mod", Modulate(dim, 6)); +#if defined(USE_NPU) + const torch::TensorOptions options = context.get_tensor_options(); + img_norm1_ = register_module( + "img_norm1", + layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options)); + img_norm2_ = register_module( + "img_norm2", + layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options)); +#endif img_mlp_ = register_module( "img_mlp", qwenimage::FeedForward(context, dim, dim, /*mult=*/4)); txt_mod_ = register_module("txt_mod", Modulate(dim, 6)); +#if defined(USE_NPU) + txt_norm1_ = register_module( + "txt_norm1", + layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options)); + txt_norm2_ = register_module( + "txt_norm2", + layer::AdaLayerNorm(dim, eps, /*elementwise_affine=*/false, options)); +#endif txt_mlp_ = register_module( "txt_mlp", qwenimage::FeedForward(context, dim, dim, /*mult=*/4)); attn_ = register_module( @@ -668,32 +753,50 @@ class TransformerBlockImpl final : public torch::nn::Module { auto ehs = encoder_hidden_states; // --- attention --- +#if defined(USE_NPU) + auto img_modulated = img_norm1_->forward(hs, img_c1, img_s1); + auto txt_modulated = txt_norm1_->forward(ehs, txt_c1, txt_s1); +#else auto img_normed = fp32_norm(hs, eps_); auto txt_normed = fp32_norm(ehs, eps_); auto img_modulated = img_normed * (1 + img_c1.unsqueeze(1)) + img_s1.unsqueeze(1); auto txt_modulated = txt_normed * (1 + txt_c1.unsqueeze(1)) + txt_s1.unsqueeze(1); +#endif 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); + img_attn.mul_(img_g1.unsqueeze(1)); + img_attn.add_(hs); + hs = img_attn; + txt_attn.mul_(txt_g1.unsqueeze(1)); + txt_attn.add_(ehs); + ehs = txt_attn; // --- FFN --- +#if defined(USE_NPU) + auto img_ffn_in = img_norm2_->forward(hs, img_c2, img_s2); + auto txt_ffn_in = txt_norm2_->forward(ehs, txt_c2, txt_s2); +#else auto img_ffn_normed = fp32_norm(hs, eps_); auto txt_ffn_normed = fp32_norm(ehs, eps_); auto img_ffn_in = img_ffn_normed * (1 + img_c2.unsqueeze(1)) + img_s2.unsqueeze(1); auto txt_ffn_in = txt_ffn_normed * (1 + txt_c2.unsqueeze(1)) + txt_s2.unsqueeze(1); +#endif auto img_ffn = img_mlp_->forward(img_ffn_in); auto txt_ffn = txt_mlp_->forward(txt_ffn_in); - hs = hs + img_ffn * img_g2.unsqueeze(1); - ehs = ehs + txt_ffn * txt_g2.unsqueeze(1); + img_ffn.mul_(img_g2.unsqueeze(1)); + img_ffn.add_(hs); + hs = img_ffn; + txt_ffn.mul_(txt_g2.unsqueeze(1)); + txt_ffn.add_(ehs); + ehs = txt_ffn; return std::make_tuple(hs, ehs); } @@ -717,6 +820,12 @@ class TransformerBlockImpl final : public torch::nn::Module { double eps_; Modulate img_mod_{nullptr}; Modulate txt_mod_{nullptr}; +#if defined(USE_NPU) + layer::AdaLayerNorm img_norm1_{nullptr}; + layer::AdaLayerNorm img_norm2_{nullptr}; + layer::AdaLayerNorm txt_norm1_{nullptr}; + layer::AdaLayerNorm txt_norm2_{nullptr}; +#endif qwenimage::FeedForward img_mlp_{nullptr}; qwenimage::FeedForward txt_mlp_{nullptr}; JoyAttention attn_{nullptr}; From 67a56eff2adc02fcb4697a2045b3916f3d04af80 Mon Sep 17 00:00:00 2001 From: shan-chen-feng Date: Fri, 31 Jul 2026 14:23:02 +0800 Subject: [PATCH 3/4] refact sequence parallel code --- .../parallel_state/parallel_state.cpp | 37 +- .../framework/parallel_state/parallel_state.h | 15 +- .../dit/attn_processor/attn_processor.h | 60 --- .../attn_processor/attn_processor_factory.h | 79 ---- xllm/models/dit/parallel_mode.h | 47 -- .../pipelines/pipeline_joyimage_edit_plus.h | 53 +-- .../sequence_parallel_context.h | 157 +++++++ .../sequence_parallel_mixin.h | 69 +-- .../transformer_joyimage_edit_plus.h | 440 +++++++----------- xllm/models/dit/utils/dit_parallel_mixin.h | 3 - 10 files changed, 367 insertions(+), 593 deletions(-) delete mode 100644 xllm/models/dit/attn_processor/attn_processor.h delete mode 100644 xllm/models/dit/attn_processor/attn_processor_factory.h delete mode 100644 xllm/models/dit/parallel_mode.h create mode 100644 xllm/models/dit/sequence_parallel/sequence_parallel_context.h diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index ea6bb2f460..91c20b4c86 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -16,7 +16,6 @@ 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" @@ -302,9 +301,7 @@ 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, - const std::string& tensor_name) { + ProcessGroup* process_group) { if (!process_group) { return [input]() { return input; }; } @@ -353,13 +350,7 @@ std::function all_to_all_4D(const torch::Tensor& input, .transpose(0, 1) .contiguous() .reshape({bs, seqlen, shard_head_num, head_size}); - 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; - }; + return [output]() { return output; }; } else { c10::intrusive_ptr all2all_work; process_group->all_to_all_single(output, @@ -373,32 +364,18 @@ std::function all_to_all_4D(const torch::Tensor& input, bs, seqlen, shard_head_num, - head_size, - enable_sp_pad, - tensor_name]() mutable -> torch::Tensor { + head_size]() mutable -> torch::Tensor { all2all_work->wait(); - auto comm_output = - output.reshape({seqlen, bs, shard_head_num, head_size}) - .transpose(0, 1) - .contiguous() - .reshape({bs, seqlen, shard_head_num, head_size}); - if (enable_sp_pad) { - xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( - comm_output, tensor_name, /*dim=*/1); - } - return comm_output; + return output.reshape({seqlen, bs, shard_head_num, head_size}) + .transpose(0, 1) + .contiguous() + .reshape({bs, seqlen, shard_head_num, head_size}); }; } } else if (scatter_idx == 1 && gather_idx == 2) { // 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 9953e30b47..38e8b3aed2 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -15,7 +15,7 @@ limitations under the License. #pragma once -#include +#include #include "parallel_args.h" #include "process_group.h" @@ -80,14 +80,11 @@ 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 async_ops, - ProcessGroup* process_group, - bool enable_sp_pad = false, - const std::string& tensor_name = ""); +std::function all_to_all_4D(const torch::Tensor& input, + int32_t scatter_idx, + int32_t gather_idx, + bool async_ops, + ProcessGroup* process_group); // Create a process group where each process has a single device // devices: list of devices to create process groups on. diff --git a/xllm/models/dit/attn_processor/attn_processor.h b/xllm/models/dit/attn_processor/attn_processor.h deleted file mode 100644 index 736afe4d5d..0000000000 --- a/xllm/models/dit/attn_processor/attn_processor.h +++ /dev/null @@ -1,60 +0,0 @@ -/* 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 deleted file mode 100644 index 9bf45e3d53..0000000000 --- a/xllm/models/dit/attn_processor/attn_processor_factory.h +++ /dev/null @@ -1,79 +0,0 @@ -/* 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/parallel_mode.h b/xllm/models/dit/parallel_mode.h deleted file mode 100644 index 224eee6e4c..0000000000 --- a/xllm/models/dit/parallel_mode.h +++ /dev/null @@ -1,47 +0,0 @@ -/* 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 index 92c7b9a215..118c5a0ae8 100644 --- a/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h +++ b/xllm/models/dit/pipelines/pipeline_joyimage_edit_plus.h @@ -413,12 +413,12 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module, const TransformerForwardContext& transformer_context = is_positive ? prompt_transformer_context : negative_transformer_context; - return sequence_parallel_forward(latents, - t_expand, - encoder_hidden_states, - transformer_context, - /*use_cfg=*/!is_positive, - /*step_index=*/i + 1); + return transformer_forward(latents, + t_expand, + encoder_hidden_states, + transformer_context, + /*use_cfg=*/!is_positive, + /*step_index=*/i + 1); }); auto comb = uncond + guidance_scale * (cond - uncond); // Norm-rescale (diffusers): comb * (||cond|| / ||comb||) over channel @@ -430,12 +430,12 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module, noise_pred = comb * (cond_norm / noise_norm.clamp_min(1e-6)); } else { torch::Tensor t_expand = t.repeat({batch_size}); - noise_pred = sequence_parallel_forward(latents, - t_expand, - prompt_embeds, - prompt_transformer_context, - /*use_cfg=*/false, - /*step_index=*/i + 1); + noise_pred = transformer_forward(latents, + t_expand, + prompt_embeds, + prompt_transformer_context, + /*use_cfg=*/false, + /*step_index=*/i + 1); } latents = scheduler_->step(noise_pred, t, latents).to(latents.dtype()); @@ -705,32 +705,21 @@ class JoyImageEditPlusPipelineImpl : public torch::nn::Module, return {torch::cat(embeddings, /*dim=*/0), torch::cat(masks, /*dim=*/0)}; } - torch::Tensor sequence_parallel_forward( + torch::Tensor transformer_forward( const torch::Tensor& hidden_states, const torch::Tensor& timestep, const torch::Tensor& encoder_hidden_states, const TransformerForwardContext& forward_context, bool use_cfg, int64_t step_index) { - 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"); + return transformer_->forward(hidden_states, + timestep, + encoder_hidden_states, + forward_context.rope_cos, + forward_context.rope_sin, + forward_context.attention_mask, + use_cfg, + step_index); } TransformerForwardContext prepare_transformer_context( diff --git a/xllm/models/dit/sequence_parallel/sequence_parallel_context.h b/xllm/models/dit/sequence_parallel/sequence_parallel_context.h new file mode 100644 index 0000000000..d02c808c74 --- /dev/null +++ b/xllm/models/dit/sequence_parallel/sequence_parallel_context.h @@ -0,0 +1,157 @@ +/* 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 +#include + +#include "core/framework/parallel_state/parallel_state.h" + +namespace xllm::dit { + +using SequenceParallelWork = std::function; + +class SequenceParallelContext final { + public: + explicit SequenceParallelContext(ProcessGroup* process_group) + : process_group_(process_group) {} + + int32_t world_size() const { + return process_group_ == nullptr ? 1 : process_group_->world_size(); + } + + bool enabled() const { return world_size() > 1; } + + torch::Tensor scatter_sequence(const torch::Tensor& input, + const std::string& tensor_name, + int64_t sequence_dim) { + if (!enabled() || !input.defined()) { + return input; + } + + const int64_t sequence_length = input.size(sequence_dim); + const int64_t padding_length = + (world_size() - sequence_length % world_size()) % world_size(); + padding_lengths_[tensor_name] = padding_length; + torch::Tensor padded_input = pad_right(input, padding_length, sequence_dim); + return parallel_state::scatter( + padded_input, process_group_, static_cast(sequence_dim)); + } + + torch::Tensor gather_sequence(const torch::Tensor& input, + const std::string& tensor_name, + int64_t sequence_dim) const { + if (!enabled() || !input.defined()) { + return input; + } + + torch::Tensor output = parallel_state::gather( + input.contiguous(), process_group_, static_cast(sequence_dim)); + return unpad_right(output, padding_length(tensor_name), sequence_dim); + } + + SequenceParallelWork launch_sequence_to_head( + const torch::Tensor& input, + const std::string& tensor_name) const { + if (!enabled()) { + return [input]() { return input; }; + } + + SequenceParallelWork work = + parallel_state::all_to_all_4D(input, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + process_group_); + const int64_t padding = padding_length(tensor_name); + return [work = std::move(work), padding]() mutable { + return unpad_right(work(), padding, /*dim=*/1); + }; + } + + SequenceParallelWork launch_head_to_sequence( + const torch::Tensor& input, + const std::string& tensor_name) const { + if (!enabled()) { + return [input]() { return input; }; + } + + torch::Tensor padded_input = + pad_right(input, padding_length(tensor_name), /*dim=*/1); + return parallel_state::all_to_all_4D(padded_input, + /*scatter_idx=*/1, + /*gather_idx=*/2, + /*async_ops=*/true, + process_group_); + } + + private: + static int64_t normalize_dim(const torch::Tensor& input, int64_t dim) { + const int64_t normalized_dim = dim < 0 ? input.dim() + dim : dim; + CHECK_GE(normalized_dim, 0) << "Invalid tensor dimension: " << dim; + CHECK_LT(normalized_dim, input.dim()) + << "Invalid tensor dimension: " << dim; + return normalized_dim; + } + + static torch::Tensor pad_right(const torch::Tensor& input, + int64_t padding_length, + int64_t dim) { + if (!input.defined() || padding_length == 0) { + return input; + } + + const int64_t normalized_dim = normalize_dim(input, dim); + std::vector padding(static_cast(input.dim() * 2), 0); + const int64_t padding_index = 2 * (input.dim() - normalized_dim - 1) + 1; + padding[static_cast(padding_index)] = padding_length; + return torch::pad(input, padding, "constant", 0); + } + + static torch::Tensor unpad_right(const torch::Tensor& input, + int64_t padding_length, + int64_t dim) { + if (!input.defined() || padding_length == 0) { + return input; + } + + const int64_t normalized_dim = normalize_dim(input, dim); + CHECK_GE(input.size(normalized_dim), padding_length) + << "Padding length exceeds tensor size"; + return input.narrow(normalized_dim, + /*start=*/0, + input.size(normalized_dim) - padding_length); + } + + int64_t padding_length(const std::string& tensor_name) const { + auto padding_it = padding_lengths_.find(tensor_name); + CHECK(padding_it != padding_lengths_.end()) + << "Missing sequence-parallel padding metadata: " << tensor_name; + return padding_it->second; + } + + ProcessGroup* process_group_{nullptr}; + std::unordered_map padding_lengths_; +}; + +} // namespace xllm::dit diff --git a/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h b/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h index 13dd09ad77..a1ad80dc08 100644 --- a/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h +++ b/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h @@ -23,8 +23,7 @@ limitations under the License. #include #include -#include "core/framework/parallel_state/parallel_state.h" -#include "models/dit/utils/sequence_parallel_pad_manager.h" +#include "models/dit/sequence_parallel/sequence_parallel_context.h" namespace xllm::dit { @@ -33,7 +32,7 @@ using SequenceParallelTensorMap = using SequenceParallelTensorDims = std::unordered_map; class SequenceParallelMixin { - public: + protected: SequenceParallelMixin(ProcessGroup* process_group, SequenceParallelTensorDims input_sequence_dims, SequenceParallelTensorDims output_sequence_dims) @@ -45,67 +44,29 @@ class SequenceParallelMixin { 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; - } - + SequenceParallelContext context(process_group_); + SequenceParallelTensorMap local_inputs = 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()) + auto tensor_it = local_inputs.find(tensor_name); + CHECK(tensor_it != local_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; + tensor_it->second = context.scatter_sequence( + tensor_it->second, tensor_name, sequence_dim); } + SequenceParallelTensorMap outputs = + std::forward(forward_fn)(local_inputs, context); for (const auto& [tensor_name, sequence_dim] : output_sequence_dims_) { - auto tensor_it = gathered_outputs.find(tensor_name); - CHECK(tensor_it != gathered_outputs.end()) + auto tensor_it = outputs.find(tensor_name); + CHECK(tensor_it != 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); + context.gather_sequence(tensor_it->second, tensor_name, sequence_dim); } - return gathered_outputs; - } - - bool sequence_parallel_enabled() const { - return process_group_ != nullptr && process_group_->world_size() > 1; + return outputs; } + private: ProcessGroup* process_group_{nullptr}; SequenceParallelTensorDims input_sequence_dims_; SequenceParallelTensorDims output_sequence_dims_; diff --git a/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h index 834855c535..983869bdd3 100644 --- a/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h +++ b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h @@ -37,8 +37,8 @@ limitations under the License. #include "core/layers/common/ada_layer_norm.h" #include "core/layers/common/add_matmul.h" #include "core/layers/common/rms_norm.h" -#include "models/dit/attn_processor/attn_processor.h" -#include "models/dit/attn_processor/attn_processor_factory.h" +#include "models/dit/sequence_parallel/sequence_parallel_context.h" +#include "models/dit/sequence_parallel/sequence_parallel_mixin.h" #include "models/dit/transformers/transformer_qwen_image.h" #include "models/model_registry.h" @@ -286,8 +286,8 @@ class TimeTextEmbeddingImpl final : public torch::nn::Module { }; TORCH_MODULE(TimeTextEmbedding); -// Internal module that only registers parameters used by Joy attention. -// Attention forward is implemented by processor classes. +using JoyAttentionOutput = std::tuple; + class JoyAttentionImpl final : public torch::nn::Module { public: JoyAttentionImpl(const ModelContext& context, @@ -335,6 +335,85 @@ class JoyAttentionImpl final : public torch::nn::Module { layer::AddMatmulWeightTransposed(inner, dim, true, options_)); } + JoyAttentionOutput forward( + const torch::Tensor& hidden_states, + const torch::Tensor& encoder_hidden_states, + const torch::Tensor& rope_cos, + const torch::Tensor& rope_sin, + const torch::Tensor& attn_mask, + xllm::dit::SequenceParallelContext& sequence_parallel_context) { + std::vector reshape = {heads_, -1}; + torch::Tensor img_q = + img_attn_q_->forward(hidden_states).unflatten(-1, reshape); + auto img_q_work = sequence_parallel_context.launch_sequence_to_head( + img_q, /*tensor_name=*/"hidden_states"); + torch::Tensor img_k = + img_attn_k_->forward(hidden_states).unflatten(-1, reshape); + auto img_k_work = sequence_parallel_context.launch_sequence_to_head( + img_k, /*tensor_name=*/"hidden_states"); + torch::Tensor img_v = + img_attn_v_->forward(hidden_states).unflatten(-1, reshape); + auto img_v_work = sequence_parallel_context.launch_sequence_to_head( + img_v, /*tensor_name=*/"hidden_states"); + + torch::Tensor txt_q = + txt_attn_q_->forward(encoder_hidden_states).unflatten(-1, reshape); + auto txt_q_work = sequence_parallel_context.launch_sequence_to_head( + txt_q, /*tensor_name=*/"encoder_hidden_states"); + torch::Tensor txt_k = + txt_attn_k_->forward(encoder_hidden_states).unflatten(-1, reshape); + auto txt_k_work = sequence_parallel_context.launch_sequence_to_head( + txt_k, /*tensor_name=*/"encoder_hidden_states"); + torch::Tensor txt_v = + txt_attn_v_->forward(encoder_hidden_states).unflatten(-1, reshape); + auto txt_v_work = sequence_parallel_context.launch_sequence_to_head( + txt_v, /*tensor_name=*/"encoder_hidden_states"); + + img_q = img_q_work(); + img_k = img_k_work(); + txt_q = txt_q_work(); + txt_k = txt_k_work(); + + img_q = std::get<0>(img_attn_q_norm_->forward(img_q)); + img_k = std::get<0>(img_attn_k_norm_->forward(img_k)); + txt_q = std::get<0>(txt_attn_q_norm_->forward(txt_q)); + txt_k = std::get<0>(txt_attn_k_norm_->forward(txt_k)); + + if (rope_cos.defined()) { + img_q = apply_rotary_emb_batched(img_q, rope_cos, rope_sin); + img_k = apply_rotary_emb_batched(img_k, rope_cos, rope_sin); + } + + img_v = img_v_work(); + txt_v = txt_v_work(); + + torch::Tensor joint_q = torch::cat({img_q, txt_q}, /*dim=*/1); + torch::Tensor joint_k = torch::cat({img_k, txt_k}, /*dim=*/1); + torch::Tensor joint_v = torch::cat({img_v, txt_v}, /*dim=*/1); + const int64_t local_heads = heads_ / sequence_parallel_context.world_size(); + torch::Tensor joint = joyimage_joint_attention( + joint_q, joint_k, joint_v, local_heads, attn_mask); + joint = joint.to(joint_q.dtype()); + + const int64_t image_sequence_length = img_q.size(1); + torch::Tensor img_out = joint.slice( + /*dim=*/1, /*start=*/0, /*end=*/image_sequence_length); + torch::Tensor txt_out = joint.slice(/*dim=*/1, + /*start=*/image_sequence_length, + /*end=*/joint.size(1)); + + auto img_out_work = sequence_parallel_context.launch_head_to_sequence( + img_out, /*tensor_name=*/"hidden_states"); + auto txt_out_work = sequence_parallel_context.launch_head_to_sequence( + txt_out, /*tensor_name=*/"encoder_hidden_states"); + + img_out = img_out_work().flatten(/*start_dim=*/2, /*end_dim=*/3); + img_out = img_attn_proj_->forward(img_out); + txt_out = txt_out_work().flatten(/*start_dim=*/2, /*end_dim=*/3); + txt_out = txt_attn_proj_->forward(txt_out); + return std::make_tuple(img_out, txt_out); + } + void load_state_dict(const StateDict& state_dict) { load_qkv_projections(state_dict, /*prefix=*/"img_attn_qkv", @@ -374,7 +453,7 @@ class JoyAttentionImpl final : public torch::nn::Module { txt_attn_proj_->verify_loaded_weights(prefix + "txt_attn_proj."); } - public: + private: int64_t heads_; int64_t head_dim_; layer::AddMatmulWeightTransposed img_attn_q_{nullptr}; @@ -434,244 +513,6 @@ class JoyAttentionImpl final : public torch::nn::Module { }; 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(); - std::vector reshape = {attn->heads_, -1}; - auto img_q = - attn->img_attn_q_->forward(hidden_states).unflatten(-1, reshape); - auto img_k = - attn->img_attn_k_->forward(hidden_states).unflatten(-1, reshape); - auto img_v = - attn->img_attn_v_->forward(hidden_states).unflatten(-1, reshape); - auto txt_q = attn->txt_attn_q_->forward(encoder_hidden_states) - .unflatten(-1, reshape); - auto txt_k = attn->txt_attn_k_->forward(encoder_hidden_states) - .unflatten(-1, reshape); - auto txt_v = attn->txt_attn_v_->forward(encoder_hidden_states) - .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(); - std::vector reshape = {attn->heads_, -1}; - auto img_q = - attn->img_attn_q_->forward(hidden_states).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 = - attn->img_attn_k_->forward(hidden_states).unflatten(-1, reshape); - 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 = - attn->img_attn_v_->forward(hidden_states).unflatten(-1, reshape); - 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_q = attn->txt_attn_q_->forward(encoder_hidden_states) - .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 = attn->txt_attn_k_->forward(encoder_hidden_states) - .unflatten(-1, reshape); - 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 = attn->txt_attn_v_->forward(encoder_hidden_states) - .unflatten(-1, reshape); - 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: @@ -680,9 +521,6 @@ class TransformerBlockImpl final : public torch::nn::Module { 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); @@ -712,9 +550,6 @@ class TransformerBlockImpl final : public torch::nn::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) } @@ -732,7 +567,8 @@ class TransformerBlockImpl final : public torch::nn::Module { const torch::Tensor& temb, // [B, hidden] const torch::Tensor& rope_cos, const torch::Tensor& rope_sin, - const torch::Tensor& attn_mask) { + const torch::Tensor& attn_mask, + xllm::dit::SequenceParallelContext& sequence_parallel_context) { auto img_mod = img_mod_->forward(temb); // 6 x [B, hidden] auto txt_mod = txt_mod_->forward(temb); @@ -765,8 +601,12 @@ class TransformerBlockImpl final : public torch::nn::Module { txt_normed * (1 + txt_c1.unsqueeze(1)) + txt_s1.unsqueeze(1); #endif - auto attn_out = attn_processor_->forward( - img_modulated, txt_modulated, rope_cos, rope_sin, attn_mask); + JoyAttentionOutput attn_out = attn_->forward(img_modulated, + txt_modulated, + rope_cos, + rope_sin, + attn_mask, + sequence_parallel_context); auto img_attn = std::get<0>(attn_out); auto txt_attn = std::get<1>(attn_out); @@ -829,11 +669,10 @@ class TransformerBlockImpl final : public torch::nn::Module { qwenimage::FeedForward img_mlp_{nullptr}; qwenimage::FeedForward txt_mlp_{nullptr}; JoyAttention attn_{nullptr}; - std::unique_ptr attn_processor_; }; TORCH_MODULE(TransformerBlock); -class JoyImageEditPlusTransformer3DModelImpl +class JoyImageEditPlusTransformer3DModelImpl final : public torch::nn::Module, public xllm::dit::SequenceParallelMixin { public: @@ -845,10 +684,6 @@ class JoyImageEditPlusTransformer3DModelImpl {{"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(); @@ -863,6 +698,11 @@ class JoyImageEditPlusTransformer3DModelImpl head_dim_ = hidden_size_ / num_heads_; CHECK_EQ(hidden_size_ % num_heads_, 0) << "hidden_size must be divisible by num_attention_heads"; + const int32_t sp_size = parallel_args.dit_sp_group_ == nullptr + ? 1 + : parallel_args.dit_sp_group_->world_size(); + CHECK_EQ(num_heads_ % sp_size, 0) + << "JoyImageEditPlus attention heads must be divisible by sp_size"; // Conv3d patchifier: kernel == stride == patch_size. img_in_ = register_module( @@ -885,14 +725,8 @@ class JoyImageEditPlusTransformer3DModelImpl 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_); + auto block = TransformerBlock( + context, hidden_size_, num_heads_, head_dim_, mlp_width_ratio); double_blocks_->push_back(block); block_layers_.push_back(block); } @@ -915,15 +749,56 @@ class JoyImageEditPlusTransformer3DModelImpl 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); + xllm::dit::SequenceParallelTensorMap model_outputs = + sequence_parallel_forward( + {{"hidden_states", hidden_states}, + {"encoder_hidden_states", encoder_hidden_states}}, + [this, + ×tep, + &rope_cos, + &rope_sin, + &attention_mask, + use_cfg, + step_index]( + const xllm::dit::SequenceParallelTensorMap& model_inputs, + xllm::dit::SequenceParallelContext& sequence_parallel_context) { + torch::Tensor output = + forward_impl(model_inputs.at("hidden_states"), + timestep, + model_inputs.at("encoder_hidden_states"), + rope_cos, + rope_sin, + attention_mask, + use_cfg, + step_index, + sequence_parallel_context); + return xllm::dit::SequenceParallelTensorMap{ + {"hidden_states", output}}; + }); + return model_outputs.at("hidden_states"); + } + + private: + torch::Tensor forward_impl( + const torch::Tensor& local_hidden_states, + const torch::Tensor& timestep, + const torch::Tensor& local_encoder_hidden_states, + const torch::Tensor& rope_cos, + const torch::Tensor& rope_sin, + const torch::Tensor& attention_mask, + bool use_cfg, + int64_t step_index, + xllm::dit::SequenceParallelContext& sequence_parallel_context) { + int64_t batch_size = local_hidden_states.size(0); + int64_t sequence_length = local_hidden_states.size(1); + int64_t channels = local_hidden_states.size(2); + int64_t pt = local_hidden_states.size(3); + int64_t ph = local_hidden_states.size(4); + int64_t pw = local_hidden_states.size(5); // 1. Condition embeddings. - auto cond = condition_embedder_->forward(timestep, encoder_hidden_states); + auto cond = + condition_embedder_->forward(timestep, local_encoder_hidden_states); auto vec = std::get<1>(cond); // [B, hidden*6] -> but we use per-block temb auto txt = std::get<2>(cond); // [B, L, hidden] // vec is the projected timestep [B, 6*hidden]; blocks re-add modulate_table @@ -935,9 +810,10 @@ class JoyImageEditPlusTransformer3DModelImpl 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}); + auto x = local_hidden_states.reshape( + {batch_size * sequence_length, channels, pt, ph, pw}); x = img_in_->forward(x); // [B*N, hidden, 1, 1, 1] - auto img = x.reshape({B, N, hidden_size_}); + auto img = x.reshape({batch_size, sequence_length, hidden_size_}); // 3. Blocks with optional DiT cache. torch::Tensor original_img = img; @@ -959,8 +835,14 @@ class JoyImageEditPlusTransformer3DModelImpl 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); + std::tie(img, txt) = + block_layers_[block_index]->forward(img, + txt, + temb6, + rope_cos, + rope_sin, + attention_mask, + sequence_parallel_context); } TensorMap block_after_map = { @@ -988,11 +870,12 @@ class JoyImageEditPlusTransformer3DModelImpl // 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_}) + img = img.reshape({batch_size, sequence_length, pt, ph, pw, out_channels_}) .permute({0, 1, 5, 2, 3, 4}); return img; } + public: static torch::Tensor fp32_norm_out(const torch::Tensor& x) { auto xf = x.to(torch::kFloat32); auto out = @@ -1047,7 +930,6 @@ class JoyImageEditPlusTransformer3DModelImpl 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}; diff --git a/xllm/models/dit/utils/dit_parallel_mixin.h b/xllm/models/dit/utils/dit_parallel_mixin.h index cbbb52b00e..6f68ed6891 100644 --- a/xllm/models/dit/utils/dit_parallel_mixin.h +++ b/xllm/models/dit/utils/dit_parallel_mixin.h @@ -70,8 +70,5 @@ class CFGParallelMixin { // Mixin for VAE parallelism (to be implemented). class VaeParallelMixin {}; -// Mixin for sequence parallelism (to be implemented). -class SpParallelMixin {}; - } // namespace dit } // namespace xllm From d6bb9027ce899df2d70df44239ae293dbfa685bc Mon Sep 17 00:00:00 2001 From: shan-chen-feng Date: Fri, 31 Jul 2026 19:11:55 +0800 Subject: [PATCH 4/4] refact: refact code for operations. --- xllm/core/kernels/npu/attention.cpp | 11 ++ xllm/core/kernels/npu/npu_ops_api.h | 9 +- xllm/core/kernels/npu/rope.cpp | 11 +- xllm/core/kernels/ops_api.cpp | 13 +- xllm/core/layers/common/qwen2_attention.cpp | 13 -- xllm/core/layers/common/rms_norm.cpp | 2 + xllm/core/layers/npu_torch/attention.cpp | 31 ++-- xllm/models/dit/pipelines/pipeline_wan_i2v.h | 4 +- .../sequence_parallel_mixin.h | 75 --------- .../transformer_joyimage_edit_plus.h | 157 ++++++++++-------- .../sequence_parallel_mixin.h} | 119 +++++++------ 11 files changed, 210 insertions(+), 235 deletions(-) delete mode 100644 xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h rename xllm/models/dit/{sequence_parallel/sequence_parallel_context.h => utils/sequence_parallel_mixin.h} (56%) diff --git a/xllm/core/kernels/npu/attention.cpp b/xllm/core/kernels/npu/attention.cpp index cd6a0e9f13..663cf44dcd 100644 --- a/xllm/core/kernels/npu/attention.cpp +++ b/xllm/core/kernels/npu/attention.cpp @@ -23,6 +23,17 @@ void reshape_paged_cache(torch::Tensor& key, torch::Tensor& k_cache, std::optional& v_cache, const torch::Tensor& slot_mapping) { + // KV-cache-free full prefill does not need to populate paged caches. + if (!k_cache.defined()) { + CHECK(!v_cache.has_value() || !v_cache->defined()) + << "NPU key and value caches must either both be defined or both be " + "empty."; + return; + } + CHECK(value.has_value() && value->defined()) + << "NPU value tensor must be defined when key cache is defined."; + CHECK(v_cache.has_value() && v_cache->defined()) + << "NPU value cache must be defined when key cache is defined."; atb::npu_reshape_and_cache( key, value.value(), k_cache, v_cache.value(), slot_mapping); } diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index 032d10536d..fac965877c 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -51,10 +51,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); +void apply_rotary(torch::Tensor& query, + torch::Tensor& key, + const torch::Tensor& cos, + const torch::Tensor& sin, + const std::string& input_layout); std::tuple npu_fused_infer_attention( const torch::Tensor& query, diff --git a/xllm/core/kernels/npu/rope.cpp b/xllm/core/kernels/npu/rope.cpp index 2baf629558..06ec8aabf6 100644 --- a/xllm/core/kernels/npu/rope.cpp +++ b/xllm/core/kernels/npu/rope.cpp @@ -20,10 +20,11 @@ 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) { +void apply_rotary(torch::Tensor& query, + torch::Tensor& key, + const torch::Tensor& cos, + const torch::Tensor& sin, + const std::string& input_layout) { CHECK(cos.defined() && sin.defined()); CHECK(cos.sizes() == sin.sizes()); CHECK_GT(cos.dim(), 0); @@ -49,7 +50,7 @@ void apply_rotary_pos_emb(torch::Tensor& query, 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_view, key_view, cos_view, sin_view, input_layout); query = std::get<0>(rotary_result).view(query_shape); key = std::get<1>(rotary_result).view(key_shape); } diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index d383aa4788..56497a39b5 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -115,8 +115,17 @@ void apply_rotary(RotaryParams& params) { params.dynamic_ntk, params.max_query_len); #elif defined(USE_NPU) - npu::apply_rotary( - params.q, params.k, params.cos_sin, params.position_ids.value()); + if (!params.position_ids.has_value() && params.cos.defined() && + params.sin.defined()) { + npu::apply_rotary(params.q, params.k, params.cos, params.sin, "BSND"); + } else { + CHECK(params.position_ids.has_value()) + << "NPU rotary embedding requires position_ids when precomputed " + "cos/sin " + "are unavailable"; + npu::apply_rotary( + params.q, params.k, params.cos_sin, params.position_ids.value()); + } #elif defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_DCU) bool is_neox = !params.interleaved; torch::Tensor pos_ids; diff --git a/xllm/core/layers/common/qwen2_attention.cpp b/xllm/core/layers/common/qwen2_attention.cpp index ec351e4b38..1def5d5d36 100644 --- a/xllm/core/layers/common/qwen2_attention.cpp +++ b/xllm/core/layers/common/qwen2_attention.cpp @@ -19,9 +19,6 @@ 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 @@ -183,17 +180,7 @@ 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/rms_norm.cpp b/xllm/core/layers/common/rms_norm.cpp index 19b06c2cdb..978b955d4f 100644 --- a/xllm/core/layers/common/rms_norm.cpp +++ b/xllm/core/layers/common/rms_norm.cpp @@ -47,6 +47,8 @@ std::tuple> RMSNormImpl::forward( std::optional inplace_output) { auto org_shape = input.sizes().vec(); + // Qwen3-VL vision blocks use LayerNorm, while the NPU fused normalization + // kernel only supports RMSNorm. Use Torch LayerNorm for the NPU TORCH path. if (Platform::is_npu() && mode_ == kLayerNormMode) { torch::Tensor norm_input = input; std::optional residual_out; diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index 438ad4aad2..df5fe7283b 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -52,26 +52,17 @@ std::tuple> AttentionImpl::forward( const bool only_prefill = attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; - torch::Tensor k_cache; - std::optional v_cache; - if (kv_cache.empty()) { - CHECK(attn_metadata.is_prefill) - << "KV-cache-free NPU attention only supports full prefill."; - CHECK(!attn_metadata.use_expanded_decode_for_spec_verify_attention) - << "KV-cache-free NPU attention does not support expanded decode."; - } else { - k_cache = kv_cache.get_k_cache(); - v_cache = kv_cache.get_v_cache(); - - xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; - reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_}); - reshape_paged_cache_params.value = - value.view({-1, num_kv_heads_, head_size_}); - reshape_paged_cache_params.k_cache = k_cache; - reshape_paged_cache_params.v_cache = v_cache; - reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; - xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); - } + torch::Tensor k_cache = kv_cache.get_k_cache(); + torch::Tensor v = value.view({-1, num_kv_heads_, head_size_}); + std::optional v_cache = kv_cache.get_v_cache(); + + xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; + reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_}); + reshape_paged_cache_params.value = v; + reshape_paged_cache_params.k_cache = k_cache; + reshape_paged_cache_params.v_cache = v_cache; + reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; + xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); if (attn_metadata.use_expanded_decode_for_spec_verify_attention) { decoder_forward(query, output, k_cache, v_cache, attn_metadata); diff --git a/xllm/models/dit/pipelines/pipeline_wan_i2v.h b/xllm/models/dit/pipelines/pipeline_wan_i2v.h index e25eaf9dc0..b14e9c5817 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(), options_.dtype()); + video_condition = video_condition.to(options_.device()).to(torch::kFloat32); 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(options_.dtype())).sample; + video = vae_->decode(prepared_latents.to(torch::kFloat32)).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 deleted file mode 100644 index a1ad80dc08..0000000000 --- a/xllm/models/dit/sequence_parallel/sequence_parallel_mixin.h +++ /dev/null @@ -1,75 +0,0 @@ -/* 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/sequence_parallel/sequence_parallel_context.h" - -namespace xllm::dit { - -using SequenceParallelTensorMap = - std::unordered_map; -using SequenceParallelTensorDims = std::unordered_map; - -class SequenceParallelMixin { - protected: - SequenceParallelMixin(ProcessGroup* process_group, - SequenceParallelTensorDims input_sequence_dims, - SequenceParallelTensorDims output_sequence_dims) - : process_group_(process_group), - input_sequence_dims_(std::move(input_sequence_dims)), - output_sequence_dims_(std::move(output_sequence_dims)) {} - - template - SequenceParallelTensorMap sequence_parallel_forward( - const SequenceParallelTensorMap& inputs, - ForwardFn&& forward_fn) const { - SequenceParallelContext context(process_group_); - SequenceParallelTensorMap local_inputs = inputs; - for (const auto& [tensor_name, sequence_dim] : input_sequence_dims_) { - auto tensor_it = local_inputs.find(tensor_name); - CHECK(tensor_it != local_inputs.end()) - << "Missing registered sequence-parallel input: " << tensor_name; - tensor_it->second = context.scatter_sequence( - tensor_it->second, tensor_name, sequence_dim); - } - - SequenceParallelTensorMap outputs = - std::forward(forward_fn)(local_inputs, context); - for (const auto& [tensor_name, sequence_dim] : output_sequence_dims_) { - auto tensor_it = outputs.find(tensor_name); - CHECK(tensor_it != outputs.end()) - << "Missing registered sequence-parallel output: " << tensor_name; - tensor_it->second = - context.gather_sequence(tensor_it->second, tensor_name, sequence_dim); - } - return outputs; - } - - private: - 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 index 983869bdd3..86388b2b18 100644 --- a/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h +++ b/xllm/models/dit/transformers/transformer_joyimage_edit_plus.h @@ -37,9 +37,8 @@ limitations under the License. #include "core/layers/common/ada_layer_norm.h" #include "core/layers/common/add_matmul.h" #include "core/layers/common/rms_norm.h" -#include "models/dit/sequence_parallel/sequence_parallel_context.h" -#include "models/dit/sequence_parallel/sequence_parallel_mixin.h" #include "models/dit/transformers/transformer_qwen_image.h" +#include "models/dit/utils/sequence_parallel_mixin.h" #include "models/model_registry.h" namespace xllm { @@ -294,9 +293,11 @@ class JoyAttentionImpl final : public torch::nn::Module { int64_t dim, int64_t num_heads, int64_t head_dim, + ProcessGroup* sp_group, double eps = 1e-6) : heads_(num_heads), head_dim_(head_dim), + sp_group_(sp_group), options_(context.get_tensor_options()) { const int64_t inner = num_heads * head_dim; @@ -335,44 +336,64 @@ class JoyAttentionImpl final : public torch::nn::Module { layer::AddMatmulWeightTransposed(inner, dim, true, options_)); } - JoyAttentionOutput forward( - const torch::Tensor& hidden_states, - const torch::Tensor& encoder_hidden_states, - const torch::Tensor& rope_cos, - const torch::Tensor& rope_sin, - const torch::Tensor& attn_mask, - xllm::dit::SequenceParallelContext& sequence_parallel_context) { + JoyAttentionOutput forward(const torch::Tensor& hidden_states, + const torch::Tensor& encoder_hidden_states, + const torch::Tensor& rope_cos, + const torch::Tensor& rope_sin, + const torch::Tensor& attn_mask) { std::vector reshape = {heads_, -1}; torch::Tensor img_q = img_attn_q_->forward(hidden_states).unflatten(-1, reshape); - auto img_q_work = sequence_parallel_context.launch_sequence_to_head( - img_q, /*tensor_name=*/"hidden_states"); + auto img_q_work = parallel_state::all_to_all_4D(img_q, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + sp_group_); torch::Tensor img_k = img_attn_k_->forward(hidden_states).unflatten(-1, reshape); - auto img_k_work = sequence_parallel_context.launch_sequence_to_head( - img_k, /*tensor_name=*/"hidden_states"); + auto img_k_work = parallel_state::all_to_all_4D(img_k, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + sp_group_); torch::Tensor img_v = img_attn_v_->forward(hidden_states).unflatten(-1, reshape); - auto img_v_work = sequence_parallel_context.launch_sequence_to_head( - img_v, /*tensor_name=*/"hidden_states"); + auto img_v_work = parallel_state::all_to_all_4D(img_v, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + sp_group_); torch::Tensor txt_q = txt_attn_q_->forward(encoder_hidden_states).unflatten(-1, reshape); - auto txt_q_work = sequence_parallel_context.launch_sequence_to_head( - txt_q, /*tensor_name=*/"encoder_hidden_states"); + auto txt_q_work = parallel_state::all_to_all_4D(txt_q, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + sp_group_); torch::Tensor txt_k = txt_attn_k_->forward(encoder_hidden_states).unflatten(-1, reshape); - auto txt_k_work = sequence_parallel_context.launch_sequence_to_head( - txt_k, /*tensor_name=*/"encoder_hidden_states"); + auto txt_k_work = parallel_state::all_to_all_4D(txt_k, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + sp_group_); torch::Tensor txt_v = txt_attn_v_->forward(encoder_hidden_states).unflatten(-1, reshape); - auto txt_v_work = sequence_parallel_context.launch_sequence_to_head( - txt_v, /*tensor_name=*/"encoder_hidden_states"); - - img_q = img_q_work(); - img_k = img_k_work(); - txt_q = txt_q_work(); - txt_k = txt_k_work(); + auto txt_v_work = parallel_state::all_to_all_4D(txt_v, + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/true, + sp_group_); + + img_q = xllm::dit::SequenceParallelMixin::unpad_tensor( + img_q_work(), /*tensor_name=*/"hidden_states", /*dim=*/1); + img_k = xllm::dit::SequenceParallelMixin::unpad_tensor( + img_k_work(), /*tensor_name=*/"hidden_states", /*dim=*/1); + txt_q = xllm::dit::SequenceParallelMixin::unpad_tensor( + txt_q_work(), /*tensor_name=*/"encoder_hidden_states", /*dim=*/1); + txt_k = xllm::dit::SequenceParallelMixin::unpad_tensor( + txt_k_work(), /*tensor_name=*/"encoder_hidden_states", /*dim=*/1); img_q = std::get<0>(img_attn_q_norm_->forward(img_q)); img_k = std::get<0>(img_attn_k_norm_->forward(img_k)); @@ -384,13 +405,16 @@ class JoyAttentionImpl final : public torch::nn::Module { img_k = apply_rotary_emb_batched(img_k, rope_cos, rope_sin); } - img_v = img_v_work(); - txt_v = txt_v_work(); + img_v = xllm::dit::SequenceParallelMixin::unpad_tensor( + img_v_work(), /*tensor_name=*/"hidden_states", /*dim=*/1); + txt_v = xllm::dit::SequenceParallelMixin::unpad_tensor( + txt_v_work(), /*tensor_name=*/"encoder_hidden_states", /*dim=*/1); torch::Tensor joint_q = torch::cat({img_q, txt_q}, /*dim=*/1); torch::Tensor joint_k = torch::cat({img_k, txt_k}, /*dim=*/1); torch::Tensor joint_v = torch::cat({img_v, txt_v}, /*dim=*/1); - const int64_t local_heads = heads_ / sequence_parallel_context.world_size(); + const int32_t sp_size = sp_group_ == nullptr ? 1 : sp_group_->world_size(); + const int64_t local_heads = heads_ / sp_size; torch::Tensor joint = joyimage_joint_attention( joint_q, joint_k, joint_v, local_heads, attn_mask); joint = joint.to(joint_q.dtype()); @@ -402,10 +426,20 @@ class JoyAttentionImpl final : public torch::nn::Module { /*start=*/image_sequence_length, /*end=*/joint.size(1)); - auto img_out_work = sequence_parallel_context.launch_head_to_sequence( - img_out, /*tensor_name=*/"hidden_states"); - auto txt_out_work = sequence_parallel_context.launch_head_to_sequence( - txt_out, /*tensor_name=*/"encoder_hidden_states"); + img_out = xllm::dit::SequenceParallelMixin::pad_tensor( + img_out, /*tensor_name=*/"hidden_states", /*dim=*/1); + auto img_out_work = parallel_state::all_to_all_4D(img_out, + /*scatter_idx=*/1, + /*gather_idx=*/2, + /*async_ops=*/true, + sp_group_); + txt_out = xllm::dit::SequenceParallelMixin::pad_tensor( + txt_out, /*tensor_name=*/"encoder_hidden_states", /*dim=*/1); + auto txt_out_work = parallel_state::all_to_all_4D(txt_out, + /*scatter_idx=*/1, + /*gather_idx=*/2, + /*async_ops=*/true, + sp_group_); img_out = img_out_work().flatten(/*start_dim=*/2, /*end_dim=*/3); img_out = img_attn_proj_->forward(img_out); @@ -456,6 +490,7 @@ class JoyAttentionImpl final : public torch::nn::Module { private: int64_t heads_; int64_t head_dim_; + ProcessGroup* sp_group_{nullptr}; layer::AddMatmulWeightTransposed img_attn_q_{nullptr}; layer::AddMatmulWeightTransposed img_attn_k_{nullptr}; layer::AddMatmulWeightTransposed img_attn_v_{nullptr}; @@ -521,6 +556,7 @@ class TransformerBlockImpl final : public torch::nn::Module { int64_t num_heads, int64_t head_dim, double mlp_width_ratio, + ProcessGroup* sp_group, double eps = 1e-6) : eps_(eps) { int64_t mlp_hidden = static_cast(dim * mlp_width_ratio); @@ -549,7 +585,7 @@ class TransformerBlockImpl final : public torch::nn::Module { 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", JoyAttention(context, dim, num_heads, head_dim, sp_group, eps)); (void)mlp_hidden; // FeedForward uses dim*4 internally (== dim*ratio) } @@ -567,8 +603,7 @@ class TransformerBlockImpl final : public torch::nn::Module { const torch::Tensor& temb, // [B, hidden] const torch::Tensor& rope_cos, const torch::Tensor& rope_sin, - const torch::Tensor& attn_mask, - xllm::dit::SequenceParallelContext& sequence_parallel_context) { + const torch::Tensor& attn_mask) { auto img_mod = img_mod_->forward(temb); // 6 x [B, hidden] auto txt_mod = txt_mod_->forward(temb); @@ -601,12 +636,8 @@ class TransformerBlockImpl final : public torch::nn::Module { txt_normed * (1 + txt_c1.unsqueeze(1)) + txt_s1.unsqueeze(1); #endif - JoyAttentionOutput attn_out = attn_->forward(img_modulated, - txt_modulated, - rope_cos, - rope_sin, - attn_mask, - sequence_parallel_context); + JoyAttentionOutput attn_out = attn_->forward( + img_modulated, txt_modulated, rope_cos, rope_sin, attn_mask); auto img_attn = std::get<0>(attn_out); auto txt_attn = std::get<1>(attn_out); @@ -725,8 +756,12 @@ class JoyImageEditPlusTransformer3DModelImpl final 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); + auto block = TransformerBlock(context, + hidden_size_, + num_heads_, + head_dim_, + mlp_width_ratio, + parallel_args.dit_sp_group_); double_blocks_->push_back(block); block_layers_.push_back(block); } @@ -760,8 +795,7 @@ class JoyImageEditPlusTransformer3DModelImpl final &attention_mask, use_cfg, step_index]( - const xllm::dit::SequenceParallelTensorMap& model_inputs, - xllm::dit::SequenceParallelContext& sequence_parallel_context) { + const xllm::dit::SequenceParallelTensorMap& model_inputs) { torch::Tensor output = forward_impl(model_inputs.at("hidden_states"), timestep, @@ -770,8 +804,7 @@ class JoyImageEditPlusTransformer3DModelImpl final rope_sin, attention_mask, use_cfg, - step_index, - sequence_parallel_context); + step_index); return xllm::dit::SequenceParallelTensorMap{ {"hidden_states", output}}; }); @@ -779,16 +812,14 @@ class JoyImageEditPlusTransformer3DModelImpl final } private: - torch::Tensor forward_impl( - const torch::Tensor& local_hidden_states, - const torch::Tensor& timestep, - const torch::Tensor& local_encoder_hidden_states, - const torch::Tensor& rope_cos, - const torch::Tensor& rope_sin, - const torch::Tensor& attention_mask, - bool use_cfg, - int64_t step_index, - xllm::dit::SequenceParallelContext& sequence_parallel_context) { + torch::Tensor forward_impl(const torch::Tensor& local_hidden_states, + const torch::Tensor& timestep, + const torch::Tensor& local_encoder_hidden_states, + const torch::Tensor& rope_cos, + const torch::Tensor& rope_sin, + const torch::Tensor& attention_mask, + bool use_cfg, + int64_t step_index) { int64_t batch_size = local_hidden_states.size(0); int64_t sequence_length = local_hidden_states.size(1); int64_t channels = local_hidden_states.size(2); @@ -835,14 +866,8 @@ class JoyImageEditPlusTransformer3DModelImpl final 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, - sequence_parallel_context); + std::tie(img, txt) = block_layers_[block_index]->forward( + img, txt, temb6, rope_cos, rope_sin, attention_mask); } TensorMap block_after_map = { diff --git a/xllm/models/dit/sequence_parallel/sequence_parallel_context.h b/xllm/models/dit/utils/sequence_parallel_mixin.h similarity index 56% rename from xllm/models/dit/sequence_parallel/sequence_parallel_context.h rename to xllm/models/dit/utils/sequence_parallel_mixin.h index d02c808c74..8fef8490c0 100644 --- a/xllm/models/dit/sequence_parallel/sequence_parallel_context.h +++ b/xllm/models/dit/utils/sequence_parallel_mixin.h @@ -19,7 +19,6 @@ limitations under the License. #include #include -#include #include #include #include @@ -29,23 +28,77 @@ limitations under the License. namespace xllm::dit { -using SequenceParallelWork = std::function; +using SequenceParallelTensorMap = + std::unordered_map; +using SequenceParallelTensorDims = std::unordered_map; -class SequenceParallelContext final { +class SequenceParallelMixin { public: - explicit SequenceParallelContext(ProcessGroup* process_group) - : process_group_(process_group) {} + static torch::Tensor pad_tensor(const torch::Tensor& input, + const std::string& tensor_name, + int64_t dim) { + if (!input.defined()) { + return input; + } + + return pad_right(input, padding_length(tensor_name), dim); + } + + static torch::Tensor unpad_tensor(const torch::Tensor& input, + const std::string& tensor_name, + int64_t dim) { + if (!input.defined()) { + return input; + } + + return unpad_right(input, padding_length(tensor_name), dim); + } + + protected: + SequenceParallelMixin(ProcessGroup* process_group, + SequenceParallelTensorDims input_sequence_dims, + SequenceParallelTensorDims output_sequence_dims) + : process_group_(process_group), + input_sequence_dims_(std::move(input_sequence_dims)), + output_sequence_dims_(std::move(output_sequence_dims)) {} + + template + SequenceParallelTensorMap sequence_parallel_forward( + const SequenceParallelTensorMap& inputs, + ForwardFn&& forward_fn) { + padding_lengths_.clear(); + SequenceParallelTensorMap local_inputs = inputs; + for (const auto& [tensor_name, sequence_dim] : input_sequence_dims_) { + auto tensor_it = local_inputs.find(tensor_name); + CHECK(tensor_it != local_inputs.end()) + << "Missing registered sequence-parallel input: " << tensor_name; + tensor_it->second = + scatter_sequence(tensor_it->second, tensor_name, sequence_dim); + } + + SequenceParallelTensorMap outputs = + std::forward(forward_fn)(local_inputs); + for (const auto& [tensor_name, sequence_dim] : output_sequence_dims_) { + auto tensor_it = outputs.find(tensor_name); + CHECK(tensor_it != outputs.end()) + << "Missing registered sequence-parallel output: " << tensor_name; + tensor_it->second = + gather_sequence(tensor_it->second, tensor_name, sequence_dim); + } + return outputs; + } + private: int32_t world_size() const { return process_group_ == nullptr ? 1 : process_group_->world_size(); } - bool enabled() const { return world_size() > 1; } + bool sequence_parallel_enabled() const { return world_size() > 1; } torch::Tensor scatter_sequence(const torch::Tensor& input, const std::string& tensor_name, int64_t sequence_dim) { - if (!enabled() || !input.defined()) { + if (!input.defined()) { return input; } @@ -53,7 +106,11 @@ class SequenceParallelContext final { const int64_t padding_length = (world_size() - sequence_length % world_size()) % world_size(); padding_lengths_[tensor_name] = padding_length; - torch::Tensor padded_input = pad_right(input, padding_length, sequence_dim); + if (!sequence_parallel_enabled()) { + return input; + } + + torch::Tensor padded_input = pad_tensor(input, tensor_name, sequence_dim); return parallel_state::scatter( padded_input, process_group_, static_cast(sequence_dim)); } @@ -61,51 +118,15 @@ class SequenceParallelContext final { torch::Tensor gather_sequence(const torch::Tensor& input, const std::string& tensor_name, int64_t sequence_dim) const { - if (!enabled() || !input.defined()) { + if (!sequence_parallel_enabled() || !input.defined()) { return input; } torch::Tensor output = parallel_state::gather( input.contiguous(), process_group_, static_cast(sequence_dim)); - return unpad_right(output, padding_length(tensor_name), sequence_dim); + return unpad_tensor(output, tensor_name, sequence_dim); } - SequenceParallelWork launch_sequence_to_head( - const torch::Tensor& input, - const std::string& tensor_name) const { - if (!enabled()) { - return [input]() { return input; }; - } - - SequenceParallelWork work = - parallel_state::all_to_all_4D(input, - /*scatter_idx=*/2, - /*gather_idx=*/1, - /*async_ops=*/true, - process_group_); - const int64_t padding = padding_length(tensor_name); - return [work = std::move(work), padding]() mutable { - return unpad_right(work(), padding, /*dim=*/1); - }; - } - - SequenceParallelWork launch_head_to_sequence( - const torch::Tensor& input, - const std::string& tensor_name) const { - if (!enabled()) { - return [input]() { return input; }; - } - - torch::Tensor padded_input = - pad_right(input, padding_length(tensor_name), /*dim=*/1); - return parallel_state::all_to_all_4D(padded_input, - /*scatter_idx=*/1, - /*gather_idx=*/2, - /*async_ops=*/true, - process_group_); - } - - private: static int64_t normalize_dim(const torch::Tensor& input, int64_t dim) { const int64_t normalized_dim = dim < 0 ? input.dim() + dim : dim; CHECK_GE(normalized_dim, 0) << "Invalid tensor dimension: " << dim; @@ -143,7 +164,7 @@ class SequenceParallelContext final { input.size(normalized_dim) - padding_length); } - int64_t padding_length(const std::string& tensor_name) const { + static int64_t padding_length(const std::string& tensor_name) { auto padding_it = padding_lengths_.find(tensor_name); CHECK(padding_it != padding_lengths_.end()) << "Missing sequence-parallel padding metadata: " << tensor_name; @@ -151,7 +172,9 @@ class SequenceParallelContext final { } ProcessGroup* process_group_{nullptr}; - std::unordered_map padding_lengths_; + const SequenceParallelTensorDims input_sequence_dims_; + const SequenceParallelTensorDims output_sequence_dims_; + inline static SequenceParallelTensorDims padding_lengths_; }; } // namespace xllm::dit