diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index d9071c6e69..25edaf4303 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -360,6 +360,12 @@ 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_double(dit_regione_cache_threshold); +DECLARE_bool(dit_regione_use_avd_gamma); +DECLARE_bool(dit_regione_erosion_dilation); + 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..9f8ce4768d 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,31 @@ 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_double(dit_regione_cache_threshold, + 0.02, + "RegionE: AVDCache error threshold δ (paper Eq.8). " + "Reuse velocity while 1-accumulate <= threshold. " + "Default is 0.02."); + +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_sp_communication_overlap, true, "Communication & Computation overlap for sequence parallel"); @@ -131,6 +156,11 @@ 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_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_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_laser_attention_enabled); @@ -157,6 +187,11 @@ 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_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_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_laser_attention_enabled); @@ -194,6 +229,16 @@ 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_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_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..8693b1cb42 100644 --- a/xllm/core/framework/config/dit_config.h +++ b/xllm/core/framework/config/dit_config.h @@ -51,6 +51,11 @@ class DiTConfig final { "dit_cache_end_steps", "dit_cache_start_blocks", "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_sp_communication_overlap", "dit_debug_print", "dit_laser_attention_enabled", @@ -87,6 +92,16 @@ 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(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_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..851bf3785f 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,21 @@ struct FBCacheTaylorSeerOptions : public DiTBaseCacheOptions { int n_derivatives = 3; }; +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. + 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; +}; + struct ResidualCacheOptions { // The number of steps to skip at the start. int64_t dit_cache_start_steps = 5; @@ -85,6 +105,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..86da2864d6 --- /dev/null +++ b/xllm/core/framework/dit_cache/regione.cpp @@ -0,0 +1,825 @@ +/* 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 + +#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_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(); +} + +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, + double timestep, + 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. + // 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; + 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 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, + 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; + 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), timestep-delta normalized γ: + // ratio = gamma(step) * (1 + (t - t_prev) / 1000) + // accumulate *= ratio; error = 1 - accumulate + // 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; + regione_avd_raw_gamma_ = raw_gamma; + regione_avd_gamma_ = gamma; + regione_avd_gamma_exponent_ = gamma_exponent; + + regione_avd_accumulate_ *= ratio; + 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 + } + return false; // reuse velocity cache * ratio +} + +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_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(); + 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 { + // 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 { + // 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_; +} + +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 { + // 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); +} + +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) { + 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()); + // 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); + } + 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 (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_}); + 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, + double timestep, + double prev_timestep) { + 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, timestep, prev_timestep) && + 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) { + // 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); + } + } + 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)) { + 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}; +} + +} // 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..2fdbfe6895 --- /dev/null +++ b/xllm/core/framework/dit_cache/regione.h @@ -0,0 +1,233 @@ +/* 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 + +#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; + // 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; + 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, + double timestep = 0.0, + double prev_timestep = 0.0); + 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_; + double regione_avd_accumulate_ = 1.0; + double regione_avd_ratio_ = 1.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_; + 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..5a1603ec88 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,25 @@ 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(); + 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(); } 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..dc8c23a4ef 100644 --- a/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h +++ b/xllm/models/dit/pipelines/pipeline_qwenimage_edit_plus.h @@ -1055,35 +1055,74 @@ 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()) { + 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, t_value, prev_t_value); + 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 +1132,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 +1156,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); @@ -1143,9 +1182,49 @@ class QwenImageEditPlusPipelineImpl : public torch::nn::Module { noise_pred = comb_pred * (cond_norm / noise_norm); } } + 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..d70ed3b776 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; } @@ -1584,6 +1585,22 @@ class QwenDoubleStreamAttnProcessor2_0Impl : public torch::nn::Module { 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) = + 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); auto txt_key = attn_->add_k_proj_->forward(encoder_hidden_states); @@ -1625,15 +1642,14 @@ 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); - - 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,10 +2368,15 @@ class QwenImageTransformer2DModelImpl : public torch::nn::Module { block_attention_kwargs["attention_mask"] = joint_attention_mask; } + auto* regione = DiTCache::get_instance().regione(); + // 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) { - 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, @@ -2373,6 +2394,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 +2415,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) =