From 21191656fbe8b6e2bc8e00aafdc90560c58adb7b Mon Sep 17 00:00:00 2001 From: BurnWan Date: Thu, 30 Jul 2026 15:34:30 +0800 Subject: [PATCH 1/3] feat: add RegionE support for Qwen-Image-Edit. --- xllm/core/common/global_flags.h | 3 + xllm/core/framework/config/dit_config.cpp | 18 +- xllm/core/framework/config/dit_config.h | 6 + xllm/core/framework/dit_cache/CMakeLists.txt | 3 + xllm/core/framework/dit_cache/dit_cache.cpp | 6 + xllm/core/framework/dit_cache/dit_cache.h | 19 +- .../framework/dit_cache/dit_cache_config.h | 17 +- .../framework/dit_cache/dit_cache_impl.cpp | 2 + xllm/core/framework/dit_cache/regione.cpp | 701 ++++++++++++++++++ xllm/core/framework/dit_cache/regione.h | 220 ++++++ xllm/core/runtime/dit_worker_impl.cpp | 24 + .../pipelines/pipeline_qwenimage_edit_plus.h | 99 ++- .../dit/transformers/transformer_qwen_image.h | 41 +- 13 files changed, 1132 insertions(+), 27 deletions(-) create mode 100644 xllm/core/framework/dit_cache/regione.cpp create mode 100644 xllm/core/framework/dit_cache/regione.h diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index d9071c6e69..daebb8b4ce 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -360,6 +360,9 @@ DECLARE_int64(dit_cache_start_blocks); DECLARE_int64(dit_cache_end_blocks); +DECLARE_string(dit_regione_refresh_steps); +DECLARE_double(dit_regione_region_threshold); + DECLARE_bool(dit_sp_communication_overlap); DECLARE_int64(dit_generation_image_area_max); diff --git a/xllm/core/framework/config/dit_config.cpp b/xllm/core/framework/config/dit_config.cpp index c213c36e32..a60a26213c 100644 --- a/xllm/core/framework/config/dit_config.cpp +++ b/xllm/core/framework/config/dit_config.cpp @@ -23,7 +23,7 @@ DEFINE_int32(max_requests_per_batch, 1, "Max number of request per batch."); DEFINE_string(dit_cache_policy, "TaylorSeer", "The policy of dit cache(e.g. None, FBCache, TaylorSeer, " - "FBCacheTaylorSeer, ResidualCache)."); + "FBCacheTaylorSeer, ResidualCache, RegionE)."); DEFINE_int64(dit_cache_warmup_steps, 0, "The number of warmup steps."); @@ -53,6 +53,14 @@ DEFINE_int64(dit_cache_end_blocks, 5, "The number of blocks to skip at the end."); +DEFINE_string(dit_regione_refresh_steps, + "16", + "RegionE: comma-separated full-image refresh steps in RAGS."); + +DEFINE_double(dit_regione_region_threshold, + 0.80, + "RegionE: cosine threshold for adaptive region partition."); + DEFINE_bool(dit_sp_communication_overlap, true, "Communication & Computation overlap for sequence parallel"); @@ -131,6 +139,8 @@ void DiTConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_steps); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_start_blocks); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_blocks); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_refresh_steps); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_region_threshold); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_laser_attention_enabled); @@ -157,6 +167,8 @@ void DiTConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_steps); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_start_blocks); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_blocks); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_refresh_steps); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_region_threshold); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_laser_attention_enabled); @@ -194,6 +206,10 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const { config_json, default_config, dit_cache_start_blocks); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_cache_end_blocks); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_regione_refresh_steps); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_regione_region_threshold); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_sp_communication_overlap); 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 6eaa3209f1..98beb313eb 100644 --- a/xllm/core/framework/config/dit_config.h +++ b/xllm/core/framework/config/dit_config.h @@ -51,6 +51,8 @@ class DiTConfig final { "dit_cache_end_steps", "dit_cache_start_blocks", "dit_cache_end_blocks", + "dit_regione_refresh_steps", + "dit_regione_region_threshold", "dit_sp_communication_overlap", "dit_debug_print", "dit_laser_attention_enabled", @@ -87,6 +89,10 @@ class DiTConfig final { PROPERTY(int64_t, dit_cache_end_blocks) = 5; + PROPERTY(std::string, dit_regione_refresh_steps) = "16"; + + PROPERTY(double, dit_regione_region_threshold) = 0.80; + PROPERTY(bool, dit_sp_communication_overlap) = true; PROPERTY(bool, dit_debug_print) = false; diff --git a/xllm/core/framework/dit_cache/CMakeLists.txt b/xllm/core/framework/dit_cache/CMakeLists.txt index e5e6ce2744..326c6caaa6 100644 --- a/xllm/core/framework/dit_cache/CMakeLists.txt +++ b/xllm/core/framework/dit_cache/CMakeLists.txt @@ -13,6 +13,7 @@ cc_library( fbcache_taylorseer.h taylorseer.h residual_cache.h + regione.h SRCS dit_cache_impl.cpp dit_cache.cpp @@ -21,8 +22,10 @@ cc_library( fbcache_taylorseer.cpp taylorseer.cpp residual_cache.cpp + regione.cpp DEPS torch + $<$:torch_npu> glog::glog Folly::folly parallel_state diff --git a/xllm/core/framework/dit_cache/dit_cache.cpp b/xllm/core/framework/dit_cache/dit_cache.cpp index 738ef136fa..f33c1cf74b 100644 --- a/xllm/core/framework/dit_cache/dit_cache.cpp +++ b/xllm/core/framework/dit_cache/dit_cache.cpp @@ -19,6 +19,12 @@ namespace xllm { bool DiTCache::init(const DiTCacheConfig& cfg, const ParallelArgs& parallel_args) { + regione_cache_.reset(); + if (cfg.selected_policy == PolicyType::RegionE) { + regione_cache_ = std::make_unique(); + regione_cache_->init(cfg); + } + active_cache_ = create_dit_cache(cfg); active_cond_cache_ = create_dit_cache(cfg); if (!active_cache_ || !active_cond_cache_) { diff --git a/xllm/core/framework/dit_cache/dit_cache.h b/xllm/core/framework/dit_cache/dit_cache.h index a0f34bfaa9..073ebe0b1d 100644 --- a/xllm/core/framework/dit_cache/dit_cache.h +++ b/xllm/core/framework/dit_cache/dit_cache.h @@ -14,7 +14,12 @@ limitations under the License. ==============================================================================*/ #pragma once + +#include +#include + #include "dit_cache_impl.h" +#include "regione.h" namespace xllm { @@ -36,24 +41,30 @@ class DiTCache { bool init(const DiTCacheConfig& cfg, const ParallelArgs& parallel_args); bool on_before_block(const CacheBlockIn& blockin, bool use_cfg = false); - CacheBlockOut on_after_block(const CacheBlockIn& blockin, bool use_cfg = false); - bool on_before_step(const CacheStepIn& stepin, bool use_cfg = false); - CacheStepOut on_after_step(const CacheStepIn& stepin, bool use_cfg = false); void set_context(const CacheContext& context) { + if (regione_cache_) { + regione_cache_->set_infer_steps(context.infer_steps); + regione_cache_->set_num_blocks(context.num_blocks); + } active_cache_->set_context(context); active_cond_cache_->set_context(context); } + RegionECache* regione() { return regione_cache_.get(); } + const RegionECache* regione() const { return regione_cache_.get(); } + private: - torch::Tensor get_tensor_or_empty(const TensorMap& m, const std::string& k); + static torch::Tensor get_tensor_or_empty(const TensorMap& m, + const std::string& k); std::unique_ptr active_cache_; std::unique_ptr active_cond_cache_; + std::unique_ptr regione_cache_; }; } // namespace xllm diff --git a/xllm/core/framework/dit_cache/dit_cache_config.h b/xllm/core/framework/dit_cache/dit_cache_config.h index 217f989cd7..42d55202f5 100644 --- a/xllm/core/framework/dit_cache/dit_cache_config.h +++ b/xllm/core/framework/dit_cache/dit_cache_config.h @@ -15,6 +15,10 @@ limitations under the License. #pragma once +#include +#include +#include + namespace xllm { enum class PolicyType { @@ -22,7 +26,8 @@ enum class PolicyType { FBCache, TaylorSeer, FBCacheTaylorSeer, - ResidualCache + ResidualCache, + RegionE }; struct DiTBaseCacheOptions { @@ -51,6 +56,13 @@ struct FBCacheTaylorSeerOptions : public DiTBaseCacheOptions { int n_derivatives = 3; }; +struct RegionEOptions : public DiTBaseCacheOptions { + int64_t skip_interval_steps = 3; + int64_t tail_steps = 1; + std::vector refresh_steps = {16}; + float region_threshold = 0.80f; +}; + struct ResidualCacheOptions { // The number of steps to skip at the start. int64_t dit_cache_start_steps = 5; @@ -85,6 +97,9 @@ struct DiTCacheConfig { // the configuration for ResidualCache policy. ResidualCacheOptions residual_cache; + + // the configuration for RegionE policy. + RegionEOptions regione; }; } // namespace xllm diff --git a/xllm/core/framework/dit_cache/dit_cache_impl.cpp b/xllm/core/framework/dit_cache/dit_cache_impl.cpp index 60d52d2fed..426ba4975f 100644 --- a/xllm/core/framework/dit_cache/dit_cache_impl.cpp +++ b/xllm/core/framework/dit_cache/dit_cache_impl.cpp @@ -82,6 +82,8 @@ std::unique_ptr create_dit_cache(const DiTCacheConfig& cfg) { return std::make_unique(); case PolicyType::ResidualCache: return std::make_unique(); + case PolicyType::RegionE: + return std::make_unique(); default: return std::make_unique(); } diff --git a/xllm/core/framework/dit_cache/regione.cpp b/xllm/core/framework/dit_cache/regione.cpp new file mode 100644 index 0000000000..9c09b858b5 --- /dev/null +++ b/xllm/core/framework/dit_cache/regione.cpp @@ -0,0 +1,701 @@ +/* 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. +==============================================================================*/ + +#include "regione.h" + +#include + +#include + +namespace xllm { +namespace { +bool tensor_cache_ready(const std::vector& cache, + int64_t block_id) { + return block_id >= 0 && block_id < static_cast(cache.size()) && + cache[block_id].defined(); +} + +torch::Tensor regione_to_cpu_cache(const torch::Tensor& tensor) { + auto src = tensor.detach().contiguous(); + auto cpu_options = src.options().device(torch::kCPU).pinned_memory(true); + auto cpu_tensor = torch::empty(src.sizes(), cpu_options); + cpu_tensor.copy_(src, /*non_blocking=*/false); + return cpu_tensor; +} +} // namespace + +void RegionECache::init(const DiTCacheConfig& cfg) { + config_ = cfg; + regione_enabled_ = cfg.selected_policy == PolicyType::RegionE; + regione_velocity_cache_ = torch::Tensor(); + regione_current_block_ = -1; + regione_current_use_cfg_ = false; + regione_current_step_ = 0; + regione_infer_steps_ = 0; + regione_num_blocks_ = 0; + regione_partial_mode_ = false; + regione_target_seq_len_ = 0; + regione_grid_h_ = 0; + regione_grid_w_ = 0; + regione_image_seq_len_ = 0; + regione_sp_rank_ = 0; + regione_sp_size_ = 1; + regione_local_start_ = 0; + regione_local_end_ = 0; + regione_condition_latents_ = torch::Tensor(); + regione_edited_ids_ = torch::Tensor(); + regione_unedited_ids_ = torch::Tensor(); + regione_local_edited_global_ids_ = torch::Tensor(); + regione_local_edited_cache_ids_ = torch::Tensor(); + regione_local_image_global_ids_ = torch::Tensor(); + regione_clear_all_prefetch_slots(); +} + +void RegionECache::set_num_blocks(int64_t num_blocks) { + regione_num_blocks_ = num_blocks; + ensure_regione_kv_size(num_blocks); +} + +bool RegionECache::regione_is_refresh_step(int64_t step) const { + for (const auto refresh_step : config_.regione.refresh_steps) { + const auto refresh_index = + refresh_step > 0 ? refresh_step - 1 : refresh_step; + if (step == refresh_index) return true; + } + return false; +} + +bool RegionECache::regione_is_tail_step(int64_t step) const { + return config_.regione.tail_steps > 0 && regione_infer_steps_ > 0 && + step >= regione_infer_steps_ - config_.regione.tail_steps; +} + +bool RegionECache::regione_should_run_full_step(int64_t step) const { + if (!regione_enabled_) return true; + if (!regione_has_regions()) return true; + if (step < config_.regione.warmup_steps) return true; + if (regione_is_tail_step(step)) return true; + return regione_is_refresh_step(step); +} + +bool RegionECache::regione_should_compute_velocity(int64_t step) const { + if (!regione_enabled_) return true; + if (regione_should_run_full_step(step)) return true; + const auto interval = + std::max(1, config_.regione.skip_interval_steps); + return ((step - config_.regione.warmup_steps) % interval) == 0; +} + +bool RegionECache::regione_should_direct_unedited(int64_t step) const { + if (!regione_enabled_ || !regione_has_regions()) return false; + if (regione_is_tail_step(step)) return false; + return step == config_.regione.warmup_steps - 1 || + regione_is_refresh_step(step); +} + +int64_t RegionECache::regione_next_direct_step(int64_t step) const { + int64_t tail_start = regione_infer_steps_; + if (config_.regione.tail_steps > 0 && regione_infer_steps_ > 0) { + tail_start = + std::max(0, regione_infer_steps_ - config_.regione.tail_steps); + } + int64_t next_step = tail_start; + for (const auto refresh_step : config_.regione.refresh_steps) { + const auto refresh_index = + refresh_step > 0 ? refresh_step - 1 : refresh_step; + if (refresh_index > step && refresh_index < next_step) + next_step = refresh_index; + } + if (next_step <= step) next_step = step + 1; + if (regione_infer_steps_ > 0) + next_step = std::min(next_step, regione_infer_steps_); + return next_step; +} + +void RegionECache::regione_prepare_inference( + const torch::Tensor& latents, + const torch::Tensor& condition_latents, + int64_t grid_h, + int64_t grid_w, + int64_t sp_rank, + int64_t sp_size) { + if (!regione_enabled_) return; + regione_target_seq_len_ = + latents.defined() && latents.dim() > 1 ? latents.size(1) : 0; + const auto condition_seq_len = + condition_latents.defined() && condition_latents.dim() > 1 + ? condition_latents.size(1) + : 0; + regione_image_seq_len_ = regione_target_seq_len_ + condition_seq_len; + regione_grid_h_ = grid_h; + regione_grid_w_ = grid_w; + regione_sp_rank_ = sp_rank; + regione_sp_size_ = std::max(1, sp_size); + const auto shard = regione_sp_size_ > 0 + ? regione_image_seq_len_ / regione_sp_size_ + : regione_image_seq_len_; + regione_local_start_ = regione_sp_rank_ * shard; + regione_local_end_ = regione_sp_rank_ == regione_sp_size_ - 1 + ? regione_image_seq_len_ + : regione_local_start_ + shard; + if (regione_image_seq_len_ > 0 && regione_local_end_ > regione_local_start_) { + regione_local_image_global_ids_ = + torch::arange(regione_local_start_, + regione_local_end_, + latents.options().dtype(torch::kLong)); + } else { + regione_local_image_global_ids_ = torch::Tensor(); + } + regione_condition_latents_ = condition_latents; + regione_edited_ids_ = torch::Tensor(); + regione_unedited_ids_ = torch::Tensor(); + regione_velocity_cache_ = torch::Tensor(); + regione_partial_mode_ = false; + regione_local_edited_global_ids_ = torch::Tensor(); + regione_local_edited_cache_ids_ = torch::Tensor(); + for (auto& cache : regione_k_cache_cpu_) cache = torch::Tensor(); + for (auto& cache : regione_v_cache_cpu_) cache = torch::Tensor(); + for (auto& cache : regione_cond_k_cache_cpu_) cache = torch::Tensor(); + for (auto& cache : regione_cond_v_cache_cpu_) cache = torch::Tensor(); + regione_clear_all_prefetch_slots(); +} + +torch::Tensor RegionECache::regione_normalize_ids( + const torch::Tensor& ids, + const torch::Device& device) const { + if (!ids.defined()) return torch::Tensor(); + auto out = ids; + if (out.dim() > 1) out = out.reshape({-1}); + return out.to(device, torch::kLong, /*non_blocking=*/false, /*copy=*/false); +} + +torch::Tensor RegionECache::regione_gather_ids(const torch::Tensor& tensor, + const torch::Tensor& ids, + int64_t dim) const { + if (!tensor.defined() || !ids.defined()) return tensor; + return tensor.index_select(dim, regione_normalize_ids(ids, tensor.device())); +} + +torch::Tensor RegionECache::regione_scatter_ids(const torch::Tensor& values, + const torch::Tensor& ids, + const torch::Tensor& base, + int64_t dim) const { + if (!values.defined() || !ids.defined() || !base.defined()) return base; + auto out = base.clone(); + out.index_copy_(dim, regione_normalize_ids(ids, base.device()), values); + return out; +} + +torch::Tensor RegionECache::regione_active_edited_ids() const { + if (regione_is_partial_sp_mode() && + regione_local_edited_global_ids_.defined()) { + return regione_local_edited_global_ids_; + } + return regione_edited_ids_; +} + +torch::Tensor RegionECache::regione_kv_update_ids() const { + if (regione_is_partial_sp_mode() && + regione_local_edited_cache_ids_.defined()) { + return regione_local_edited_cache_ids_; + } + return regione_edited_ids_; +} + +torch::Tensor RegionECache::regione_gather_edited( + const torch::Tensor& tensor) const { + return regione_gather_ids(tensor, regione_active_edited_ids(), 1); +} + +torch::Tensor RegionECache::regione_gather_unedited( + const torch::Tensor& tensor) const { + return regione_gather_ids(tensor, regione_unedited_ids_, 1); +} + +torch::Tensor RegionECache::regione_scatter_edited( + const torch::Tensor& edited, + const torch::Tensor& base) const { + return regione_scatter_ids(edited, regione_active_edited_ids(), base, 1); +} + +torch::Tensor RegionECache::regione_scatter_unedited( + const torch::Tensor& unedited, + const torch::Tensor& base) const { + return regione_scatter_ids(unedited, regione_unedited_ids_, base, 1); +} + +torch::Tensor RegionECache::regione_gather_query_rope( + const torch::Tensor& image_rope) const { + auto ids = regione_active_edited_ids(); + if (!image_rope.defined() || !ids.defined()) return image_rope; + if (image_rope.size(0) == ids.numel()) return image_rope; + return regione_gather_ids(image_rope, ids, 0); +} + +torch::Tensor RegionECache::regione_gather_key_rope( + const torch::Tensor& image_rope, + int64_t key_len) const { + if (!image_rope.defined() || !regione_is_partial_sp_mode()) return image_rope; + if (image_rope.size(0) == key_len) return image_rope; + if (!regione_local_image_global_ids_.defined() || + regione_local_image_global_ids_.numel() != key_len) { + return image_rope; + } + return regione_gather_ids(image_rope, regione_local_image_global_ids_, 0); +} + +torch::Tensor RegionECache::regione_local_update_mask( + const torch::Tensor& base) const { + if (!base.defined()) return torch::Tensor(); + auto mask = torch::zeros({base.size(0), base.size(1), 1}, base.options()); + auto ids = regione_active_edited_ids(); + if (ids.defined() && ids.numel() > 0) { + auto ones = torch::ones({base.size(0), ids.numel(), 1}, base.options()); + mask.index_copy_(1, regione_normalize_ids(ids, base.device()), ones); + } + return mask; +} + +void RegionECache::regione_select_regions(const torch::Tensor& sample, + const torch::Tensor& model_output, + const torch::Tensor& sigmas, + int64_t step) { + if (!regione_enabled_ || regione_has_regions()) return; + if (!sample.defined() || !model_output.defined() || + !regione_condition_latents_.defined()) + return; + if (sample.dim() != 3 || sample.size(0) != 1) { + regione_edited_ids_ = + torch::arange(sample.size(1), sample.options().dtype(torch::kLong)); + regione_unedited_ids_ = + torch::empty({0}, sample.options().dtype(torch::kLong)); + regione_update_local_ids(); + return; + } + auto condition = regione_condition_latents_; + if (condition.dim() != 3 || condition.size(1) < sample.size(1)) { + regione_edited_ids_ = + torch::arange(sample.size(1), sample.options().dtype(torch::kLong)); + regione_unedited_ids_ = + torch::empty({0}, sample.options().dtype(torch::kLong)); + regione_update_local_ids(); + return; + } + condition = condition.slice(1, 0, sample.size(1)).to(sample.dtype()); + auto sigma = sigmas.index({step}).to(sample.device()).to(sample.dtype()); + auto sigma_final = sigmas.index({-1}).to(sample.device()).to(sample.dtype()); + auto estimate = sample + (sigma_final - sigma) * model_output; + auto estimate_norm = + estimate / + torch::sqrt(torch::sum(estimate * estimate, -1, true)).clamp_min(1e-6); + auto condition_norm = + condition / + torch::sqrt(torch::sum(condition * condition, -1, true)).clamp_min(1e-6); + auto similarity = torch::sum(estimate_norm * condition_norm, -1); + auto selected_mask = similarity <= config_.regione.region_threshold; + if (regione_grid_h_ > 0 && regione_grid_w_ > 0 && + regione_grid_h_ * regione_grid_w_ == sample.size(1)) { + auto mask2d = selected_mask.to(torch::kFloat) + .view({1, 1, regione_grid_h_, regione_grid_w_}); + auto vertical_pool_opts = + torch::nn::functional::MaxPool2dFuncOptions({3, 1}).stride(1).padding( + {1, 0}); + auto horizontal_pool_opts = + torch::nn::functional::MaxPool2dFuncOptions({1, 3}).stride(1).padding( + {0, 1}); + auto eroded_vertical = + -torch::nn::functional::max_pool2d(-mask2d, vertical_pool_opts); + auto eroded_horizontal = + -torch::nn::functional::max_pool2d(-mask2d, horizontal_pool_opts); + auto eroded = eroded_vertical * eroded_horizontal; + auto dilation_pool_opts = + torch::nn::functional::MaxPool2dFuncOptions({5, 5}).stride(1).padding( + 2); + auto dilated = + torch::nn::functional::max_pool2d(eroded, dilation_pool_opts); + selected_mask = dilated.view({1, -1}) > 0.5; + } + auto edited = torch::nonzero(selected_mask[0]) + .reshape({-1}) + .to(sample.device(), torch::kLong); + if (edited.numel() == 0) { + edited = std::get<1>(similarity[0].min(0, false)) + .reshape({1}) + .to(sample.device(), torch::kLong); + } + auto unedited_mask = + torch::ones({sample.size(1)}, sample.options().dtype(torch::kBool)); + unedited_mask.index_fill_(0, edited, false); + auto unedited = torch::nonzero(unedited_mask) + .reshape({-1}) + .to(sample.device(), torch::kLong); + regione_edited_ids_ = edited; + regione_unedited_ids_ = unedited; + regione_update_local_ids(); +} + +void RegionECache::regione_update_local_ids() { + regione_local_edited_global_ids_ = torch::Tensor(); + regione_local_edited_cache_ids_ = torch::Tensor(); + if (!regione_edited_ids_.defined() || regione_sp_size_ <= 1) return; + auto ids = + regione_normalize_ids(regione_edited_ids_, regione_edited_ids_.device()); + auto mask = (ids >= regione_local_start_) & (ids < regione_local_end_); + auto local_global = ids.index({mask}); + regione_local_edited_global_ids_ = local_global; + regione_local_edited_cache_ids_ = local_global - regione_local_start_; +} + +void RegionECache::regione_update_velocity_cache(const torch::Tensor& value) { + if (regione_enabled_ && value.defined()) regione_velocity_cache_ = value; +} + +torch::Tensor RegionECache::regione_velocity_cache() const { + return regione_velocity_cache_; +} + +RegionEStepPlan RegionECache::begin_step(int64_t step) { + RegionEStepPlan plan; + plan.enabled = regione_enabled_; + if (!regione_enabled_) return plan; + + regione_set_current_step(step); + plan.full_step = regione_should_run_full_step(step); + plan.partial_step = regione_has_regions() && !plan.full_step; + regione_set_partial_mode(plan.partial_step); + plan.use_velocity_cache = !regione_should_compute_velocity(step) && + regione_velocity_cache_.defined(); + plan.run_partition = !regione_has_regions() && + step == regione_warmup_steps() - 1 && + !plan.use_velocity_cache; + plan.direct_unedited = regione_should_direct_unedited(step); + return plan; +} + +RegionEStepInput RegionECache::prepare_step_input( + const torch::Tensor& latents, + const torch::Tensor& condition_latents, + const std::vector>& main_shape, + const RegionEStepPlan& plan) const { + RegionEStepInput input; + input.step_latents = + plan.partial_step ? regione_gather_edited(latents) : latents; + input.latent_model_input = input.step_latents; + if (!plan.partial_step && condition_latents.defined()) { + input.latent_model_input = torch::cat({latents, condition_latents}, 1); + } + input.main_shape = main_shape; + if (plan.partial_step) { + input.main_shape = {{1, input.step_latents.size(1), 1}}; + } + input.use_cached_velocity = plan.use_velocity_cache; + if (plan.use_velocity_cache) { + input.cached_velocity = regione_velocity_cache_; + if (plan.partial_step && input.cached_velocity.defined() && + input.cached_velocity.size(1) != input.step_latents.size(1)) { + input.cached_velocity = regione_gather_edited(input.cached_velocity); + } + } + return input; +} + +void RegionECache::observe_velocity(const torch::Tensor& latents, + const torch::Tensor& noise_pred, + const torch::Tensor& sigmas, + int64_t step, + const RegionEStepPlan& plan) { + if (!regione_enabled_ || plan.use_velocity_cache) return; + if (plan.run_partition) { + regione_select_regions(latents, noise_pred, sigmas, step); + } + regione_update_velocity_cache(noise_pred); +} + +torch::Tensor RegionECache::apply_direct_unedited( + const torch::Tensor& prev_latents, + const torch::Tensor& latents, + const torch::Tensor& noise_pred, + const torch::Tensor& sigmas, + int64_t step) const { + auto sigma = sigmas.index({step}).to(latents.device()).to(latents.dtype()); + auto next_direct_step = regione_next_direct_step(step); + auto sigma_direct = + sigmas.index({next_direct_step}).to(latents.device()).to(latents.dtype()); + auto unedited_direct = + regione_gather_unedited(latents) + + (sigma_direct - sigma) * regione_gather_unedited(noise_pred); + return regione_scatter_unedited(unedited_direct, prev_latents); +} + +void RegionECache::regione_prefetch_img_kv(int64_t block_id, + bool use_cfg, + const torch::Tensor& reference) { + if (!reference.defined()) return; + regione_prefetch_img_kv( + block_id, use_cfg, reference.device(), reference.scalar_type()); +} + +void RegionECache::regione_set_current_block(int64_t block_id, + bool use_cfg, + const torch::Tensor& reference) { + regione_current_block_ = block_id; + regione_current_use_cfg_ = use_cfg; + if (reference.defined()) { + regione_prefetch_img_kv( + block_id + 1, use_cfg, reference.device(), reference.scalar_type()); + } +} + +void RegionECache::regione_set_current_step(int64_t step) { + regione_current_step_ = step; +} + +void RegionECache::regione_set_partial_mode(bool partial_mode) { + regione_partial_mode_ = regione_enabled_ && partial_mode; +} + +bool RegionECache::regione_is_partial_sp_mode() const { + return regione_enabled_ && regione_partial_mode_ && regione_sp_size_ > 1; +} + +bool RegionECache::regione_should_store_kv() const { + return regione_enabled_ && !regione_partial_mode_; +} + +bool RegionECache::regione_should_patch_kv() const { + return regione_enabled_ && regione_partial_mode_; +} + +void RegionECache::ensure_regione_kv_size(int64_t num_blocks) { + if (num_blocks <= 0) return; + regione_k_cache_cpu_.resize(num_blocks); + regione_v_cache_cpu_.resize(num_blocks); + regione_cond_k_cache_cpu_.resize(num_blocks); + regione_cond_v_cache_cpu_.resize(num_blocks); +} + +std::vector& +RegionECache::regione_prefetch_slots(bool use_cfg) { + auto& slots = + use_cfg ? regione_cond_prefetch_slots_ : regione_prefetch_slots_; + if (slots.empty()) slots.resize(2); + return slots; +} + +void RegionECache::regione_clear_prefetch_slot(RegionEPrefetchedKV& slot) { + slot.block_id = -1; + slot.key = torch::Tensor(); + slot.value = torch::Tensor(); +#if defined(USE_NPU) + slot.ready_event.reset(); +#endif +} + +void RegionECache::regione_clear_prefetch_block(bool use_cfg, + int64_t block_id) { + for (auto& slot : regione_prefetch_slots(use_cfg)) { + if (slot.block_id == block_id) regione_clear_prefetch_slot(slot); + } +} + +void RegionECache::regione_clear_all_prefetch_slots() { + for (auto& slot : regione_prefetch_slots_) { + regione_clear_prefetch_slot(slot); + } + for (auto& slot : regione_cond_prefetch_slots_) { + regione_clear_prefetch_slot(slot); + } +} + +void RegionECache::regione_prefetch_img_kv(int64_t block_id, + bool use_cfg, + const torch::Device& device, + c10::ScalarType dtype) { + if (!regione_enabled_ || !regione_partial_mode_ || block_id < 0 || + block_id >= regione_num_blocks_) { + return; + } + auto& k_cache = use_cfg ? regione_cond_k_cache_cpu_ : regione_k_cache_cpu_; + auto& v_cache = use_cfg ? regione_cond_v_cache_cpu_ : regione_v_cache_cpu_; + if (!tensor_cache_ready(k_cache, block_id) || + !tensor_cache_ready(v_cache, block_id)) { + return; + } + + auto& slots = regione_prefetch_slots(use_cfg); + for (const auto& slot : slots) { + if (slot.block_id == block_id && slot.key.defined() && + slot.value.defined() && slot.key.device() == device && + slot.key.scalar_type() == dtype) { + return; + } + } + + RegionEPrefetchedKV* target = nullptr; + for (auto& slot : slots) { + if (slot.block_id < 0 || !slot.key.defined() || !slot.value.defined()) { + target = &slot; + break; + } + } + if (target == nullptr) { + target = &slots[static_cast(block_id) % slots.size()]; + } + regione_clear_prefetch_slot(*target); + +#if defined(USE_NPU) + if (device.is_privateuseone()) { + auto stream = c10_npu::getStreamFromPool(false, device.index()); + { + c10_npu::NPUStreamGuard stream_guard(stream); + target->key = k_cache[block_id] + .to(device, dtype, /*non_blocking=*/true, /*copy=*/true) + .contiguous(); + target->value = + v_cache[block_id] + .to(device, dtype, /*non_blocking=*/true, /*copy=*/true) + .contiguous(); + target->ready_event = + std::make_shared(ACL_EVENT_EXTERNAL); + target->ready_event->record(stream); + } + target->block_id = block_id; + return; + } +#endif + + target->key = k_cache[block_id] + .to(device, dtype, /*non_blocking=*/false, /*copy=*/true) + .contiguous(); + target->value = v_cache[block_id] + .to(device, dtype, /*non_blocking=*/false, /*copy=*/true) + .contiguous(); + target->block_id = block_id; +} + +bool RegionECache::regione_take_prefetched_img_kv(int64_t block_id, + bool use_cfg, + const torch::Device& device, + c10::ScalarType dtype, + torch::Tensor* key, + torch::Tensor* value) { + for (auto& slot : regione_prefetch_slots(use_cfg)) { + if (slot.block_id != block_id || !slot.key.defined() || + !slot.value.defined() || slot.key.device() != device || + slot.key.scalar_type() != dtype) { + continue; + } +#if defined(USE_NPU) + if (slot.ready_event != nullptr && device.is_privateuseone()) { + auto current_stream = c10_npu::getCurrentNPUStream(device.index()); + slot.ready_event->block(current_stream); + } +#endif + *key = slot.key; + *value = slot.value; + regione_clear_prefetch_slot(slot); + return true; + } + return false; +} + +void RegionECache::regione_store_img_kv(int64_t block_id, + bool use_cfg, + const torch::Tensor& key, + const torch::Tensor& value) { + if (!regione_enabled_ || block_id < 0) return; + ensure_regione_kv_size(std::max(regione_num_blocks_, block_id + 1)); + auto& k_cache = use_cfg ? regione_cond_k_cache_cpu_ : regione_k_cache_cpu_; + auto& v_cache = use_cfg ? regione_cond_v_cache_cpu_ : regione_v_cache_cpu_; + k_cache[block_id] = regione_to_cpu_cache(key); + v_cache[block_id] = regione_to_cpu_cache(value); + regione_clear_prefetch_block(use_cfg, block_id); +} + +std::pair RegionECache::process_image_kv( + const torch::Tensor& key, + const torch::Tensor& value) { + if (!regione_enabled_) return {key, value}; + if (regione_should_store_kv()) { + regione_store_img_kv( + regione_current_block_, regione_current_use_cfg_, key, value); + return {key, value}; + } + if (regione_should_patch_kv()) { + return regione_patch_img_kv( + regione_current_block_, regione_current_use_cfg_, key, value); + } + return {key, value}; +} + +std::pair RegionECache::adjust_image_rope( + const torch::Tensor& image_rope, + int64_t key_len) const { + auto query_rope = regione_partial_mode_ + ? regione_gather_query_rope(image_rope) + : image_rope; + auto key_rope = regione_gather_key_rope(image_rope, key_len); + return {query_rope, key_rope}; +} + +std::pair RegionECache::regione_patch_img_kv( + int64_t block_id, + bool use_cfg, + const torch::Tensor& key, + const torch::Tensor& value) { + if (!regione_enabled_ || block_id < 0) return {key, value}; + auto& k_cache = use_cfg ? regione_cond_k_cache_cpu_ : regione_k_cache_cpu_; + auto& v_cache = use_cfg ? regione_cond_v_cache_cpu_ : regione_v_cache_cpu_; + if (!tensor_cache_ready(k_cache, block_id) || + !tensor_cache_ready(v_cache, block_id)) { + regione_store_img_kv(block_id, use_cfg, key, value); + return {key, value}; + } + torch::Tensor full_key; + torch::Tensor full_value; + const auto took_prefetched = regione_take_prefetched_img_kv(block_id, + use_cfg, + key.device(), + key.scalar_type(), + &full_key, + &full_value); + if (!took_prefetched) { + full_key = k_cache[block_id] + .to(key.device(), + key.scalar_type(), + /*non_blocking=*/true, + /*copy=*/true) + .contiguous(); + full_value = v_cache[block_id] + .to(value.device(), + value.scalar_type(), + /*non_blocking=*/true, + /*copy=*/true) + .contiguous(); + } + if (full_key.sizes() == key.sizes()) { + regione_store_img_kv(block_id, use_cfg, key, value); + return {key, value}; + } + auto update_ids = regione_kv_update_ids(); + if (update_ids.defined() && key.dim() >= 2 && + full_key.size(0) == key.size(0) && key.size(1) == update_ids.numel()) { + full_key = regione_scatter_ids(key, update_ids, full_key, 1); + full_value = regione_scatter_ids(value, update_ids, full_value, 1); + } + return {full_key, full_value}; +} + +} // namespace xllm diff --git a/xllm/core/framework/dit_cache/regione.h b/xllm/core/framework/dit_cache/regione.h new file mode 100644 index 0000000000..4de3840674 --- /dev/null +++ b/xllm/core/framework/dit_cache/regione.h @@ -0,0 +1,220 @@ +/* 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 "dit_cache_config.h" + +#if defined(USE_NPU) +#include +#include +#endif + +namespace xllm { + +struct RegionEStepPlan { + bool enabled = false; + bool full_step = true; + bool partial_step = false; + bool use_velocity_cache = false; + bool run_partition = false; + bool direct_unedited = false; +}; + +struct RegionEStepInput { + torch::Tensor step_latents; + torch::Tensor latent_model_input; + std::vector> main_shape; + torch::Tensor cached_velocity; + bool use_cached_velocity = false; +}; + +class RegionECache { + private: + struct RegionEPrefetchedKV { + int64_t block_id = -1; + torch::Tensor key; + torch::Tensor value; +#if defined(USE_NPU) + std::shared_ptr ready_event; +#endif + }; + + public: + RegionECache() = default; + ~RegionECache() = default; + + RegionECache(const RegionECache&) = delete; + RegionECache& operator=(const RegionECache&) = delete; + RegionECache(RegionECache&&) = default; + RegionECache& operator=(RegionECache&&) = default; + + void init(const DiTCacheConfig& cfg); + void set_infer_steps(int64_t infer_steps) { + regione_infer_steps_ = infer_steps; + } + void set_num_blocks(int64_t num_blocks); + + bool is_enabled() const { return regione_enabled_; } + int64_t regione_warmup_steps() const { return config_.regione.warmup_steps; } + bool regione_has_regions() const { return regione_edited_ids_.defined(); } + bool regione_is_partial_mode() const { return regione_partial_mode_; } + bool regione_is_partial_sp_mode() const; + bool regione_should_compute_velocity(int64_t step) const; + bool regione_should_run_full_step(int64_t step) const; + bool regione_should_direct_unedited(int64_t step) const; + int64_t regione_next_direct_step(int64_t step) const; + void regione_prepare_inference(const torch::Tensor& latents, + const torch::Tensor& condition_latents, + int64_t grid_h, + int64_t grid_w, + int64_t sp_rank = 0, + int64_t sp_size = 1); + + RegionEStepPlan begin_step(int64_t step); + RegionEStepInput prepare_step_input( + const torch::Tensor& latents, + const torch::Tensor& condition_latents, + const std::vector>& main_shape, + const RegionEStepPlan& plan) const; + void observe_velocity(const torch::Tensor& latents, + const torch::Tensor& noise_pred, + const torch::Tensor& sigmas, + int64_t step, + const RegionEStepPlan& plan); + torch::Tensor apply_direct_unedited(const torch::Tensor& prev_latents, + const torch::Tensor& latents, + const torch::Tensor& noise_pred, + const torch::Tensor& sigmas, + int64_t step) const; + + torch::Tensor regione_gather_edited(const torch::Tensor& tensor) const; + torch::Tensor regione_gather_unedited(const torch::Tensor& tensor) const; + torch::Tensor regione_scatter_edited(const torch::Tensor& edited, + const torch::Tensor& base) const; + torch::Tensor regione_scatter_unedited(const torch::Tensor& unedited, + const torch::Tensor& base) const; + torch::Tensor regione_local_update_mask(const torch::Tensor& base) const; + + void regione_prefetch_img_kv(int64_t block_id, + bool use_cfg, + const torch::Tensor& reference); + void regione_set_current_block( + int64_t block_id, + bool use_cfg, + const torch::Tensor& reference = torch::Tensor()); + std::pair process_image_kv( + const torch::Tensor& key, + const torch::Tensor& value); + std::pair adjust_image_rope( + const torch::Tensor& image_rope, + int64_t key_len) const; + + private: + void regione_select_regions(const torch::Tensor& sample, + const torch::Tensor& model_output, + const torch::Tensor& sigmas, + int64_t step); + torch::Tensor regione_gather_query_rope( + const torch::Tensor& image_rope) const; + torch::Tensor regione_gather_key_rope(const torch::Tensor& image_rope, + int64_t key_len) const; + void regione_update_velocity_cache(const torch::Tensor& value); + torch::Tensor regione_velocity_cache() const; + void regione_set_current_step(int64_t step); + void regione_set_partial_mode(bool partial_mode); + int64_t regione_current_block() const { return regione_current_block_; } + bool regione_current_use_cfg() const { return regione_current_use_cfg_; } + bool regione_should_store_kv() const; + bool regione_should_patch_kv() const; + void regione_store_img_kv(int64_t block_id, + bool use_cfg, + const torch::Tensor& key, + const torch::Tensor& value); + std::pair regione_patch_img_kv( + int64_t block_id, + bool use_cfg, + const torch::Tensor& key, + const torch::Tensor& value); + + private: + bool regione_is_refresh_step(int64_t step) const; + bool regione_is_tail_step(int64_t step) const; + torch::Tensor regione_normalize_ids(const torch::Tensor& ids, + const torch::Device& device) const; + torch::Tensor regione_active_edited_ids() const; + torch::Tensor regione_kv_update_ids() const; + void regione_update_local_ids(); + torch::Tensor regione_gather_ids(const torch::Tensor& tensor, + const torch::Tensor& ids, + int64_t dim) const; + torch::Tensor regione_scatter_ids(const torch::Tensor& values, + const torch::Tensor& ids, + const torch::Tensor& base, + int64_t dim) const; + void ensure_regione_kv_size(int64_t num_blocks); + std::vector& regione_prefetch_slots(bool use_cfg); + void regione_clear_prefetch_slot(RegionEPrefetchedKV& slot); + void regione_clear_prefetch_block(bool use_cfg, int64_t block_id); + void regione_clear_all_prefetch_slots(); + void regione_prefetch_img_kv(int64_t block_id, + bool use_cfg, + const torch::Device& device, + c10::ScalarType dtype); + bool regione_take_prefetched_img_kv(int64_t block_id, + bool use_cfg, + const torch::Device& device, + c10::ScalarType dtype, + torch::Tensor* key, + torch::Tensor* value); + + DiTCacheConfig config_; + bool regione_enabled_ = false; + int64_t regione_infer_steps_ = 0; + int64_t regione_num_blocks_ = 0; + int64_t regione_current_step_ = 0; + int64_t regione_current_block_ = -1; + bool regione_current_use_cfg_ = false; + bool regione_partial_mode_ = false; + int64_t regione_target_seq_len_ = 0; + int64_t regione_grid_h_ = 0; + int64_t regione_grid_w_ = 0; + int64_t regione_image_seq_len_ = 0; + int64_t regione_sp_rank_ = 0; + int64_t regione_sp_size_ = 1; + int64_t regione_local_start_ = 0; + int64_t regione_local_end_ = 0; + torch::Tensor regione_condition_latents_; + torch::Tensor regione_edited_ids_; + torch::Tensor regione_unedited_ids_; + torch::Tensor regione_local_edited_global_ids_; + torch::Tensor regione_local_edited_cache_ids_; + torch::Tensor regione_local_image_global_ids_; + torch::Tensor regione_velocity_cache_; + std::vector regione_k_cache_cpu_; + std::vector regione_v_cache_cpu_; + std::vector regione_cond_k_cache_cpu_; + std::vector regione_cond_v_cache_cpu_; + std::vector regione_prefetch_slots_; + std::vector regione_cond_prefetch_slots_; +}; + +} // namespace xllm diff --git a/xllm/core/runtime/dit_worker_impl.cpp b/xllm/core/runtime/dit_worker_impl.cpp index 08c69ce9a6..13b04954ac 100644 --- a/xllm/core/runtime/dit_worker_impl.cpp +++ b/xllm/core/runtime/dit_worker_impl.cpp @@ -23,6 +23,7 @@ limitations under the License. #include #include +#include #include #include "common/device_monitor.h" @@ -42,6 +43,16 @@ limitations under the License. namespace xllm { namespace { +std::vector parse_regione_refresh_steps(const std::string& text) { + std::vector steps; + std::stringstream ss(text); + std::string item; + while (std::getline(ss, item, ',')) { + if (!item.empty()) steps.push_back(std::stoll(item)); + } + return steps; +} + DiTCacheConfig parse_dit_cache_from_flags() { DiTCacheConfig cache_config; if (::xllm::DiTConfig::get_instance().dit_cache_policy() == "FBCache") { @@ -81,6 +92,19 @@ DiTCacheConfig parse_dit_cache_from_flags() { ::xllm::DiTConfig::get_instance().dit_cache_end_blocks(); cache_config.residual_cache.skip_interval_steps = ::xllm::DiTConfig::get_instance().dit_cache_skip_interval_steps(); + } else if (::xllm::DiTConfig::get_instance().dit_cache_policy() == + "RegionE") { + cache_config.selected_policy = PolicyType::RegionE; + cache_config.regione.warmup_steps = + ::xllm::DiTConfig::get_instance().dit_cache_warmup_steps(); + cache_config.regione.skip_interval_steps = + ::xllm::DiTConfig::get_instance().dit_cache_skip_interval_steps(); + cache_config.regione.tail_steps = + ::xllm::DiTConfig::get_instance().dit_cache_end_steps(); + cache_config.regione.refresh_steps = parse_regione_refresh_steps( + ::xllm::DiTConfig::get_instance().dit_regione_refresh_steps()); + cache_config.regione.region_threshold = + ::xllm::DiTConfig::get_instance().dit_regione_region_threshold(); } else if (::xllm::DiTConfig::get_instance().dit_cache_policy() == "None") { cache_config.selected_policy = PolicyType::None; } diff --git a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h index 3c75bfda13..0662fbcd74 100644 --- a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h +++ b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -1055,35 +1056,64 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { do_true_cfg ? get_image_rotary_emb(negative_prompt_embeds.size(1)) : image_rotary_emb_pos; + auto* regione = DiTCache::get_instance().regione(); + if (regione) { + const auto sp_group = parallel_args_.dit_sp_group_; + const int64_t sp_rank = sp_group ? sp_group->rank() : 0; + const int64_t sp_size = sp_group ? sp_group->world_size() : 1; + regione->regione_prepare_inference(final_latents, + image_latents, + main_shape[0][1], + main_shape[0][2], + sp_rank, + sp_size); + } + for (int64_t i = 0; i < timesteps.size(0); ++i) { auto t = timesteps[i]; current_timestep_ = t; - auto latent_model_input = final_latents; - if (image_latents.defined()) { + RegionEStepPlan regione_plan; + RegionEStepInput regione_input; + if (regione) { + regione_plan = regione->begin_step(i); + regione_input = regione->prepare_step_input( + final_latents, image_latents, main_shape, regione_plan); + } + + auto step_latents = + regione_plan.enabled ? regione_input.step_latents : final_latents; + auto latent_model_input = step_latents; + if (regione_plan.enabled) { + latent_model_input = regione_input.latent_model_input; + } else if (image_latents.defined()) { latent_model_input = torch::cat({final_latents, image_latents}, 1); } + auto step_main_shape = + regione_plan.enabled ? regione_input.main_shape : main_shape; auto timestep_expanded = - t.expand({final_latents.size(0)}).to(final_latents.dtype()); + t.expand({step_latents.size(0)}).to(step_latents.dtype()); torch::Tensor noise_pred; torch::Tensor neg_noise_pred; torch::Tensor pos_neg_noise_preds; - if (::xllm::ParallelConfig::get_instance().cfg_size() == 2 && - do_true_cfg) { + if (regione_plan.use_velocity_cache) { + noise_pred = regione_input.cached_velocity; + } else if (::xllm::ParallelConfig::get_instance().cfg_size() == 2 && + do_true_cfg) { auto rank = parallel_args_.dit_cfg_group_->rank(); if (rank == 0) { noise_pred = transformer_->forward(latent_model_input, prompt_embeds, prompt_embeds_mask, timestep_expanded / 1000.0, - main_shape, + step_main_shape, txt_seq_lens, image_rotary_emb_pos, /*use_cfg=*/false, /*step_index=*/i + 1); - noise_pred = noise_pred.slice(1, 0, final_latents.size(1)); + noise_pred = noise_pred.slice(1, 0, step_latents.size(1)); pos_neg_noise_preds = xllm::parallel_state::gather(noise_pred, parallel_args_.dit_cfg_group_, @@ -1093,13 +1123,13 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { negative_prompt_embeds, negative_prompt_embeds_mask, timestep_expanded / 1000.0, - main_shape, + step_main_shape, negative_txt_seq_lens, image_rotary_emb_neg, /*use_cfg=*/true, /*step_index=*/i + 1); - neg_noise_pred = neg_noise_pred.slice(1, 0, final_latents.size(1)); + neg_noise_pred = neg_noise_pred.slice(1, 0, step_latents.size(1)); pos_neg_noise_preds = xllm::parallel_state::gather(neg_noise_pred, parallel_args_.dit_cfg_group_, @@ -1117,24 +1147,24 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { prompt_embeds, prompt_embeds_mask, timestep_expanded / 1000.0, - main_shape, + step_main_shape, txt_seq_lens, image_rotary_emb_pos, /*use_cfg=*/false, /*step_index=*/i + 1); - noise_pred = noise_pred.slice(1, 0, final_latents.size(1)); + noise_pred = noise_pred.slice(1, 0, step_latents.size(1)); if (do_true_cfg) { neg_noise_pred = transformer_->forward(latent_model_input, negative_prompt_embeds, negative_prompt_embeds_mask, timestep_expanded / 1000.0, - main_shape, + step_main_shape, negative_txt_seq_lens, image_rotary_emb_neg, /*use_cfg=*/true, /*step_index=*/i + 1); - neg_noise_pred = neg_noise_pred.slice(1, 0, final_latents.size(1)); + neg_noise_pred = neg_noise_pred.slice(1, 0, step_latents.size(1)); auto comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred); @@ -1144,8 +1174,49 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { } } + if (regione_plan.enabled) { + regione->observe_velocity( + final_latents, noise_pred, scheduler_->sigmas(), i, regione_plan); + if (regione_plan.run_partition) { + regione_plan.direct_unedited = + regione->regione_should_direct_unedited(i); + } + } + auto latents_dtype = final_latents.dtype(); - final_latents = scheduler_->step(noise_pred, t, final_latents); + if (regione_plan.enabled) { + if (regione_plan.partial_step) { + auto edited_prev = scheduler_->step(noise_pred, t, step_latents); + auto partial_latents = + regione->regione_scatter_edited(edited_prev, final_latents); + if (::xllm::ParallelConfig::get_instance().sp_size() > 1) { + auto update_mask = + regione->regione_local_update_mask(final_latents); + auto update_values = + (partial_latents - final_latents) * update_mask; + auto reduced_values = xllm::parallel_state::reduce( + update_values, parallel_args_.dit_sp_group_); + auto reduced_mask = xllm::parallel_state::reduce( + update_mask, parallel_args_.dit_sp_group_) + .clamp(0, 1); + final_latents = final_latents + reduced_values * reduced_mask; + } else { + final_latents = partial_latents; + } + } else { + auto prev_latents = scheduler_->step(noise_pred, t, final_latents); + if (regione_plan.direct_unedited) { + prev_latents = regione->apply_direct_unedited(prev_latents, + final_latents, + noise_pred, + scheduler_->sigmas(), + i); + } + final_latents = prev_latents; + } + } else { + final_latents = scheduler_->step(noise_pred, t, final_latents); + } if (final_latents.dtype() != latents_dtype) { final_latents = final_latents.to(latents_dtype); } diff --git a/xllm/models/dit/transformers/transformer_qwen_image.h b/xllm/models/dit/transformers/transformer_qwen_image.h index 75bb57dc9c..f0e24411fb 100644 --- a/xllm/models/dit/transformers/transformer_qwen_image.h +++ b/xllm/models/dit/transformers/transformer_qwen_image.h @@ -64,6 +64,7 @@ namespace xllm { inline bool use_dit_sp_communication_overlap() { return DiTConfig::get_instance().dit_sp_communication_overlap() && + DiTConfig::get_instance().dit_cache_policy() != "RegionE" && ParallelConfig::get_instance().sp_size() > 1; } @@ -1583,6 +1584,11 @@ class QwenDoubleStreamAttnProcessor2_0Impl : public torch::nn::Module { auto img_query = attn_->to_q_->forward(hidden_states); auto img_key = attn_->to_k_->forward(hidden_states); auto img_value = attn_->to_v_->forward(hidden_states); + auto* regione = DiTCache::get_instance().regione(); + if (regione) { + std::tie(img_key, img_value) = + regione->process_image_kv(img_key, img_value); + } // Compute QKV for text stream (context projections) auto txt_query = attn_->add_q_proj_->forward(encoder_hidden_states); @@ -1632,8 +1638,14 @@ class QwenDoubleStreamAttnProcessor2_0Impl : public torch::nn::Module { xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( img_value, /*tensor_name=*/"hidden_states", /*dim=*/1); - img_query = apply_rotary_emb_qwen(img_query, img_freqs, false); - img_key = apply_rotary_emb_qwen(img_key, img_freqs, false); + auto img_query_freqs = img_freqs; + auto img_key_freqs = img_freqs; + if (regione) { + std::tie(img_query_freqs, img_key_freqs) = + regione->adjust_image_rope(img_freqs, img_key.size(1)); + } + img_query = apply_rotary_emb_qwen(img_query, img_query_freqs, false); + img_key = apply_rotary_emb_qwen(img_key, img_key_freqs, false); txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, false); txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, false); @@ -2352,15 +2364,21 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { block_attention_kwargs["attention_mask"] = joint_attention_mask; } + auto* regione = DiTCache::get_instance().regione(); + const bool regione_partial_sp_mode = + regione && regione->regione_is_partial_sp_mode(); if (::xllm::ParallelConfig::get_instance().sp_size() > 1) { - new_hidden_states = dit::sp_split_sequence(new_hidden_states, - /*dim=*/1, - parallel_args_.dit_sp_group_); + if (!regione_partial_sp_mode) { + new_hidden_states = + dit::sp_split_sequence(new_hidden_states, + /*dim=*/1, + parallel_args_.dit_sp_group_); + } new_encoder_hidden_states = dit::sp_split_sequence(new_encoder_hidden_states, /*dim=*/1, parallel_args_.dit_sp_group_); - if (modulate_index.defined()) { + if (modulate_index.defined() && !regione_partial_sp_mode) { modulate_index = dit::sp_split_sequence(modulate_index, /*dim=*/1, parallel_args_.dit_sp_group_); @@ -2373,6 +2391,10 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { bool use_step_cache = false; bool use_block_cache = false; + if (regione) { + regione->regione_prefetch_img_kv(0, use_cfg, new_hidden_states); + } + torch::Tensor original_hidden_states = new_hidden_states; torch::Tensor original_encoder_hidden_states = new_encoder_hidden_states; // Step start: prepare inputs (hidden_states, original_hidden_states) @@ -2390,6 +2412,10 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { CacheBlockIn blockin_before(index_block, block_in_before_map); use_block_cache = DiTCache::get_instance().on_before_block(blockin_before, use_cfg); + if (regione) { + regione->regione_set_current_block( + index_block, use_cfg, new_hidden_states); + } if (!use_block_cache) { std::tie(new_hidden_states, new_encoder_hidden_states) = @@ -2434,7 +2460,8 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { new_hidden_states = norm_out_->forward(new_hidden_states, temb); new_hidden_states = proj_out_->forward(new_hidden_states); - if (::xllm::ParallelConfig::get_instance().sp_size() > 1) { + if (::xllm::ParallelConfig::get_instance().sp_size() > 1 && + !regione_partial_sp_mode) { new_hidden_states = dit::sp_gather_sequence( new_hidden_states, /*dim=*/1, parallel_args_.dit_sp_group_); } From 9663dcbb2c6f57964fdf3f918e72f0c9f1ca5b40 Mon Sep 17 00:00:00 2001 From: "yuchuanhao.3" Date: Fri, 7 Aug 2026 15:22:19 +0800 Subject: [PATCH 2/3] feat: add regione gamma related dit cache changes. --- xllm/core/common/global_flags.h | 4 + xllm/core/framework/config/dit_config.cpp | 38 +++ xllm/core/framework/config/dit_config.h | 12 + .../framework/dit_cache/dit_cache_config.h | 10 + xllm/core/framework/dit_cache/regione.cpp | 233 ++++++++++++++++-- xllm/core/framework/dit_cache/regione.h | 48 +++- xllm/core/runtime/dit_worker_impl.cpp | 8 + .../pipelines/pipeline_qwenimage_edit_plus.h | 85 ++++++- .../dit/transformers/transformer_qwen_image.h | 38 +-- 9 files changed, 433 insertions(+), 43 deletions(-) diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index daebb8b4ce..0459e1458a 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -362,6 +362,10 @@ DECLARE_int64(dit_cache_end_blocks); DECLARE_string(dit_regione_refresh_steps); DECLARE_double(dit_regione_region_threshold); +DECLARE_double(dit_regione_cache_threshold); +DECLARE_bool(dit_regione_use_avd_gamma); +DECLARE_bool(dit_regione_erosion_dilation); +DECLARE_bool(dit_regione_profile); DECLARE_bool(dit_sp_communication_overlap); diff --git a/xllm/core/framework/config/dit_config.cpp b/xllm/core/framework/config/dit_config.cpp index a60a26213c..aec04be866 100644 --- a/xllm/core/framework/config/dit_config.cpp +++ b/xllm/core/framework/config/dit_config.cpp @@ -61,6 +61,28 @@ DEFINE_double(dit_regione_region_threshold, 0.80, "RegionE: cosine threshold for adaptive region partition."); +DEFINE_double(dit_regione_cache_threshold, + 0.03, + "RegionE: AVDCache error threshold δ (paper Eq.8). " + "Reuse velocity while 1-accumulate <= threshold. " + "Qwen-Image-Edit default in RegionE inplace.py is 0.03."); + +DEFINE_bool(dit_regione_use_avd_gamma, + true, + "RegionE: use AVDCache with gamma (paper/inplace.py method). " + "Uses the original diffusers 28-step gamma curve, linearly " + "upsampled/downsampled to the actual inference step count. " + "Set false to use fixed skip_interval instead."); + +DEFINE_bool(dit_regione_erosion_dilation, + true, + "RegionE: enable erosion/dilation for region mask cleanup."); + +DEFINE_bool(dit_regione_profile, + false, + "RegionE: print per-step timing breakdown for partial/full DiT and " + "K/V CPU offload."); + DEFINE_bool(dit_sp_communication_overlap, true, "Communication & Computation overlap for sequence parallel"); @@ -141,6 +163,10 @@ void DiTConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_blocks); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_refresh_steps); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_region_threshold); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_cache_threshold); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_use_avd_gamma); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_erosion_dilation); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_profile); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_laser_attention_enabled); @@ -169,6 +195,10 @@ void DiTConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_blocks); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_refresh_steps); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_region_threshold); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_cache_threshold); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_use_avd_gamma); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_erosion_dilation); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_profile); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_laser_attention_enabled); @@ -210,6 +240,14 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const { config_json, default_config, dit_regione_refresh_steps); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_regione_region_threshold); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_regione_cache_threshold); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_regione_use_avd_gamma); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_regione_erosion_dilation); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_regione_profile); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_sp_communication_overlap); 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 98beb313eb..79d847413c 100644 --- a/xllm/core/framework/config/dit_config.h +++ b/xllm/core/framework/config/dit_config.h @@ -53,6 +53,10 @@ class DiTConfig final { "dit_cache_end_blocks", "dit_regione_refresh_steps", "dit_regione_region_threshold", + "dit_regione_cache_threshold", + "dit_regione_use_avd_gamma", + "dit_regione_erosion_dilation", + "dit_regione_profile", "dit_sp_communication_overlap", "dit_debug_print", "dit_laser_attention_enabled", @@ -93,6 +97,14 @@ class DiTConfig final { PROPERTY(double, dit_regione_region_threshold) = 0.80; + PROPERTY(double, dit_regione_cache_threshold) = 0.03; + + PROPERTY(bool, dit_regione_use_avd_gamma) = true; + + PROPERTY(bool, dit_regione_erosion_dilation) = true; + + PROPERTY(bool, dit_regione_profile) = false; + PROPERTY(bool, dit_sp_communication_overlap) = true; PROPERTY(bool, dit_debug_print) = false; diff --git a/xllm/core/framework/dit_cache/dit_cache_config.h b/xllm/core/framework/dit_cache/dit_cache_config.h index 42d55202f5..06e930dc7d 100644 --- a/xllm/core/framework/dit_cache/dit_cache_config.h +++ b/xllm/core/framework/dit_cache/dit_cache_config.h @@ -57,10 +57,20 @@ struct FBCacheTaylorSeerOptions : public DiTBaseCacheOptions { }; struct RegionEOptions : public DiTBaseCacheOptions { + // Fallback fixed-interval AVD when gamma is disabled or step count + // mismatches. int64_t skip_interval_steps = 3; int64_t tail_steps = 1; std::vector refresh_steps = {16}; float region_threshold = 0.80f; + // AVDCache δ in paper Eq.8/9 / inplace.py cache_threshold (Qwen default + // 0.03). + float cache_threshold = 0.03f; + // Use fitted γ_t AVDCache (paper) instead of fixed skip_interval. + bool use_avd_gamma = true; + // Enable erosion/dilation morphological cleanup after ARP mask selection. + bool erosion_dilation = true; + bool profile = false; }; struct ResidualCacheOptions { diff --git a/xllm/core/framework/dit_cache/regione.cpp b/xllm/core/framework/dit_cache/regione.cpp index 9c09b858b5..6433df6d2d 100644 --- a/xllm/core/framework/dit_cache/regione.cpp +++ b/xllm/core/framework/dit_cache/regione.cpp @@ -15,6 +15,7 @@ limitations under the License. #include "regione.h" +#include #include #include @@ -90,12 +91,78 @@ bool RegionECache::regione_should_run_full_step(int64_t step) const { return regione_is_refresh_step(step); } -bool RegionECache::regione_should_compute_velocity(int64_t step) const { - if (!regione_enabled_) return true; - if (regione_should_run_full_step(step)) return true; - const auto interval = - std::max(1, config_.regione.skip_interval_steps); - return ((step - config_.regione.warmup_steps) % interval) == 0; +bool RegionECache::regione_should_compute_velocity(int64_t step, + double timestep, + double prev_timestep) { + if (!regione_enabled_) { + regione_avd_ratio_ = 1.0; + return true; + } + // STS / SMS / forced refresh: always run DiT and reset AVD accumulator. + // Original inplace.py also disables AVD at step == warmup (partial DiT still + // runs); keep that separate from regione_should_run_full_step so ARP→partial + // transition at warmup is unchanged. + if (regione_should_run_full_step(step) || + step <= config_.regione.warmup_steps) { + regione_avd_accumulate_ = 1.0; + regione_avd_ratio_ = 1.0; + return true; + } + + // Diffusers 28-step RegionE transition gamma (inplace.py), 27 values for + // transitions between 28 steps. Linearly upsample/downsample onto the + // actual (infer_steps - 1) transitions. + static constexpr double kRegionEGammaRef[] = { + 1.0186, 1.0241, 1.0236, 1.0205, 1.0298, 1.0221, 1.0248, 1.0246, 1.0269, + 1.0275, 1.0323, 1.0311, 1.0298, 1.0353, 1.0343, 1.0397, 1.0387, 1.0393, + 1.0404, 1.0458, 1.0507, 1.0418, 1.0518, 1.0426, 1.0311, 1.0068, 0.7628}; + static constexpr int64_t kGammaRefLen = static_cast( + sizeof(kRegionEGammaRef) / sizeof(kRegionEGammaRef[0])); + + auto sample_gamma = [&](int64_t cur_step) -> double { + const int64_t n_steps = + regione_infer_steps_ > 1 ? regione_infer_steps_ : 28; + const int64_t n_trans = std::max(1, n_steps - 1); + const int64_t idx = std::max(0, cur_step - 1); + const double pos = static_cast(idx) * + static_cast(kGammaRefLen - 1) / + static_cast(std::max(1, n_trans - 1)); + const int64_t lo = std::min(static_cast(pos), kGammaRefLen - 1); + const int64_t hi = std::min(lo + 1, kGammaRefLen - 1); + const double frac = pos - static_cast(lo); + return kRegionEGammaRef[lo] * (1.0 - frac) + kRegionEGammaRef[hi] * frac; + }; + + if (!config_.regione.use_avd_gamma || step < 1) { + const auto interval = + std::max(1, config_.regione.skip_interval_steps); + const bool compute = + ((step - config_.regione.warmup_steps) % interval) == 0; + regione_avd_ratio_ = 1.0; + if (compute) regione_avd_accumulate_ = 1.0; + return compute; + } + + // AVDCache (paper Eq.7-9 / inplace.py), step-count agnostic via resampled γ: + // ratio = gamma(step) * (1 + (t - t_prev) / 1000) + // accumulate *= ratio; error = 1 - accumulate + // reuse velocity while error <= cache_threshold and ratio < 1 + const double gamma = sample_gamma(step); + const double ratio = gamma * (1.0 + (timestep - prev_timestep) / 1000.0); + regione_avd_ratio_ = ratio; + + if (ratio >= 1.0) { + regione_avd_accumulate_ = 1.0; + return true; // recompute DiT + } + + regione_avd_accumulate_ *= ratio; + const double error = 1.0 - regione_avd_accumulate_; + if (error > static_cast(config_.regione.cache_threshold)) { + regione_avd_accumulate_ = 1.0; + return true; // recompute DiT + } + return false; // reuse velocity cache * ratio } bool RegionECache::regione_should_direct_unedited(int64_t step) const { @@ -162,6 +229,8 @@ void RegionECache::regione_prepare_inference( regione_edited_ids_ = torch::Tensor(); regione_unedited_ids_ = torch::Tensor(); regione_velocity_cache_ = torch::Tensor(); + regione_avd_accumulate_ = 1.0; + regione_avd_ratio_ = 1.0; regione_partial_mode_ = false; regione_local_edited_global_ids_ = torch::Tensor(); regione_local_edited_cache_ids_ = torch::Tensor(); @@ -199,18 +268,18 @@ torch::Tensor RegionECache::regione_scatter_ids(const torch::Tensor& values, } torch::Tensor RegionECache::regione_active_edited_ids() const { - if (regione_is_partial_sp_mode() && - regione_local_edited_global_ids_.defined()) { - return regione_local_edited_global_ids_; - } + // Partial+SP feeds the full edited set into every rank, then relies on the + // normal SP split/all-to-all path. Do not return per-shard local ids here — + // a rank with 0 local edited tokens would produce an empty sequence and + // break equal-length all_to_all_4D. return regione_edited_ids_; } torch::Tensor RegionECache::regione_kv_update_ids() const { - if (regione_is_partial_sp_mode() && - regione_local_edited_cache_ids_.defined()) { - return regione_local_edited_cache_ids_; - } + // Full-mode image K/V is stored after SP QKV all-to-all, so the cache is + // indexed by global sequence positions. Partial-step `key` is likewise + // assembled across SP ranks into global edited_ids_ order — update with + // those global indices rather than per-shard local cache ids. return regione_edited_ids_; } @@ -238,7 +307,11 @@ torch::Tensor RegionECache::regione_scatter_unedited( torch::Tensor RegionECache::regione_gather_query_rope( const torch::Tensor& image_rope) const { - auto ids = regione_active_edited_ids(); + // SP attention (non-CMO) runs QKV all-to-all before RoPE, so img_query holds + // every rank's edited tokens concatenated in contiguous-shard order. With + // sorted edited_ids_ that order matches the global edited set — gather by + // global ids, not the per-rank local edited subset. + auto ids = regione_edited_ids_; if (!image_rope.defined() || !ids.defined()) return image_rope; if (image_rope.size(0) == ids.numel()) return image_rope; return regione_gather_ids(image_rope, ids, 0); @@ -248,6 +321,9 @@ torch::Tensor RegionECache::regione_gather_key_rope( const torch::Tensor& image_rope, int64_t key_len) const { if (!image_rope.defined() || !regione_is_partial_sp_mode()) return image_rope; + // Patched image K is the full-sequence cache (post all-to-all store). Prefer + // a length-matched rope; only fall back to per-shard gather when the cache + // is still local-shard sized. if (image_rope.size(0) == key_len) return image_rope; if (!regione_local_image_global_ids_.defined() || regione_local_image_global_ids_.numel() != key_len) { @@ -260,7 +336,15 @@ torch::Tensor RegionECache::regione_local_update_mask( const torch::Tensor& base) const { if (!base.defined()) return torch::Tensor(); auto mask = torch::zeros({base.size(0), base.size(1), 1}, base.options()); - auto ids = regione_active_edited_ids(); + // SP partial DiT runs the full edited set on every rank; when reducing + // latent updates, each rank must only own its image shard's edited tokens + // so reduce(SUM) does not double-count. + torch::Tensor ids; + if (regione_sp_size_ > 1 && regione_local_edited_global_ids_.defined()) { + ids = regione_local_edited_global_ids_; + } else { + ids = regione_edited_ids_; + } if (ids.defined() && ids.numel() > 0) { auto ones = torch::ones({base.size(0), ids.numel(), 1}, base.options()); mask.index_copy_(1, regione_normalize_ids(ids, base.device()), ones); @@ -305,7 +389,8 @@ void RegionECache::regione_select_regions(const torch::Tensor& sample, torch::sqrt(torch::sum(condition * condition, -1, true)).clamp_min(1e-6); auto similarity = torch::sum(estimate_norm * condition_norm, -1); auto selected_mask = similarity <= config_.regione.region_threshold; - if (regione_grid_h_ > 0 && regione_grid_w_ > 0 && + if (config_.regione.erosion_dilation && regione_grid_h_ > 0 && + regione_grid_w_ > 0 && regione_grid_h_ * regione_grid_w_ == sample.size(1)) { auto mask2d = selected_mask.to(torch::kFloat) .view({1, 1, regione_grid_h_, regione_grid_w_}); @@ -366,7 +451,9 @@ torch::Tensor RegionECache::regione_velocity_cache() const { return regione_velocity_cache_; } -RegionEStepPlan RegionECache::begin_step(int64_t step) { +RegionEStepPlan RegionECache::begin_step(int64_t step, + double timestep, + double prev_timestep) { RegionEStepPlan plan; plan.enabled = regione_enabled_; if (!regione_enabled_) return plan; @@ -375,8 +462,9 @@ RegionEStepPlan RegionECache::begin_step(int64_t step) { plan.full_step = regione_should_run_full_step(step); plan.partial_step = regione_has_regions() && !plan.full_step; regione_set_partial_mode(plan.partial_step); - plan.use_velocity_cache = !regione_should_compute_velocity(step) && - regione_velocity_cache_.defined(); + plan.use_velocity_cache = + !regione_should_compute_velocity(step, timestep, prev_timestep) && + regione_velocity_cache_.defined(); plan.run_partition = !regione_has_regions() && step == regione_warmup_steps() - 1 && !plan.use_velocity_cache; @@ -402,7 +490,8 @@ RegionEStepInput RegionECache::prepare_step_input( } input.use_cached_velocity = plan.use_velocity_cache; if (plan.use_velocity_cache) { - input.cached_velocity = regione_velocity_cache_; + // inplace.py: noise_pred = cache * ratio + input.cached_velocity = regione_velocity_cache_ * regione_avd_ratio_; if (plan.partial_step && input.cached_velocity.defined() && input.cached_velocity.size(1) != input.step_latents.size(1)) { input.cached_velocity = regione_gather_edited(input.cached_velocity); @@ -691,11 +780,111 @@ std::pair RegionECache::regione_patch_img_kv( } auto update_ids = regione_kv_update_ids(); if (update_ids.defined() && key.dim() >= 2 && - full_key.size(0) == key.size(0) && key.size(1) == update_ids.numel()) { + full_key.size(0) == key.size(0)) { + CHECK_EQ(key.size(1), update_ids.numel()) + << "RegionE partial KV length must match edited_ids before scatter; " + "key_len=" + << key.size(1) << " edited=" << update_ids.numel() + << " full_key_len=" << full_key.size(1) + << " (likely SP pad was applied after KV patch)"; full_key = regione_scatter_ids(key, update_ids, full_key, 1); full_value = regione_scatter_ids(value, update_ids, full_value, 1); } return {full_key, full_value}; } +bool RegionECache::regione_profile_enabled() const { + return regione_enabled_ && config_.regione.profile; +} + +void RegionECache::regione_profile_reset_step(int64_t step, + bool partial_step, + bool full_step, + bool velocity_cache, + int64_t step_tokens, + int64_t full_tokens) { + regione_profile_step_ = step; + regione_profile_partial_step_ = partial_step; + regione_profile_full_step_ = full_step; + regione_profile_velocity_cache_ = velocity_cache; + regione_profile_step_tokens_ = step_tokens; + regione_profile_full_tokens_ = full_tokens; + regione_profile_kv_store_count_ = 0; + regione_profile_prefetch_issue_count_ = 0; + regione_profile_prefetch_hit_count_ = 0; + regione_profile_prefetch_miss_count_ = 0; + regione_profile_fallback_h2d_count_ = 0; + regione_profile_patch_scatter_count_ = 0; + regione_profile_kv_store_cpu_ms_ = 0.0; + regione_profile_prefetch_issue_ms_ = 0.0; + regione_profile_prefetch_wait_ms_ = 0.0; + regione_profile_fallback_h2d_ms_ = 0.0; + regione_profile_patch_scatter_ms_ = 0.0; +} + +void RegionECache::regione_profile_log_step(double transformer_ms, + double arp_ms, + double scheduler_ms, + double total_ms) const { + if (!regione_profile_enabled()) return; + LOG(INFO) << "[RegionEProfile] step=" << regione_profile_step_ << " mode=" + << (regione_profile_partial_step_ + ? "partial" + : (regione_profile_full_step_ ? "full" : "reuse")) + << " velocity_cache=" << regione_profile_velocity_cache_ + << " tokens=" << regione_profile_step_tokens_ << "/" + << regione_profile_full_tokens_ << " total_ms=" << total_ms + << " transformer_ms=" << transformer_ms + << " scheduler_ms=" << scheduler_ms << " arp_ms=" << arp_ms + << " kv_store_cpu_ms=" << regione_profile_kv_store_cpu_ms_ + << " kv_store_count=" << regione_profile_kv_store_count_ + << " kv_prefetch_issue_ms=" << regione_profile_prefetch_issue_ms_ + << " kv_prefetch_issue_count=" + << regione_profile_prefetch_issue_count_ + << " kv_prefetch_wait_ms=" << regione_profile_prefetch_wait_ms_ + << " kv_prefetch_hit_count=" << regione_profile_prefetch_hit_count_ + << " kv_prefetch_miss_count=" + << regione_profile_prefetch_miss_count_ + << " kv_fallback_h2d_ms=" << regione_profile_fallback_h2d_ms_ + << " kv_fallback_h2d_count=" << regione_profile_fallback_h2d_count_ + << " kv_patch_scatter_ms=" << regione_profile_patch_scatter_ms_ + << " kv_patch_scatter_count=" + << regione_profile_patch_scatter_count_; +} + +void RegionECache::regione_profile_add_kv_store(double ms) { + if (!regione_profile_enabled()) return; + regione_profile_kv_store_cpu_ms_ += ms; + ++regione_profile_kv_store_count_; +} + +void RegionECache::regione_profile_add_prefetch_issue(double ms) { + if (!regione_profile_enabled()) return; + regione_profile_prefetch_issue_ms_ += ms; + ++regione_profile_prefetch_issue_count_; +} + +void RegionECache::regione_profile_add_prefetch_hit(double wait_ms) { + if (!regione_profile_enabled()) return; + regione_profile_prefetch_wait_ms_ += wait_ms; + ++regione_profile_prefetch_hit_count_; +} + +void RegionECache::regione_profile_add_prefetch_miss() { + if (!regione_profile_enabled()) return; + ++regione_profile_prefetch_miss_count_; +} + +void RegionECache::regione_profile_add_fallback_h2d(double ms) { + if (!regione_profile_enabled()) return; + regione_profile_fallback_h2d_ms_ += ms; + ++regione_profile_fallback_h2d_count_; +} + +void RegionECache::regione_profile_add_patch_scatter(double ms) { + if (!regione_profile_enabled()) return; + regione_profile_patch_scatter_ms_ += ms; + ++regione_profile_patch_scatter_count_; +} + } // namespace xllm diff --git a/xllm/core/framework/dit_cache/regione.h b/xllm/core/framework/dit_cache/regione.h index 4de3840674..3e013791a2 100644 --- a/xllm/core/framework/dit_cache/regione.h +++ b/xllm/core/framework/dit_cache/regione.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include @@ -78,7 +79,11 @@ class RegionECache { bool regione_has_regions() const { return regione_edited_ids_.defined(); } bool regione_is_partial_mode() const { return regione_partial_mode_; } bool regione_is_partial_sp_mode() const; - bool regione_should_compute_velocity(int64_t step) const; + // AVDCache: true = run DiT, false = reuse velocity cache * scale. + bool regione_should_compute_velocity(int64_t step, + double timestep, + double prev_timestep); + double regione_velocity_scale() const { return regione_avd_ratio_; } bool regione_should_run_full_step(int64_t step) const; bool regione_should_direct_unedited(int64_t step) const; int64_t regione_next_direct_step(int64_t step) const; @@ -89,7 +94,9 @@ class RegionECache { int64_t sp_rank = 0, int64_t sp_size = 1); - RegionEStepPlan begin_step(int64_t step); + RegionEStepPlan begin_step(int64_t step, + double timestep = 0.0, + double prev_timestep = 0.0); RegionEStepInput prepare_step_input( const torch::Tensor& latents, const torch::Tensor& condition_latents, @@ -128,6 +135,24 @@ class RegionECache { const torch::Tensor& image_rope, int64_t key_len) const; + bool regione_profile_enabled() const; + void regione_profile_reset_step(int64_t step, + bool partial_step, + bool full_step, + bool velocity_cache, + int64_t step_tokens, + int64_t full_tokens); + void regione_profile_log_step(double transformer_ms, + double arp_ms, + double scheduler_ms, + double total_ms) const; + void regione_profile_add_kv_store(double ms); + void regione_profile_add_prefetch_issue(double ms); + void regione_profile_add_prefetch_hit(double wait_ms); + void regione_profile_add_prefetch_miss(); + void regione_profile_add_fallback_h2d(double ms); + void regione_profile_add_patch_scatter(double ms); + private: void regione_select_regions(const torch::Tensor& sample, const torch::Tensor& model_output, @@ -209,6 +234,25 @@ class RegionECache { torch::Tensor regione_local_edited_cache_ids_; torch::Tensor regione_local_image_global_ids_; torch::Tensor regione_velocity_cache_; + double regione_avd_accumulate_ = 1.0; + double regione_avd_ratio_ = 1.0; + int64_t regione_profile_step_ = -1; + bool regione_profile_partial_step_ = false; + bool regione_profile_full_step_ = false; + bool regione_profile_velocity_cache_ = false; + int64_t regione_profile_step_tokens_ = 0; + int64_t regione_profile_full_tokens_ = 0; + int64_t regione_profile_kv_store_count_ = 0; + int64_t regione_profile_prefetch_issue_count_ = 0; + int64_t regione_profile_prefetch_hit_count_ = 0; + int64_t regione_profile_prefetch_miss_count_ = 0; + int64_t regione_profile_fallback_h2d_count_ = 0; + int64_t regione_profile_patch_scatter_count_ = 0; + double regione_profile_kv_store_cpu_ms_ = 0.0; + double regione_profile_prefetch_issue_ms_ = 0.0; + double regione_profile_prefetch_wait_ms_ = 0.0; + double regione_profile_fallback_h2d_ms_ = 0.0; + double regione_profile_patch_scatter_ms_ = 0.0; std::vector regione_k_cache_cpu_; std::vector regione_v_cache_cpu_; std::vector regione_cond_k_cache_cpu_; diff --git a/xllm/core/runtime/dit_worker_impl.cpp b/xllm/core/runtime/dit_worker_impl.cpp index 13b04954ac..98df223867 100644 --- a/xllm/core/runtime/dit_worker_impl.cpp +++ b/xllm/core/runtime/dit_worker_impl.cpp @@ -105,6 +105,14 @@ DiTCacheConfig parse_dit_cache_from_flags() { ::xllm::DiTConfig::get_instance().dit_regione_refresh_steps()); cache_config.regione.region_threshold = ::xllm::DiTConfig::get_instance().dit_regione_region_threshold(); + cache_config.regione.cache_threshold = static_cast( + ::xllm::DiTConfig::get_instance().dit_regione_cache_threshold()); + cache_config.regione.use_avd_gamma = + ::xllm::DiTConfig::get_instance().dit_regione_use_avd_gamma(); + cache_config.regione.erosion_dilation = + ::xllm::DiTConfig::get_instance().dit_regione_erosion_dilation(); + cache_config.regione.profile = + ::xllm::DiTConfig::get_instance().dit_regione_profile(); } else if (::xllm::DiTConfig::get_instance().dit_cache_policy() == "None") { cache_config.selected_policy = PolicyType::None; } diff --git a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h index 0662fbcd74..195f2623aa 100644 --- a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h +++ b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h @@ -1057,6 +1057,19 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { : image_rotary_emb_pos; auto* regione = DiTCache::get_instance().regione(); + const bool regione_profile_enabled = + regione != nullptr && regione->regione_profile_enabled(); + auto regione_profile_now = []() { + return std::chrono::steady_clock::now(); + }; + auto regione_profile_ms_since = + [](const std::chrono::steady_clock::time_point& start) { + return std::chrono::duration( + std::chrono::steady_clock::now() - start) + .count(); + }; + auto regione_dit_loop_start = regione_profile_now(); + if (regione) { const auto sp_group = parallel_args_.dit_sp_group_; const int64_t sp_rank = sp_group ? sp_group->rank() : 0; @@ -1072,11 +1085,24 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { for (int64_t i = 0; i < timesteps.size(0); ++i) { auto t = timesteps[i]; current_timestep_ = t; + auto regione_step_profile_start = regione_profile_now(); + double regione_transformer_ms = 0.0; + double regione_arp_ms = 0.0; + double regione_scheduler_ms = 0.0; + + double prev_t_value = 0.0; + double t_value = 0.0; + if (regione) { + t_value = t.item(); + if (i > 0) { + prev_t_value = timesteps[i - 1].item(); + } + } RegionEStepPlan regione_plan; RegionEStepInput regione_input; if (regione) { - regione_plan = regione->begin_step(i); + regione_plan = regione->begin_step(i, t_value, prev_t_value); regione_input = regione->prepare_step_input( final_latents, image_latents, main_shape, regione_plan); } @@ -1095,9 +1121,24 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { auto timestep_expanded = t.expand({step_latents.size(0)}).to(step_latents.dtype()); + if (regione_profile_enabled) { + regione->regione_profile_reset_step( + i, + regione_plan.partial_step, + regione_plan.full_step, + regione_plan.use_velocity_cache, + step_latents.defined() && step_latents.dim() > 1 + ? step_latents.size(1) + : 0, + final_latents.defined() && final_latents.dim() > 1 + ? final_latents.size(1) + : 0); + } + torch::Tensor noise_pred; torch::Tensor neg_noise_pred; torch::Tensor pos_neg_noise_preds; + auto regione_transformer_start = regione_profile_now(); if (regione_plan.use_velocity_cache) { noise_pred = regione_input.cached_velocity; } else if (::xllm::ParallelConfig::get_instance().cfg_size() == 2 && @@ -1173,16 +1214,25 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { noise_pred = comb_pred * (cond_norm / noise_norm); } } + if (regione_profile_enabled) { + regione_transformer_ms = + regione_profile_ms_since(regione_transformer_start); + } if (regione_plan.enabled) { + auto regione_arp_start = regione_profile_now(); regione->observe_velocity( final_latents, noise_pred, scheduler_->sigmas(), i, regione_plan); if (regione_plan.run_partition) { + if (regione_profile_enabled) { + regione_arp_ms = regione_profile_ms_since(regione_arp_start); + } regione_plan.direct_unedited = regione->regione_should_direct_unedited(i); } } + auto regione_scheduler_start = regione_profile_now(); auto latents_dtype = final_latents.dtype(); if (regione_plan.enabled) { if (regione_plan.partial_step) { @@ -1220,14 +1270,33 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { if (final_latents.dtype() != latents_dtype) { final_latents = final_latents.to(latents_dtype); } + if (regione_profile_enabled) { + regione_scheduler_ms = + regione_profile_ms_since(regione_scheduler_start); + regione->regione_profile_log_step( + regione_transformer_ms, + regione_arp_ms, + regione_scheduler_ms, + regione_profile_ms_since(regione_step_profile_start)); + } + } + if (regione_profile_enabled) { + LOG(INFO) << "[RegionEProfile] dit_loop_total_ms=" + << regione_profile_ms_since(regione_dit_loop_start); } current_timestep_ = torch::Tensor(); torch::Tensor output_image; + auto regione_vae_stage_start = regione_profile_now(); auto unpacked_latents = _unpack_latents(final_latents, height, width, vae_scale_factor_) .to(dtype_); + if (regione_profile_enabled) { + LOG(INFO) << "[RegionEProfile] vae_unpack_ms=" + << regione_profile_ms_since(regione_vae_stage_start); + } + regione_vae_stage_start = regione_profile_now(); auto latents_mean = torch::tensor(vae_model_args_.latents_mean(), torch::kDouble); latents_mean = @@ -1239,8 +1308,22 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { latents_std.view({1, latent_channels_, 1, 1, 1}).to(device_, dtype_); unpacked_latents = unpacked_latents / latents_std + latents_mean; + if (regione_profile_enabled) { + LOG(INFO) << "[RegionEProfile] vae_latent_norm_ms=" + << regione_profile_ms_since(regione_vae_stage_start); + } + regione_vae_stage_start = regione_profile_now(); output_image = vae_->decode(unpacked_latents).sample.squeeze(2); + if (regione_profile_enabled) { + LOG(INFO) << "[RegionEProfile] vae_decode_ms=" + << regione_profile_ms_since(regione_vae_stage_start); + } + regione_vae_stage_start = regione_profile_now(); output_image = vae_image_processor_->postprocess(output_image); + if (regione_profile_enabled) { + LOG(INFO) << "[RegionEProfile] vae_postprocess_ms=" + << regione_profile_ms_since(regione_vae_stage_start); + } auto output_chunks = torch::chunk(output_image, batch_size, /*dim=*/0); DiTForwardOutput out; out.tensors = std::move(output_chunks); diff --git a/xllm/models/dit/transformers/transformer_qwen_image.h b/xllm/models/dit/transformers/transformer_qwen_image.h index f0e24411fb..d70ed3b776 100644 --- a/xllm/models/dit/transformers/transformer_qwen_image.h +++ b/xllm/models/dit/transformers/transformer_qwen_image.h @@ -1584,6 +1584,17 @@ class QwenDoubleStreamAttnProcessor2_0Impl : public torch::nn::Module { auto img_query = attn_->to_q_->forward(hidden_states); auto img_key = attn_->to_k_->forward(hidden_states); auto img_value = attn_->to_v_->forward(hidden_states); + + // Unpad before RegionE KV patch/store. Partial steps expand K/V to the + // full cached sequence; applying the SP pad trim afterward would shrink + // that full K/V by the edited-sequence pad length and break RoPE. + xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( + img_query, /*tensor_name=*/"hidden_states", /*dim=*/1); + xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( + img_key, /*tensor_name=*/"hidden_states", /*dim=*/1); + xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( + img_value, /*tensor_name=*/"hidden_states", /*dim=*/1); + auto* regione = DiTCache::get_instance().regione(); if (regione) { std::tie(img_key, img_value) = @@ -1631,13 +1642,6 @@ class QwenDoubleStreamAttnProcessor2_0Impl : public torch::nn::Module { xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( txt_value, /*tensor_name=*/"encoder_hidden_states", /*dim=*/1); - xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( - img_query, /*tensor_name=*/"hidden_states", /*dim=*/1); - xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( - img_key, /*tensor_name=*/"hidden_states", /*dim=*/1); - xllm::dit::SequenceParallelPadManager::get_instance().unpad_tensor( - img_value, /*tensor_name=*/"hidden_states", /*dim=*/1); - auto img_query_freqs = img_freqs; auto img_key_freqs = img_freqs; if (regione) { @@ -2365,20 +2369,19 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { } auto* regione = DiTCache::get_instance().regione(); - const bool regione_partial_sp_mode = - regione && regione->regione_is_partial_sp_mode(); + // RegionE partial+SP now replicates the global edited sequence on every + // rank (see RegionECache::regione_active_edited_ids), so the normal SP + // split/gather path applies — do not skip it for partial mode. if (::xllm::ParallelConfig::get_instance().sp_size() > 1) { - if (!regione_partial_sp_mode) { - new_hidden_states = - dit::sp_split_sequence(new_hidden_states, - /*dim=*/1, - parallel_args_.dit_sp_group_); - } + new_hidden_states = + dit::sp_split_sequence(new_hidden_states, + /*dim=*/1, + parallel_args_.dit_sp_group_); new_encoder_hidden_states = dit::sp_split_sequence(new_encoder_hidden_states, /*dim=*/1, parallel_args_.dit_sp_group_); - if (modulate_index.defined() && !regione_partial_sp_mode) { + if (modulate_index.defined()) { modulate_index = dit::sp_split_sequence(modulate_index, /*dim=*/1, parallel_args_.dit_sp_group_); @@ -2460,8 +2463,7 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { new_hidden_states = norm_out_->forward(new_hidden_states, temb); new_hidden_states = proj_out_->forward(new_hidden_states); - if (::xllm::ParallelConfig::get_instance().sp_size() > 1 && - !regione_partial_sp_mode) { + if (::xllm::ParallelConfig::get_instance().sp_size() > 1) { new_hidden_states = dit::sp_gather_sequence( new_hidden_states, /*dim=*/1, parallel_args_.dit_sp_group_); } From c01bad883b7847c826732ed683e06905a5552ab6 Mon Sep 17 00:00:00 2001 From: kongweiqian Date: Fri, 14 Aug 2026 16:48:52 +0800 Subject: [PATCH 3/3] refactor: simplify regione avd gamma and remove profile flag. --- xllm/core/common/global_flags.h | 1 - xllm/core/framework/config/dit_config.cpp | 13 +- xllm/core/framework/config/dit_config.h | 5 +- .../framework/dit_cache/dit_cache_config.h | 6 +- xllm/core/framework/dit_cache/regione.cpp | 145 +++++------------- xllm/core/framework/dit_cache/regione.h | 39 +---- xllm/core/runtime/dit_worker_impl.cpp | 2 - .../pipelines/pipeline_qwenimage_edit_plus.h | 75 --------- 8 files changed, 49 insertions(+), 237 deletions(-) diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 0459e1458a..25edaf4303 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -365,7 +365,6 @@ DECLARE_double(dit_regione_region_threshold); DECLARE_double(dit_regione_cache_threshold); DECLARE_bool(dit_regione_use_avd_gamma); DECLARE_bool(dit_regione_erosion_dilation); -DECLARE_bool(dit_regione_profile); DECLARE_bool(dit_sp_communication_overlap); diff --git a/xllm/core/framework/config/dit_config.cpp b/xllm/core/framework/config/dit_config.cpp index aec04be866..9f8ce4768d 100644 --- a/xllm/core/framework/config/dit_config.cpp +++ b/xllm/core/framework/config/dit_config.cpp @@ -62,10 +62,10 @@ DEFINE_double(dit_regione_region_threshold, "RegionE: cosine threshold for adaptive region partition."); DEFINE_double(dit_regione_cache_threshold, - 0.03, + 0.02, "RegionE: AVDCache error threshold δ (paper Eq.8). " "Reuse velocity while 1-accumulate <= threshold. " - "Qwen-Image-Edit default in RegionE inplace.py is 0.03."); + "Default is 0.02."); DEFINE_bool(dit_regione_use_avd_gamma, true, @@ -78,11 +78,6 @@ DEFINE_bool(dit_regione_erosion_dilation, true, "RegionE: enable erosion/dilation for region mask cleanup."); -DEFINE_bool(dit_regione_profile, - false, - "RegionE: print per-step timing breakdown for partial/full DiT and " - "K/V CPU offload."); - DEFINE_bool(dit_sp_communication_overlap, true, "Communication & Computation overlap for sequence parallel"); @@ -166,7 +161,6 @@ void DiTConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_cache_threshold); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_use_avd_gamma); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_erosion_dilation); - XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_regione_profile); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_laser_attention_enabled); @@ -198,7 +192,6 @@ void DiTConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_cache_threshold); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_use_avd_gamma); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_erosion_dilation); - XLLM_CONFIG_ASSIGN_FROM_JSON(dit_regione_profile); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_laser_attention_enabled); @@ -246,8 +239,6 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const { config_json, default_config, dit_regione_use_avd_gamma); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_regione_erosion_dilation); - APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( - config_json, default_config, dit_regione_profile); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_sp_communication_overlap); 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 79d847413c..8693b1cb42 100644 --- a/xllm/core/framework/config/dit_config.h +++ b/xllm/core/framework/config/dit_config.h @@ -56,7 +56,6 @@ class DiTConfig final { "dit_regione_cache_threshold", "dit_regione_use_avd_gamma", "dit_regione_erosion_dilation", - "dit_regione_profile", "dit_sp_communication_overlap", "dit_debug_print", "dit_laser_attention_enabled", @@ -97,14 +96,12 @@ class DiTConfig final { PROPERTY(double, dit_regione_region_threshold) = 0.80; - PROPERTY(double, dit_regione_cache_threshold) = 0.03; + PROPERTY(double, dit_regione_cache_threshold) = 0.02; PROPERTY(bool, dit_regione_use_avd_gamma) = true; PROPERTY(bool, dit_regione_erosion_dilation) = true; - PROPERTY(bool, dit_regione_profile) = false; - PROPERTY(bool, dit_sp_communication_overlap) = true; PROPERTY(bool, dit_debug_print) = false; diff --git a/xllm/core/framework/dit_cache/dit_cache_config.h b/xllm/core/framework/dit_cache/dit_cache_config.h index 06e930dc7d..851bf3785f 100644 --- a/xllm/core/framework/dit_cache/dit_cache_config.h +++ b/xllm/core/framework/dit_cache/dit_cache_config.h @@ -63,14 +63,12 @@ struct RegionEOptions : public DiTBaseCacheOptions { int64_t tail_steps = 1; std::vector refresh_steps = {16}; float region_threshold = 0.80f; - // AVDCache δ in paper Eq.8/9 / inplace.py cache_threshold (Qwen default - // 0.03). - float cache_threshold = 0.03f; + // AVDCache δ in paper Eq.8/9. + float cache_threshold = 0.02f; // Use fitted γ_t AVDCache (paper) instead of fixed skip_interval. bool use_avd_gamma = true; // Enable erosion/dilation morphological cleanup after ARP mask selection. bool erosion_dilation = true; - bool profile = false; }; struct ResidualCacheOptions { diff --git a/xllm/core/framework/dit_cache/regione.cpp b/xllm/core/framework/dit_cache/regione.cpp index 6433df6d2d..86da2864d6 100644 --- a/xllm/core/framework/dit_cache/regione.cpp +++ b/xllm/core/framework/dit_cache/regione.cpp @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include namespace xllm { namespace { @@ -61,6 +62,12 @@ void RegionECache::init(const DiTCacheConfig& cfg) { regione_local_edited_global_ids_ = torch::Tensor(); regione_local_edited_cache_ids_ = torch::Tensor(); regione_local_image_global_ids_ = torch::Tensor(); + regione_avd_accumulate_ = 1.0; + regione_avd_ratio_ = 1.0; + regione_avd_raw_gamma_ = 1.0; + regione_avd_gamma_ = 1.0; + regione_avd_gamma_exponent_ = 1.0; + regione_avd_error_ = 0.0; regione_clear_all_prefetch_slots(); } @@ -96,6 +103,10 @@ bool RegionECache::regione_should_compute_velocity(int64_t step, double prev_timestep) { if (!regione_enabled_) { regione_avd_ratio_ = 1.0; + regione_avd_raw_gamma_ = 1.0; + regione_avd_gamma_ = 1.0; + regione_avd_gamma_exponent_ = 1.0; + regione_avd_error_ = 0.0; return true; } // STS / SMS / forced refresh: always run DiT and reset AVD accumulator. @@ -106,12 +117,17 @@ bool RegionECache::regione_should_compute_velocity(int64_t step, step <= config_.regione.warmup_steps) { regione_avd_accumulate_ = 1.0; regione_avd_ratio_ = 1.0; + regione_avd_raw_gamma_ = 1.0; + regione_avd_gamma_ = 1.0; + regione_avd_gamma_exponent_ = 1.0; + regione_avd_error_ = 0.0; return true; } // Diffusers 28-step RegionE transition gamma (inplace.py), 27 values for - // transitions between 28 steps. Linearly upsample/downsample onto the - // actual (infer_steps - 1) transitions. + // transitions between 28 steps. Linearly sample it onto the actual + // (infer_steps - 1) transitions, then temper the per-step gamma when the + // current run uses more transitions than the reference schedule. static constexpr double kRegionEGammaRef[] = { 1.0186, 1.0241, 1.0236, 1.0205, 1.0298, 1.0221, 1.0248, 1.0246, 1.0269, 1.0275, 1.0323, 1.0311, 1.0298, 1.0353, 1.0343, 1.0397, 1.0387, 1.0393, @@ -139,25 +155,34 @@ bool RegionECache::regione_should_compute_velocity(int64_t step, const bool compute = ((step - config_.regione.warmup_steps) % interval) == 0; regione_avd_ratio_ = 1.0; + regione_avd_raw_gamma_ = 1.0; + regione_avd_gamma_ = 1.0; + regione_avd_gamma_exponent_ = 1.0; + regione_avd_error_ = 0.0; if (compute) regione_avd_accumulate_ = 1.0; return compute; } - // AVDCache (paper Eq.7-9 / inplace.py), step-count agnostic via resampled γ: + // AVDCache (paper Eq.7-9 / inplace.py), timestep-delta normalized γ: // ratio = gamma(step) * (1 + (t - t_prev) / 1000) // accumulate *= ratio; error = 1 - accumulate - // reuse velocity while error <= cache_threshold and ratio < 1 - const double gamma = sample_gamma(step); + // reuse velocity while error <= cache_threshold + const double raw_gamma = sample_gamma(step); + const double ref_timestep_delta = 1000.0 / static_cast(kGammaRefLen); + const double local_timestep_delta = std::abs(timestep - prev_timestep); + const double raw_gamma_exponent = local_timestep_delta / ref_timestep_delta; + const double gamma_exponent = + std::max(0.25, std::min(3.0, raw_gamma_exponent)); + const double gamma = std::pow(raw_gamma, gamma_exponent); const double ratio = gamma * (1.0 + (timestep - prev_timestep) / 1000.0); regione_avd_ratio_ = ratio; - - if (ratio >= 1.0) { - regione_avd_accumulate_ = 1.0; - return true; // recompute DiT - } + regione_avd_raw_gamma_ = raw_gamma; + regione_avd_gamma_ = gamma; + regione_avd_gamma_exponent_ = gamma_exponent; regione_avd_accumulate_ *= ratio; - const double error = 1.0 - regione_avd_accumulate_; + const double error = std::abs(1.0 - regione_avd_accumulate_); + regione_avd_error_ = error; if (error > static_cast(config_.regione.cache_threshold)) { regione_avd_accumulate_ = 1.0; return true; // recompute DiT @@ -231,6 +256,10 @@ void RegionECache::regione_prepare_inference( regione_velocity_cache_ = torch::Tensor(); regione_avd_accumulate_ = 1.0; regione_avd_ratio_ = 1.0; + regione_avd_raw_gamma_ = 1.0; + regione_avd_gamma_ = 1.0; + regione_avd_gamma_exponent_ = 1.0; + regione_avd_error_ = 0.0; regione_partial_mode_ = false; regione_local_edited_global_ids_ = torch::Tensor(); regione_local_edited_cache_ids_ = torch::Tensor(); @@ -793,98 +822,4 @@ std::pair RegionECache::regione_patch_img_kv( return {full_key, full_value}; } -bool RegionECache::regione_profile_enabled() const { - return regione_enabled_ && config_.regione.profile; -} - -void RegionECache::regione_profile_reset_step(int64_t step, - bool partial_step, - bool full_step, - bool velocity_cache, - int64_t step_tokens, - int64_t full_tokens) { - regione_profile_step_ = step; - regione_profile_partial_step_ = partial_step; - regione_profile_full_step_ = full_step; - regione_profile_velocity_cache_ = velocity_cache; - regione_profile_step_tokens_ = step_tokens; - regione_profile_full_tokens_ = full_tokens; - regione_profile_kv_store_count_ = 0; - regione_profile_prefetch_issue_count_ = 0; - regione_profile_prefetch_hit_count_ = 0; - regione_profile_prefetch_miss_count_ = 0; - regione_profile_fallback_h2d_count_ = 0; - regione_profile_patch_scatter_count_ = 0; - regione_profile_kv_store_cpu_ms_ = 0.0; - regione_profile_prefetch_issue_ms_ = 0.0; - regione_profile_prefetch_wait_ms_ = 0.0; - regione_profile_fallback_h2d_ms_ = 0.0; - regione_profile_patch_scatter_ms_ = 0.0; -} - -void RegionECache::regione_profile_log_step(double transformer_ms, - double arp_ms, - double scheduler_ms, - double total_ms) const { - if (!regione_profile_enabled()) return; - LOG(INFO) << "[RegionEProfile] step=" << regione_profile_step_ << " mode=" - << (regione_profile_partial_step_ - ? "partial" - : (regione_profile_full_step_ ? "full" : "reuse")) - << " velocity_cache=" << regione_profile_velocity_cache_ - << " tokens=" << regione_profile_step_tokens_ << "/" - << regione_profile_full_tokens_ << " total_ms=" << total_ms - << " transformer_ms=" << transformer_ms - << " scheduler_ms=" << scheduler_ms << " arp_ms=" << arp_ms - << " kv_store_cpu_ms=" << regione_profile_kv_store_cpu_ms_ - << " kv_store_count=" << regione_profile_kv_store_count_ - << " kv_prefetch_issue_ms=" << regione_profile_prefetch_issue_ms_ - << " kv_prefetch_issue_count=" - << regione_profile_prefetch_issue_count_ - << " kv_prefetch_wait_ms=" << regione_profile_prefetch_wait_ms_ - << " kv_prefetch_hit_count=" << regione_profile_prefetch_hit_count_ - << " kv_prefetch_miss_count=" - << regione_profile_prefetch_miss_count_ - << " kv_fallback_h2d_ms=" << regione_profile_fallback_h2d_ms_ - << " kv_fallback_h2d_count=" << regione_profile_fallback_h2d_count_ - << " kv_patch_scatter_ms=" << regione_profile_patch_scatter_ms_ - << " kv_patch_scatter_count=" - << regione_profile_patch_scatter_count_; -} - -void RegionECache::regione_profile_add_kv_store(double ms) { - if (!regione_profile_enabled()) return; - regione_profile_kv_store_cpu_ms_ += ms; - ++regione_profile_kv_store_count_; -} - -void RegionECache::regione_profile_add_prefetch_issue(double ms) { - if (!regione_profile_enabled()) return; - regione_profile_prefetch_issue_ms_ += ms; - ++regione_profile_prefetch_issue_count_; -} - -void RegionECache::regione_profile_add_prefetch_hit(double wait_ms) { - if (!regione_profile_enabled()) return; - regione_profile_prefetch_wait_ms_ += wait_ms; - ++regione_profile_prefetch_hit_count_; -} - -void RegionECache::regione_profile_add_prefetch_miss() { - if (!regione_profile_enabled()) return; - ++regione_profile_prefetch_miss_count_; -} - -void RegionECache::regione_profile_add_fallback_h2d(double ms) { - if (!regione_profile_enabled()) return; - regione_profile_fallback_h2d_ms_ += ms; - ++regione_profile_fallback_h2d_count_; -} - -void RegionECache::regione_profile_add_patch_scatter(double ms) { - if (!regione_profile_enabled()) return; - regione_profile_patch_scatter_ms_ += ms; - ++regione_profile_patch_scatter_count_; -} - } // namespace xllm diff --git a/xllm/core/framework/dit_cache/regione.h b/xllm/core/framework/dit_cache/regione.h index 3e013791a2..2fdbfe6895 100644 --- a/xllm/core/framework/dit_cache/regione.h +++ b/xllm/core/framework/dit_cache/regione.h @@ -135,24 +135,6 @@ class RegionECache { const torch::Tensor& image_rope, int64_t key_len) const; - bool regione_profile_enabled() const; - void regione_profile_reset_step(int64_t step, - bool partial_step, - bool full_step, - bool velocity_cache, - int64_t step_tokens, - int64_t full_tokens); - void regione_profile_log_step(double transformer_ms, - double arp_ms, - double scheduler_ms, - double total_ms) const; - void regione_profile_add_kv_store(double ms); - void regione_profile_add_prefetch_issue(double ms); - void regione_profile_add_prefetch_hit(double wait_ms); - void regione_profile_add_prefetch_miss(); - void regione_profile_add_fallback_h2d(double ms); - void regione_profile_add_patch_scatter(double ms); - private: void regione_select_regions(const torch::Tensor& sample, const torch::Tensor& model_output, @@ -236,23 +218,10 @@ class RegionECache { torch::Tensor regione_velocity_cache_; double regione_avd_accumulate_ = 1.0; double regione_avd_ratio_ = 1.0; - int64_t regione_profile_step_ = -1; - bool regione_profile_partial_step_ = false; - bool regione_profile_full_step_ = false; - bool regione_profile_velocity_cache_ = false; - int64_t regione_profile_step_tokens_ = 0; - int64_t regione_profile_full_tokens_ = 0; - int64_t regione_profile_kv_store_count_ = 0; - int64_t regione_profile_prefetch_issue_count_ = 0; - int64_t regione_profile_prefetch_hit_count_ = 0; - int64_t regione_profile_prefetch_miss_count_ = 0; - int64_t regione_profile_fallback_h2d_count_ = 0; - int64_t regione_profile_patch_scatter_count_ = 0; - double regione_profile_kv_store_cpu_ms_ = 0.0; - double regione_profile_prefetch_issue_ms_ = 0.0; - double regione_profile_prefetch_wait_ms_ = 0.0; - double regione_profile_fallback_h2d_ms_ = 0.0; - double regione_profile_patch_scatter_ms_ = 0.0; + double regione_avd_raw_gamma_ = 1.0; + double regione_avd_gamma_ = 1.0; + double regione_avd_gamma_exponent_ = 1.0; + double regione_avd_error_ = 0.0; std::vector regione_k_cache_cpu_; std::vector regione_v_cache_cpu_; std::vector regione_cond_k_cache_cpu_; diff --git a/xllm/core/runtime/dit_worker_impl.cpp b/xllm/core/runtime/dit_worker_impl.cpp index 98df223867..5a1603ec88 100644 --- a/xllm/core/runtime/dit_worker_impl.cpp +++ b/xllm/core/runtime/dit_worker_impl.cpp @@ -111,8 +111,6 @@ DiTCacheConfig parse_dit_cache_from_flags() { ::xllm::DiTConfig::get_instance().dit_regione_use_avd_gamma(); cache_config.regione.erosion_dilation = ::xllm::DiTConfig::get_instance().dit_regione_erosion_dilation(); - cache_config.regione.profile = - ::xllm::DiTConfig::get_instance().dit_regione_profile(); } else if (::xllm::DiTConfig::get_instance().dit_cache_policy() == "None") { cache_config.selected_policy = PolicyType::None; } diff --git a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h index 195f2623aa..dc8c23a4ef 100644 --- a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h +++ b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include #include @@ -1057,18 +1056,6 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { : image_rotary_emb_pos; auto* regione = DiTCache::get_instance().regione(); - const bool regione_profile_enabled = - regione != nullptr && regione->regione_profile_enabled(); - auto regione_profile_now = []() { - return std::chrono::steady_clock::now(); - }; - auto regione_profile_ms_since = - [](const std::chrono::steady_clock::time_point& start) { - return std::chrono::duration( - std::chrono::steady_clock::now() - start) - .count(); - }; - auto regione_dit_loop_start = regione_profile_now(); if (regione) { const auto sp_group = parallel_args_.dit_sp_group_; @@ -1085,10 +1072,6 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { for (int64_t i = 0; i < timesteps.size(0); ++i) { auto t = timesteps[i]; current_timestep_ = t; - auto regione_step_profile_start = regione_profile_now(); - double regione_transformer_ms = 0.0; - double regione_arp_ms = 0.0; - double regione_scheduler_ms = 0.0; double prev_t_value = 0.0; double t_value = 0.0; @@ -1121,24 +1104,9 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { auto timestep_expanded = t.expand({step_latents.size(0)}).to(step_latents.dtype()); - if (regione_profile_enabled) { - regione->regione_profile_reset_step( - i, - regione_plan.partial_step, - regione_plan.full_step, - regione_plan.use_velocity_cache, - step_latents.defined() && step_latents.dim() > 1 - ? step_latents.size(1) - : 0, - final_latents.defined() && final_latents.dim() > 1 - ? final_latents.size(1) - : 0); - } - torch::Tensor noise_pred; torch::Tensor neg_noise_pred; torch::Tensor pos_neg_noise_preds; - auto regione_transformer_start = regione_profile_now(); if (regione_plan.use_velocity_cache) { noise_pred = regione_input.cached_velocity; } else if (::xllm::ParallelConfig::get_instance().cfg_size() == 2 && @@ -1214,25 +1182,15 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { noise_pred = comb_pred * (cond_norm / noise_norm); } } - if (regione_profile_enabled) { - regione_transformer_ms = - regione_profile_ms_since(regione_transformer_start); - } - if (regione_plan.enabled) { - auto regione_arp_start = regione_profile_now(); regione->observe_velocity( final_latents, noise_pred, scheduler_->sigmas(), i, regione_plan); if (regione_plan.run_partition) { - if (regione_profile_enabled) { - regione_arp_ms = regione_profile_ms_since(regione_arp_start); - } regione_plan.direct_unedited = regione->regione_should_direct_unedited(i); } } - auto regione_scheduler_start = regione_profile_now(); auto latents_dtype = final_latents.dtype(); if (regione_plan.enabled) { if (regione_plan.partial_step) { @@ -1270,33 +1228,14 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { if (final_latents.dtype() != latents_dtype) { final_latents = final_latents.to(latents_dtype); } - if (regione_profile_enabled) { - regione_scheduler_ms = - regione_profile_ms_since(regione_scheduler_start); - regione->regione_profile_log_step( - regione_transformer_ms, - regione_arp_ms, - regione_scheduler_ms, - regione_profile_ms_since(regione_step_profile_start)); - } - } - if (regione_profile_enabled) { - LOG(INFO) << "[RegionEProfile] dit_loop_total_ms=" - << regione_profile_ms_since(regione_dit_loop_start); } current_timestep_ = torch::Tensor(); torch::Tensor output_image; - auto regione_vae_stage_start = regione_profile_now(); auto unpacked_latents = _unpack_latents(final_latents, height, width, vae_scale_factor_) .to(dtype_); - if (regione_profile_enabled) { - LOG(INFO) << "[RegionEProfile] vae_unpack_ms=" - << regione_profile_ms_since(regione_vae_stage_start); - } - regione_vae_stage_start = regione_profile_now(); auto latents_mean = torch::tensor(vae_model_args_.latents_mean(), torch::kDouble); latents_mean = @@ -1308,22 +1247,8 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { latents_std.view({1, latent_channels_, 1, 1, 1}).to(device_, dtype_); unpacked_latents = unpacked_latents / latents_std + latents_mean; - if (regione_profile_enabled) { - LOG(INFO) << "[RegionEProfile] vae_latent_norm_ms=" - << regione_profile_ms_since(regione_vae_stage_start); - } - regione_vae_stage_start = regione_profile_now(); output_image = vae_->decode(unpacked_latents).sample.squeeze(2); - if (regione_profile_enabled) { - LOG(INFO) << "[RegionEProfile] vae_decode_ms=" - << regione_profile_ms_since(regione_vae_stage_start); - } - regione_vae_stage_start = regione_profile_now(); output_image = vae_image_processor_->postprocess(output_image); - if (regione_profile_enabled) { - LOG(INFO) << "[RegionEProfile] vae_postprocess_ms=" - << regione_profile_ms_since(regione_vae_stage_start); - } auto output_chunks = torch::chunk(output_image, batch_size, /*dim=*/0); DiTForwardOutput out; out.tensors = std::move(output_chunks);