refactor: add cfgparallelmixin for dit models.

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<float>(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<float>(current_guidance) *
+                (pos_noise.to(torch::kFloat32) - neg_noise.to(torch::kFloat32));
       } else {
         noise_pred = use_rolling_load_
                          ? rolling_forward(encoded_prompt_embeds)

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 <torch/torch.h>
+
+#include <utility>
+
+#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 <typename ForwardFn>
+  std::pair<torch::Tensor, torch::Tensor> 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