feat: support parse dit cache config from runtime flags. (jd-opensource#473)

Author: z-jun03

xllm/core/common/global_flags.cpp

@@ -436,4 +436,25 @@ DEFINE_bool(
     enable_dp_balance,
     false,
     "Whether to enable dp load balance, if true, sequences within a single "
-    "dp batch will be shuffled.");
+    "dp batch will be shuffled.");
+
+// --- dit cache config ---
+
+DEFINE_string(dit_cache_policy,
+              "TaylorSeer",
+              "The policy of dit cache(e.g. None, FBCache, TaylorSeer, "
+              "FBCacheTaylorSeer).");
+
+DEFINE_int64(dit_cache_warmup_steps, 0, "The number of warmup steps.");
+
+DEFINE_int64(dit_cache_n_derivatives,
+             3,
+             "The number of derivatives to use in TaylorSeer.");
+
+DEFINE_int64(dit_cache_skip_interval_steps,
+             3,
+             "The interval steps to skip for derivative calculation.");
+
+DEFINE_double(dit_cache_residual_diff_threshold,
+              0.09f,
+              "The residual difference threshold for cache reuse.");

xllm/core/common/global_flags.h

@@ -214,3 +214,13 @@ DECLARE_bool(enable_prefetch_weight);
 DECLARE_int32(flashinfer_workspace_buffer_size);
 
 DECLARE_bool(enable_dp_balance);
+
+DECLARE_string(dit_cache_policy);
+
+DECLARE_int64(dit_cache_warmup_steps);
+
+DECLARE_int64(dit_cache_n_derivatives);
+
+DECLARE_int64(dit_cache_skip_interval_steps);
+
+DECLARE_double(dit_cache_residual_diff_threshold);

xllm/core/framework/dit_cache/dit_cache_config.h

@@ -25,9 +25,6 @@ enum class PolicyType {
 };
 
 struct DiTBaseCacheOptions {
-  // the number of inference steps.
-  int num_inference_steps = 25;
-
   // the number of warmup steps.
   int warmup_steps = 0;
 };

xllm/core/framework/dit_cache/fbcache.cpp

@@ -20,14 +20,9 @@ namespace xllm {
 void FBCache::init(const DiTCacheConfig& cfg) {
   CHECK_GE(cfg.fbcache.residual_diff_threshold, 0.0)
       << "residual_diff_threshold must be >= 0";
-  CHECK_GT(cfg.fbcache.num_inference_steps, 0)
-      << "num_inference_steps must be > 0";
   CHECK_GE(cfg.fbcache.warmup_steps, 0) << "warmup_steps must be >= 0";
-  CHECK_LE(cfg.fbcache.warmup_steps, cfg.fbcache.num_inference_steps)
-      << "warmup_steps cannot exceed num_inference_steps";
 
   residual_diff_threshold_ = cfg.fbcache.residual_diff_threshold;
-  num_inference_steps_ = cfg.fbcache.num_inference_steps;
   warmup_steps_ = cfg.fbcache.warmup_steps;
   current_step_ = 0;
   use_cache_ = false;

xllm/core/framework/dit_cache/fbcache_taylorseer.cpp

@@ -20,18 +20,12 @@ namespace xllm {
 void FBCacheTaylorSeer::init(const DiTCacheConfig& cfg) {
   CHECK_GE(cfg.fbcachetaylorseer.residual_diff_threshold, 0.0)
       << "residual_diff_threshold must be >= 0";
-  CHECK_GT(cfg.fbcachetaylorseer.num_inference_steps, 0)
-      << "num_inference_steps must be > 0";
   CHECK_GE(cfg.fbcachetaylorseer.warmup_steps, 0)
       << "warmup_steps must be >= 0";
   CHECK_GE(cfg.fbcachetaylorseer.n_derivatives, 0)
       << "n_derivatives must be >= 0";
-  CHECK_LE(cfg.fbcachetaylorseer.warmup_steps,
-           cfg.fbcachetaylorseer.num_inference_steps)
-      << "warmup_steps cannot exceed num_inference_steps";
 
   residual_diff_threshold_ = cfg.fbcachetaylorseer.residual_diff_threshold;
-  num_inference_steps_ = cfg.fbcachetaylorseer.num_inference_steps;
   warmup_steps_ = cfg.fbcachetaylorseer.warmup_steps;
 
   if (!taylorseer) {
@@ -40,8 +34,6 @@ void FBCacheTaylorSeer::init(const DiTCacheConfig& cfg) {
 
   DiTCacheConfig ts_cfg;
   ts_cfg.taylorseer.n_derivatives = cfg.fbcachetaylorseer.n_derivatives;
-  ts_cfg.taylorseer.num_inference_steps =
-      cfg.fbcachetaylorseer.num_inference_steps;
   taylorseer->init(ts_cfg);
 }
 

xllm/core/framework/dit_cache/taylorseer.cpp

@@ -22,20 +22,11 @@ double factorial(int k) { return std::tgamma(static_cast<double>(k) + 1.0); }
 }  // namespace
 
 void TaylorSeer::init(const DiTCacheConfig& cfg) {
-  CHECK_GT(cfg.taylorseer.num_inference_steps, 0)
-      << "num_inference_steps must be > 0";
   CHECK_GE(cfg.taylorseer.warmup_steps, 0) << "warmup_steps must be >= 0";
   CHECK_GE(cfg.taylorseer.skip_interval_steps, 0)
       << "skip_interval_steps must be >= 0";
   CHECK_GE(cfg.taylorseer.n_derivatives, 0) << "n_derivatives must be >= 0";
 
-  CHECK_LT(cfg.taylorseer.skip_interval_steps,
-           cfg.taylorseer.num_inference_steps)
-      << "skip_interval_steps must be less than num_inference_steps";
-  CHECK_LT(cfg.taylorseer.warmup_steps, cfg.taylorseer.num_inference_steps)
-      << "warmup_steps must be less than num_inference_steps";
-
-  num_inference_steps_ = cfg.taylorseer.num_inference_steps;
   warmup_steps_ = cfg.taylorseer.warmup_steps;
   skip_interval_steps_ = cfg.taylorseer.skip_interval_steps;
   n_derivatives_ = cfg.taylorseer.n_derivatives;

xllm/core/runtime/dit_worker.cpp

@@ -38,6 +38,35 @@ limitations under the License.
 #include "util/utils.h"
 
 namespace xllm {
+
+namespace {
+DiTCacheConfig parse_dit_cache_from_flags() {
+  DiTCacheConfig cache_config;
+  if (FLAGS_dit_cache_policy == "FBCache") {
+    cache_config.selected_policy = PolicyType::FBCache;
+    cache_config.fbcache.warmup_steps = FLAGS_dit_cache_warmup_steps;
+    cache_config.fbcache.residual_diff_threshold =
+        FLAGS_dit_cache_residual_diff_threshold;
+  } else if (FLAGS_dit_cache_policy == "TaylorSeer") {
+    cache_config.selected_policy = PolicyType::TaylorSeer;
+    cache_config.taylorseer.n_derivatives = FLAGS_dit_cache_n_derivatives;
+    cache_config.taylorseer.skip_interval_steps =
+        FLAGS_dit_cache_skip_interval_steps;
+    cache_config.taylorseer.warmup_steps = FLAGS_dit_cache_warmup_steps;
+  } else if (FLAGS_dit_cache_policy == "FBCacheTaylorSeer") {
+    cache_config.selected_policy = PolicyType::FBCacheTaylorSeer;
+    cache_config.fbcachetaylorseer.n_derivatives =
+        FLAGS_dit_cache_n_derivatives;
+    cache_config.fbcachetaylorseer.warmup_steps = FLAGS_dit_cache_warmup_steps;
+    cache_config.fbcachetaylorseer.residual_diff_threshold =
+        FLAGS_dit_cache_residual_diff_threshold;
+  } else if (FLAGS_dit_cache_policy == "None") {
+    cache_config.selected_policy = PolicyType::TaylorSeer;
+  }
+  return cache_config;
+}
+}  // namespace
+
 DiTWorker::DiTWorker(const ParallelArgs& parallel_args,
                      const torch::Device& device,
                      const runtime::Options& options)
@@ -65,17 +94,8 @@ bool DiTWorker::init_model(const std::string& model_weights_path) {
   dit_model_executor_ =
       std::make_unique<DiTExecutor>(dit_model_.get(), options_);
 
-  DiTCacheConfig cache_config_;
-
-  // TODO: Optimize ditcache configuration initialization.
-  cache_config_.selected_policy = PolicyType::TaylorSeer;
-  cache_config_.taylorseer.n_derivatives = 3;
-  cache_config_.taylorseer.skip_interval_steps = 3;
-  cache_config_.taylorseer.num_inference_steps = 25;
-  cache_config_.taylorseer.warmup_steps = 0;
-
-  bool success = DiTCache::get_instance().init(cache_config_);
-  CHECK(success) << "DiTCache init failed";
+  DiTCacheConfig cache_config = parse_dit_cache_from_flags();
+  DiTCache::get_instance().init(cache_config);
 
   return true;
 }