diff --git a/xllm/models/dit/pipelines/pipeline_wan_i2v.h b/xllm/models/dit/pipelines/pipeline_wan_i2v.h index f86ae3797e..3bf0830083 100644 --- a/xllm/models/dit/pipelines/pipeline_wan_i2v.h +++ b/xllm/models/dit/pipelines/pipeline_wan_i2v.h @@ -39,14 +39,17 @@ limitations under the License. #if defined(USE_NPU) #include "models/dit/utils/dit_block_weight_manager.h" #endif +#include "models/dit/utils/dit_parallel_mixin.h" #include "models/model_registry.h" namespace xllm { -class WanImageToVideoPipelineImpl : public torch::nn::Module { +class WanImageToVideoPipelineImpl : public torch::nn::Module, + public dit::CFGParallelMixin { public: WanImageToVideoPipelineImpl(const DiTModelContext& context) - : parallel_args_(context.get_parallel_args()) { + : parallel_args_(context.get_parallel_args()), + dit::CFGParallelMixin(context) { options_ = context.get_tensor_options(); const auto& vae_args = context.get_model_args("vae"); zdim_ = vae_args.z_dim(); @@ -543,7 +546,6 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { } } torch::Tensor noise_pred; - torch::Tensor noise_uncond; #if defined(USE_NPU) auto& rolling = (current_model.get() == transformer_.get()) ? rolling_transformer_ @@ -566,46 +568,21 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { #endif if (do_classifier_free_guidance) { - if (ParallelConfig::get_instance().cfg_size() == 2) { - int32_t rank = parallel_args_.dit_cfg_group_->rank(); - noise_pred = - use_rolling_load_ - ? rolling_forward(rank == 0 ? encoded_prompt_embeds - : encoded_negative_embeds) - : current_model->forward(latent_model_input, - timestep_input, - rank == 0 ? encoded_prompt_embeds - : encoded_negative_embeds, - torch::Tensor(), - sparse_attn_state); - auto gathered = xllm::parallel_state::gather( - noise_pred, parallel_args_.dit_cfg_group_, /*dim=*/0); - auto chunks = torch::chunk(gathered, 2, 0); - noise_pred = chunks[0]; - noise_uncond = chunks[1]; - } else { - if (use_rolling_load_) { - noise_pred = rolling_forward(encoded_prompt_embeds); - noise_uncond = rolling_forward(encoded_negative_embeds); - } else { - noise_pred = current_model->forward(latent_model_input, - timestep_input, - encoded_prompt_embeds, - torch::Tensor(), - sparse_attn_state); - noise_uncond = current_model->forward(latent_model_input, - timestep_input, - encoded_negative_embeds, - torch::Tensor(), - sparse_attn_state); - } - } - - noise_pred = noise_uncond.to(torch::kFloat32) + - static_cast(current_guidance) * - (noise_pred.to(torch::kFloat32) - - noise_uncond.to(torch::kFloat32)); - noise_uncond.reset(); + auto [pos_noise, neg_noise] = exec_with_cfg([&](bool is_positive) { + auto& embeds = + is_positive ? encoded_prompt_embeds : encoded_negative_embeds; + return use_rolling_load_ ? rolling_forward(embeds) + : current_model->forward(latent_model_input, + timestep_input, + embeds, + torch::Tensor(), + sparse_attn_state); + }); + + noise_pred = + neg_noise.to(torch::kFloat32) + + static_cast(current_guidance) * + (pos_noise.to(torch::kFloat32) - neg_noise.to(torch::kFloat32)); } else { noise_pred = use_rolling_load_ ? rolling_forward(encoded_prompt_embeds) diff --git a/xllm/models/dit/utils/dit_parallel_mixin.h b/xllm/models/dit/utils/dit_parallel_mixin.h new file mode 100644 index 0000000000..0138a10a61 --- /dev/null +++ b/xllm/models/dit/utils/dit_parallel_mixin.h @@ -0,0 +1,72 @@ +/* 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 "core/framework/config/parallel_config.h" +#include "core/framework/model_context.h" +#include "core/framework/parallel_state/process_group.h" +#include "framework/parallel_state/parallel_state.h" + +namespace xllm { +namespace dit { + +// Mixin providing classifier-free guidance (CFG) parallelism. +// +// Usage: +// class MyPipeline : public torch::nn::Module, +// public dit::CFGParallelMixin { ... }; +class CFGParallelMixin { + public: + explicit CFGParallelMixin(const DiTModelContext& context) + : cfg_size_(ParallelConfig::get_instance().cfg_size()), + cfg_group_(context.get_parallel_args().dit_cfg_group_) {} + + // forward_fn(is_positive) -> Tensor — caller captures embeddings in lambda. + // Returns {positive_noise_pred, negative_noise_pred}. + template + std::pair exec_with_cfg( + const ForwardFn& forward_fn) const { + // Not doing CFG parallel: execute both conditionals serially. + if (cfg_group_ == nullptr || cfg_group_->world_size() != 2) { + return {forward_fn(true), forward_fn(false)}; + } + + // CFG parallel: rank 0 → positive, rank 1 → negative, gather + chunk. + int32_t rank = cfg_group_->rank(); + torch::Tensor noise_pred = forward_fn(rank == 0); + torch::Tensor gathered = + parallel_state::gather(noise_pred, cfg_group_, /*dim=*/0); + auto chunks = torch::chunk(gathered, 2, 0); + return {chunks[0], chunks[1]}; + } + + private: + int32_t cfg_size_ = 1; + ProcessGroup* cfg_group_ = nullptr; +}; + +// Mixin for VAE parallelism (to be implemented). +class VaeParallelMixin {}; + +// Mixin for sequence parallelism (to be implemented). +class SpParallelMixin {}; + +} // namespace dit +} // namespace xllm