From 750ce7c67adca382eb5f0f489334832e6299bf6f Mon Sep 17 00:00:00 2001 From: ethan686 Date: Tue, 21 Jul 2026 22:05:36 +0800 Subject: [PATCH 1/4] refactor: add cfgparallelmixin for dit models. --- xllm/models/dit/pipelines/pipeline_wan_i2v.h | 58 ++++++------------ xllm/models/dit/utils/dit_parallel_mixin.h | 64 ++++++++++++++++++++ 2 files changed, 83 insertions(+), 39 deletions(-) create mode 100644 xllm/models/dit/utils/dit_parallel_mixin.h diff --git a/xllm/models/dit/pipelines/pipeline_wan_i2v.h b/xllm/models/dit/pipelines/pipeline_wan_i2v.h index f86ae3797e..99e2746673 100644 --- a/xllm/models/dit/pipelines/pipeline_wan_i2v.h +++ b/xllm/models/dit/pipelines/pipeline_wan_i2v.h @@ -39,11 +39,13 @@ 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()) { @@ -543,7 +545,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 +567,25 @@ 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, + auto [pos_noise, neg_noise] = forward_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, - encoded_negative_embeds, + 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(); + }, + parallel_args_.dit_cfg_group_, + /*do_cfg=*/true); + + 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..2ae62bb159 --- /dev/null +++ b/xllm/models/dit/utils/dit_parallel_mixin.h @@ -0,0 +1,64 @@ +/* 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/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: + // Returns {positive_noise_pred, negative_noise_pred}. + template + std::pair forward_cfg(ForwardFn&& forward_fn, + ProcessGroup* pg, + bool do_cfg) const { + if (!do_cfg) { + return {forward_fn(true), torch::Tensor()}; + } + + if (pg != nullptr && pg->world_size() > 1) { + int32_t rank = pg->rank(); + torch::Tensor noise_pred = forward_fn(rank == 0); + torch::Tensor gathered = + parallel_state::gather(noise_pred, pg, /*dim=*/0); + auto chunks = torch::chunk(gathered, 2, 0); + return {chunks[0], chunks[1]}; + } + + return {forward_fn(true), forward_fn(false)}; + } +}; + +// Mixin for VAE parallelism (to be implemented). +class VaeParallelMixin {}; + +// Mixin for sequence parallelism (to be implemented). +class SpParallelMixin {}; + +} // namespace dit +} // namespace xllm From 84d18e3991fedcb99da446e683f7cc9928dd3fb3 Mon Sep 17 00:00:00 2001 From: ethan686 Date: Wed, 22 Jul 2026 21:41:59 +0800 Subject: [PATCH 2/4] fix: resolve the review issue. --- xllm/models/dit/pipelines/pipeline_wan_i2v.h | 27 +++++++------- xllm/models/dit/utils/dit_parallel_mixin.h | 38 ++++++++++++-------- 2 files changed, 35 insertions(+), 30 deletions(-) diff --git a/xllm/models/dit/pipelines/pipeline_wan_i2v.h b/xllm/models/dit/pipelines/pipeline_wan_i2v.h index 99e2746673..3bf0830083 100644 --- a/xllm/models/dit/pipelines/pipeline_wan_i2v.h +++ b/xllm/models/dit/pipelines/pipeline_wan_i2v.h @@ -48,7 +48,8 @@ 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(); @@ -567,20 +568,16 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module, #endif if (do_classifier_free_guidance) { - auto [pos_noise, neg_noise] = forward_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); - }, - parallel_args_.dit_cfg_group_, - /*do_cfg=*/true); + 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) + diff --git a/xllm/models/dit/utils/dit_parallel_mixin.h b/xllm/models/dit/utils/dit_parallel_mixin.h index 2ae62bb159..0138a10a61 100644 --- a/xllm/models/dit/utils/dit_parallel_mixin.h +++ b/xllm/models/dit/utils/dit_parallel_mixin.h @@ -19,6 +19,8 @@ limitations under the License. #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" @@ -32,26 +34,32 @@ namespace dit { // 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 forward_cfg(ForwardFn&& forward_fn, - ProcessGroup* pg, - bool do_cfg) const { - if (!do_cfg) { - return {forward_fn(true), torch::Tensor()}; - } - - if (pg != nullptr && pg->world_size() > 1) { - int32_t rank = pg->rank(); - torch::Tensor noise_pred = forward_fn(rank == 0); - torch::Tensor gathered = - parallel_state::gather(noise_pred, pg, /*dim=*/0); - auto chunks = torch::chunk(gathered, 2, 0); - return {chunks[0], chunks[1]}; + 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)}; } - 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). From 112a97daf642c7fa0ca139aa2fe312695ccefe4b Mon Sep 17 00:00:00 2001 From: ethan686 Date: Thu, 23 Jul 2026 11:52:33 +0800 Subject: [PATCH 3/4] fix: resolve the review issue. --- xllm/models/dit/utils/dit_parallel_mixin.h | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/xllm/models/dit/utils/dit_parallel_mixin.h b/xllm/models/dit/utils/dit_parallel_mixin.h index 0138a10a61..ef49d26aca 100644 --- a/xllm/models/dit/utils/dit_parallel_mixin.h +++ b/xllm/models/dit/utils/dit_parallel_mixin.h @@ -19,7 +19,6 @@ limitations under the License. #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" @@ -35,20 +34,29 @@ namespace dit { class CFGParallelMixin { public: explicit CFGParallelMixin(const DiTModelContext& context) - : cfg_size_(ParallelConfig::get_instance().cfg_size()), - cfg_group_(context.get_parallel_args().dit_cfg_group_) {} + : 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) { + int32_t cfg_size = 1; + if (cfg_group_ != nullptr) { + cfg_size = cfg_group_->world_size(); + } + + TORCH_CHECK(cfg_size == 1 || cfg_size == 2, + "CFGParallelMixin only supports cfg_size 1 or 2, got ", + cfg_size); + + // Serial execution: evaluate positive and negative conditionals one by one. + if (cfg_size == 1) { return {forward_fn(true), forward_fn(false)}; } - // CFG parallel: rank 0 → positive, rank 1 → negative, gather + chunk. + // CFG parallel (cfg_size == 2): 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 = @@ -58,7 +66,6 @@ class CFGParallelMixin { } private: - int32_t cfg_size_ = 1; ProcessGroup* cfg_group_ = nullptr; }; From 35d6cf34680989800cfe720d3230ae64fbd24ad8 Mon Sep 17 00:00:00 2001 From: ethan686 Date: Thu, 23 Jul 2026 14:21:14 +0800 Subject: [PATCH 4/4] fix: change torch_check to check. --- xllm/models/dit/utils/dit_parallel_mixin.h | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xllm/models/dit/utils/dit_parallel_mixin.h b/xllm/models/dit/utils/dit_parallel_mixin.h index ef49d26aca..cbbb52b00e 100644 --- a/xllm/models/dit/utils/dit_parallel_mixin.h +++ b/xllm/models/dit/utils/dit_parallel_mixin.h @@ -46,9 +46,7 @@ class CFGParallelMixin { cfg_size = cfg_group_->world_size(); } - TORCH_CHECK(cfg_size == 1 || cfg_size == 2, - "CFGParallelMixin only supports cfg_size 1 or 2, got ", - cfg_size); + CHECK(cfg_size == 1 || cfg_size == 2); // Serial execution: evaluate positive and negative conditionals one by one. if (cfg_size == 1) {