diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 46b10f196b..1a32115e91 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -374,6 +374,8 @@ DECLARE_string(dit_sparse_attention_version); DECLARE_int64(dit_sparse_attention_mask_refresh_steps); +DECLARE_bool(dit_distill_enable); + DECLARE_bool(use_audio_in_video); // --- kernel config --- diff --git a/xllm/core/framework/config/dit_config.cpp b/xllm/core/framework/config/dit_config.cpp index 585e219f6c..1a85066a89 100644 --- a/xllm/core/framework/config/dit_config.cpp +++ b/xllm/core/framework/config/dit_config.cpp @@ -61,6 +61,12 @@ DEFINE_bool(dit_debug_print, false, "whether print the debug info for dit models"); +DEFINE_bool(dit_distill_enable, + false, + "whether the DiT weights are distilled (Wan2.2 I2V): true selects " + "the FlowMatch scheduler + torch RMSNorm for norm_q/k; false uses " + "UniPC + the aclnn fused kernel."); + DEFINE_int64(dit_generation_image_area_max, 0, "Maximum allowed image area (width * height) for image generation " @@ -123,6 +129,7 @@ void DiTConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_blocks); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_distill_enable); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_generation_image_area_max); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_vae_image_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_enable_vae_tiling); @@ -147,6 +154,7 @@ void DiTConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_blocks); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_distill_enable); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_generation_image_area_max); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_vae_image_size); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_enable_vae_tiling); @@ -184,6 +192,8 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const { config_json, default_config, dit_sp_communication_overlap); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_debug_print); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_distill_enable); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_generation_image_area_max); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/dit_config.h b/xllm/core/framework/config/dit_config.h index c687ba11da..67ead78248 100644 --- a/xllm/core/framework/config/dit_config.h +++ b/xllm/core/framework/config/dit_config.h @@ -53,6 +53,7 @@ class DiTConfig final { "dit_cache_end_blocks", "dit_sp_communication_overlap", "dit_debug_print", + "dit_distill_enable", "dit_generation_image_area_max", "dit_vae_image_size", "dit_enable_vae_tiling", @@ -89,6 +90,8 @@ class DiTConfig final { PROPERTY(bool, dit_debug_print) = false; + PROPERTY(bool, dit_distill_enable) = false; + PROPERTY(int64_t, dit_generation_image_area_max) = 0; PROPERTY(int64_t, dit_vae_image_size) = 1048576; diff --git a/xllm/models/dit/pipelines/pipeline_wan_i2v.h b/xllm/models/dit/pipelines/pipeline_wan_i2v.h index ecc895e754..379c3df6ce 100644 --- a/xllm/models/dit/pipelines/pipeline_wan_i2v.h +++ b/xllm/models/dit/pipelines/pipeline_wan_i2v.h @@ -34,6 +34,7 @@ limitations under the License. #include "models/dit/autoencoders/autoencoder_kl_wan.h" #include "models/dit/encoders/umt5_encoder.h" #include "models/dit/processors/vae_video_processor.h" +#include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h" #include "models/dit/schedulers/uni_pc_multi_step_scheduler.h" #include "models/dit/transformers/transformer_wan.h" #if defined(USE_NPU) @@ -52,7 +53,6 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { zdim_ = vae_args.z_dim(); latents_mean_ = vae_args.latents_mean(); latents_std_ = vae_args.latents_std(); - const auto& scheduler_args = context.get_model_args("scheduler"); num_train_timesteps_ = scheduler_args.num_train_timesteps(); @@ -70,6 +70,11 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { rf_config_.mask_refresh_steps = dit_config.dit_sparse_attention_mask_refresh_steps(); + // is_distill from --dit_distill_enable flag: true selects FlowMatch + // scheduler + torch RMSNorm for norm_q/k; false uses UniPC + aclnn + // fused kernel. (The transformer reads the same flag from DiTConfig.) + is_distill_ = DiTConfig::get_instance().dit_distill_enable(); + LOG(INFO) << "Initializing Wan2_2I2V pipeline..."; vae_ = AutoencoderKLWan(context.get_model_context("vae")); transformer_ = WanTransformer3DModel( @@ -77,7 +82,9 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { transformer_2_ = WanTransformer3DModel( context.get_model_context("transformer_2"), rf_config_); umt5_ = UMT5EncoderModel(context.get_model_context("text_encoder")); - scheduler_ = + flow_scheduler_ = + FlowMatchEulerDiscreteScheduler(context.get_model_context("scheduler")); + unipc_scheduler_ = UniPCMultistepScheduler(context.get_model_context("scheduler")); video_processor_ = VAEVideoProcessor(context.get_model_context("vae"), true, @@ -91,7 +98,8 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { register_module("transformer", transformer_); register_module("transformer_2", transformer_2_); register_module("umt5", umt5_); - register_module("scheduler", scheduler_); + register_module("flow_scheduler", flow_scheduler_); + register_module("unipc_scheduler", unipc_scheduler_); register_module("video_processor_", video_processor_); } @@ -146,7 +154,7 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { void load_model(std::unique_ptr loader) { LOG(INFO) << "Wan2_2I2VPipeline loading model from" - << loader->model_root_path(); + << loader->model_root_path() << " is_distill=" << is_distill_; auto transformer_loader = loader->take_component_loader("transformer"); auto transformer_2_loader = loader->take_component_loader("transformer_2"); auto vae_loader = loader->take_component_loader("vae"); @@ -447,11 +455,27 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { do_classifier_free_guidance, num_videos_per_prompt, max_sequence_length); - scheduler_->set_timesteps(num_inference_steps, - options_.device(), - /*sigmas*/ std::nullopt, - /*mu*/ std::nullopt); - torch::Tensor timesteps = scheduler_->timesteps(); + torch::Tensor timesteps; + if (is_distill_) { + // Explicit raw sigmas (N-i)/N = [1,0.75,0.5,0.25] to match LightX2V. + std::vector raw_sigmas(num_inference_steps); + for (int i = 0; i < num_inference_steps; ++i) { + raw_sigmas[i] = static_cast(num_inference_steps - i) / + static_cast(num_inference_steps); + } + flow_scheduler_->set_timesteps(num_inference_steps, + options_.device(), + /*sigmas*/ raw_sigmas, + /*mu*/ std::nullopt); + timesteps = flow_scheduler_->timesteps().to(options_.device()); + } else { + // Original: UniPC computes its own sigma schedule. + unipc_scheduler_->set_timesteps(num_inference_steps, + options_.device(), + /*sigmas*/ std::nullopt, + /*mu*/ std::nullopt); + timesteps = unipc_scheduler_->timesteps().to(options_.device()); + } int64_t num_channels_latents = zdim_; torch::Tensor input_image; @@ -504,7 +528,6 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { for (int64_t i = 0; i < timesteps.numel(); ++i) { torch::Tensor t = timesteps[i]; - int64_t total_steps = timesteps.numel(); WanTransformer3DModel current_model = nullptr; float current_guidance; @@ -615,7 +638,9 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { torch::Tensor(), rf_state); } - auto prev_latents = scheduler_->step(noise_pred, t, prepared_latents); + auto prev_latents = + is_distill_ ? flow_scheduler_->step(noise_pred, t, prepared_latents) + : unipc_scheduler_->step(noise_pred, t, prepared_latents); prepared_latents = prev_latents.detach(); noise_pred.reset(); prev_latents = torch::Tensor(); @@ -736,7 +761,9 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { } #endif - UniPCMultistepScheduler scheduler_{nullptr}; + // distill mode uses flow_scheduler_; otherwise unipc_scheduler_. + FlowMatchEulerDiscreteScheduler flow_scheduler_{nullptr}; + UniPCMultistepScheduler unipc_scheduler_{nullptr}; AutoencoderKLWan vae_{nullptr}; WanTransformer3DModel transformer_{nullptr}; WanTransformer3DModel transformer_2_{nullptr}; @@ -755,6 +782,7 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module { torch::TensorOptions options_; const ParallelArgs parallel_args_; xllm::dit::RainFusionConfig rf_config_; + bool is_distill_{false}; }; TORCH_MODULE(WanImageToVideoPipeline); diff --git a/xllm/models/dit/transformers/transformer_wan.h b/xllm/models/dit/transformers/transformer_wan.h index e123150be2..f6c14c44e4 100644 --- a/xllm/models/dit/transformers/transformer_wan.h +++ b/xllm/models/dit/transformers/transformer_wan.h @@ -27,6 +27,7 @@ limitations under the License. #include #include +#include "core/framework/config/dit_config.h" #include "core/framework/config/load_config.h" #include "core/framework/config/parallel_config.h" #include "core/framework/dit_model_loader.h" @@ -758,8 +759,23 @@ class WanAttentionImpl : public torch::nn::Module { query = dit::tp_rms_norm(query, norm_q_, parallel_args_.dit_tp_group_); key = dit::tp_rms_norm(key, norm_k_, parallel_args_.dit_tp_group_); } else { - query = std::get<0>(norm_q_->forward(query)); - key = std::get<0>(norm_k_->forward(key)); + // Distill uses per-op torch RMSNorm (matches lightx2v); else fused + // kernel. + if (DiTConfig::get_instance().dit_distill_enable() && is_self_attention) { + auto torch_rms = [](const torch::Tensor& x, + const torch::Tensor& w, + double eps) { + // x * rsqrt(mean(x^2, -1) + eps) * w; w.to(x) for rolling-load + // weights. + auto var = x.pow(2).mean(-1, /*keepdim=*/true); + return (x * torch::rsqrt(var + eps)) * w.to(x.device(), x.dtype()); + }; + query = torch_rms(query, norm_q_->weight(), norm_q_->eps()); + key = torch_rms(key, norm_k_->weight(), norm_k_->eps()); + } else { + query = std::get<0>(norm_q_->forward(query)); + key = std::get<0>(norm_k_->forward(key)); + } } // ── Step 3: SP all2all for Q/K (V already done in layer) ── @@ -1032,9 +1048,8 @@ class WanTimeTextImageEmbeddingImpl : public torch::nn::Module { timestep_proj = timesteps_proj_->forward(ts).view({-1, seq_len, time_freq_dim_}); } - timestep_proj = timestep_proj.to(torch::kFloat32); - auto embed_dtype = encoder_hidden_states.dtype(); - torch::Tensor temb = time_embedder_->forward(timestep_proj.to(embed_dtype)); + // bf16-direct temb for both modes (fp32 round-trip gives no benefit). + torch::Tensor temb = time_embedder_->forward(timestep_proj); torch::Tensor timestep_proj_out = time_proj_->forward(act_fn_->forward(temb)); if (seq_len > 1) { @@ -1538,7 +1553,6 @@ class WanTransformer3DModelImpl : public torch::nn::Module { } else { timestep_proj = timestep_proj.view({batch_size, 6, -1}); } - if (encoder_hidden_states_image_embedded.defined()) { encoder_hidden_states_embedded = torch::cat({encoder_hidden_states_image_embedded, @@ -1592,12 +1606,9 @@ class WanTransformer3DModelImpl : public torch::nn::Module { shift = shift.to(hidden_states.device()); scale = scale.to(hidden_states.device()); - auto norm_result = norm_out_->forward(hidden_states, /*keep_fp32*/ true); - auto one_plus_scale = - (1 + scale.to(hidden_states.dtype())).to(torch::kFloat32); - auto shift_fp32 = shift.to(torch::kFloat32); - auto norm_out = norm_result * one_plus_scale + shift_fp32; - hidden_states = norm_out.to(hidden_states.dtype()); + // bf16 for both modes (fp32 round-trip gives no benefit). + auto norm_result = norm_out_->forward(hidden_states, /*keep_fp32*/ false); + hidden_states = norm_result * (1 + scale) + shift; if (::xllm::ParallelConfig::get_instance().sp_size() > 1 && seq_len != pad_seq_len) {