feat: add lag confidence to decouple adaptive speculative pruning from critical path.

xllm/core/common/global_flags.h

@@ -85,6 +85,8 @@ DECLARE_bool(enable_adaptive_speculative_decode);
 
 DECLARE_double(adaptive_speculative_min_gain);
 
+DECLARE_bool(enable_lag_confidence);
+
 DECLARE_int32(speculative_suffix_cache_max_depth);
 
 DECLARE_double(speculative_suffix_max_spec_factor);

xllm/core/common/options.cpp

@@ -52,6 +52,7 @@ std::string Options::to_string() const {
      << ", enable_adaptive_speculative_decode: "
      << enable_adaptive_speculative_decode()
      << ", adaptive_speculative_min_gain: " << adaptive_speculative_min_gain()
+     << ", enable_lag_confidence: " << enable_lag_confidence()
      << ", num_request_handling_threads: " << num_request_handling_threads()
      << ", communication_backend: " << communication_backend().value_or("null")
      << ", rank_tablefile: " << rank_tablefile().value_or("null")

xllm/core/common/options.h

@@ -99,6 +99,8 @@ class Options {
 
   PROPERTY(double, adaptive_speculative_min_gain) = 0.0;
 
+  PROPERTY(bool, enable_lag_confidence) = false;
+
   // thread num to handle requests
   PROPERTY(size_t, num_request_handling_threads) = 4;
 

xllm/core/distributed_runtime/master.cpp

@@ -541,6 +541,7 @@ Master::Master(const Options& options, EngineType type)
         .enable_adaptive_speculative_decode(
             options_.enable_adaptive_speculative_decode())
         .adaptive_speculative_min_gain(options_.adaptive_speculative_min_gain())
+        .enable_lag_confidence(options_.enable_lag_confidence())
         .task_type(options_.task_type())
         .enable_mla(options_.enable_mla())
         .npu_kernel_backend(options_.npu_kernel_backend())

xllm/core/framework/config/speculative_config.cpp

@@ -80,6 +80,12 @@ DEFINE_double(
     "Minimum relative throughput gain required to include a draft token in "
     "adaptive speculative validation.");
 
+DEFINE_bool(enable_lag_confidence,
+            false,
+            "Whether the adaptive controller prunes using the previous decode "
+            "step's confidence (lag-1) so the decision overlaps this step's "
+            "draft forward. Requires enable_adaptive_speculative_decode.");
+
 namespace xllm {
 
 void SpeculativeConfig::from_flags() {
@@ -97,6 +103,7 @@ void SpeculativeConfig::from_flags() {
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_atb_spec_kernel);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_adaptive_speculative_decode);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(adaptive_speculative_min_gain);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_lag_confidence);
 }
 
 void SpeculativeConfig::from_json(const JsonReader& json) {
@@ -147,6 +154,8 @@ void SpeculativeConfig::append_config_json(
       config_json, default_config, enable_adaptive_speculative_decode);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, adaptive_speculative_min_gain);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, enable_lag_confidence);
 }
 
 SpeculativeConfig& SpeculativeConfig::get_instance() {

xllm/core/framework/config/speculative_config.h

@@ -84,7 +84,8 @@ class SpeculativeConfig final {
          "enable_mtp_draft_body_tp1",
          "enable_atb_spec_kernel",
          "enable_adaptive_speculative_decode",
-         "adaptive_speculative_min_gain"}};
+         "adaptive_speculative_min_gain",
+         "enable_lag_confidence"}};
     return kOptionCategory;
   }
 
@@ -116,6 +117,13 @@ class SpeculativeConfig final {
 
   PROPERTY(double, adaptive_speculative_min_gain) = 0.0;
 
+  // When true, the adaptive controller prunes using the PREVIOUS decode step's
+  // confidence (lag-1) instead of this step's, so the decision no longer
+  // data-depends on this step's draft forward and can overlap it. Requires
+  // enable_adaptive_speculative_decode. Off by default keeps the this-step
+  // path.
+  PROPERTY(bool, enable_lag_confidence) = false;
+
  private:
   // ASCII case-insensitive equality. Mirrors the manual case handling in
   // is_mtp_algorithm rather than pulling <algorithm>/<cctype> into this header.

xllm/core/framework/speculative/embedding_cache.cpp

@@ -121,7 +121,8 @@ void EmbeddingCache::write_target_context(
     const std::vector<std::string>& request_ids,
     const torch::Tensor& accepted_tokens,
     const torch::Tensor& accepted_embeddings,
-    int32_t num_speculative_tokens) {
+    int32_t num_speculative_tokens,
+    const torch::Tensor& confidence) {
   CHECK(accepted_tokens.defined()) << "accepted target tokens are undefined";
   CHECK(accepted_embeddings.defined())
       << "accepted target embeddings are undefined";
@@ -139,6 +140,21 @@ void EmbeddingCache::write_target_context(
       << "accepted token/embedding width mismatch";
   CHECK_GE(num_speculative_tokens, 0) << "invalid speculative token count";
 
+  // Lag confidence: one blocking D2H of [batch, num_speculative_tokens] here,
+  // off the critical path (the caller invokes this past the validate sync), so
+  // the next-step read is a pure-host slot copy. Undefined when lag is off.
+  torch::Tensor confidence_cpu;
+  if (confidence.defined()) {
+    CHECK_EQ(confidence.dim(), 2)
+        << "lag confidence should be [batch, num_speculative_tokens]";
+    CHECK_EQ(confidence.size(0), static_cast<int64_t>(ids.size()))
+        << "lag confidence batch mismatch";
+    CHECK_EQ(confidence.size(1), num_speculative_tokens)
+        << "lag confidence width mismatch";
+    confidence_cpu = safe_to(confidence, torch::kCPU).to(torch::kFloat32);
+    confidence_cpu = confidence_cpu.contiguous();
+  }
+
   torch::Tensor accepted_tokens_cpu = to_cpu_int64_contiguous(accepted_tokens);
   const int64_t* accepted_tokens_data =
       accepted_tokens_cpu.const_data_ptr<int64_t>();
@@ -176,6 +192,9 @@ void EmbeddingCache::write_target_context(
     state.position_offset = last_idx;
     state.correction_token_id = correction_token;
     state.correction_position_offset = correction_offset;
+    if (confidence_cpu.defined()) {
+      state.confidence = confidence_cpu.select(/*dim=*/0, i).detach().clone();
+    }
     state.embedding = accepted_embeddings.select(/*dim=*/0, i)
                           .select(/*dim=*/0, last_idx)
                           .detach()
@@ -265,6 +284,41 @@ std::vector<int32_t> EmbeddingCache::read_accepted_prefix_lengths(
   return accepted_prefix_lengths;
 }
 
+EmbeddingCache::LaggedConfidence EmbeddingCache::read_lagged_confidence(
+    const std::vector<int32_t>& ids,
+    const std::vector<std::string>& request_ids,
+    int32_t num_speculative_tokens) const {
+  CHECK(!ids.empty()) << "decode ids should not be empty";
+  CHECK(request_ids.empty() || request_ids.size() == ids.size())
+      << "embedding_id / request_id count mismatch";
+  CHECK_GT(num_speculative_tokens, 0)
+      << "lag confidence requires positive speculative tokens";
+  const int32_t num_ids = static_cast<int32_t>(ids.size());
+  LaggedConfidence result;
+  result.valid.assign(static_cast<size_t>(num_ids), false);
+  result.confidence = torch::zeros({num_ids, num_speculative_tokens},
+                                   torch::dtype(torch::kFloat32));
+  for (int32_t i = 0; i < num_ids; ++i) {
+    const DecodeState& state = get_tail(ids[i]);
+    // Freshness gate mirrors read_accepted_prefix_lengths: a slot never
+    // written, or recycled by a later request (request_id mismatch), carries no
+    // usable lagged confidence. Width mismatch is likewise rejected. Invalid
+    // rows stay zero-filled and valid[i]=false so the caller falls back to full
+    // width.
+    if (!state.valid ||
+        (!request_ids.empty() && state.request_id != request_ids[i])) {
+      continue;
+    }
+    if (!state.confidence.defined() ||
+        state.confidence.numel() != num_speculative_tokens) {
+      continue;
+    }
+    result.confidence.select(/*dim=*/0, i).copy_(state.confidence);
+    result.valid[static_cast<size_t>(i)] = true;
+  }
+  return result;
+}
+
 void EmbeddingCache::clear(const std::vector<int32_t>& ids) {
   for (int32_t id : ids) {
     DecodeState& tail = mutable_tail(id);

xllm/core/framework/speculative/embedding_cache.h

@@ -55,6 +55,13 @@ class EmbeddingCache final {
 
     int32_t correction_token_id = 0;  // accepted token for step correction
     int32_t correction_position_offset = 0;
+
+    // Previous decode step's per-draft confidence [num_speculative_tokens],
+    // fp32 on CPU. Consumed one step later by the adaptive controller under lag
+    // confidence (enable_lag_confidence) so the pruning decision does not
+    // data-depend on this step's draft forward. Undefined when lag confidence
+    // is off or no confidence has been written for this slot yet.
+    torch::Tensor confidence;
   };
 
   EmbeddingCache(int32_t total_nums);
@@ -82,11 +89,15 @@ class EmbeddingCache final {
   // Writes target validate output after rejection sampling. accepted_tokens is
   // a contiguous accepted prefix padded by -1; accepted_embeddings keeps the
   // corresponding target hidden states for the next draft extend input.
+  // confidence, when defined, is this step's per-draft confidence
+  // [batch, num_speculative_tokens] stored per slot for lag-confidence pruning
+  // one step later; pass an undefined tensor when lag confidence is off.
   void write_target_context(const std::vector<int32_t>& embedding_ids,
                             const std::vector<std::string>& request_ids,
                             const torch::Tensor& accepted_tokens,
                             const torch::Tensor& accepted_embeddings,
-                            int32_t num_speculative_tokens);
+                            int32_t num_speculative_tokens,
+                            const torch::Tensor& confidence = torch::Tensor());
 
   // Algorithm-specific placeholder embedding for missing target context, e.g.
   // PD first decode. MTP uses hidden_size; Eagle3 uses 3 * hidden_size.
@@ -102,6 +113,20 @@ class EmbeddingCache final {
       const std::vector<int32_t>& embedding_ids,
       const std::vector<std::string>& request_ids) const;
 
+  // Lagged (previous-step) confidence per request, for adaptive lag-confidence
+  // pruning. confidence is [batch, num_speculative_tokens] fp32 on CPU (rows
+  // for invalid slots are zero-filled); valid[i] is false when slot i carries
+  // no usable lagged confidence (never written, request_id mismatch on reuse,
+  // or undefined) so the caller must fall back to full-width (no prune) for it.
+  struct LaggedConfidence {
+    torch::Tensor confidence;
+    std::vector<bool> valid;
+  };
+  LaggedConfidence read_lagged_confidence(
+      const std::vector<int32_t>& embedding_ids,
+      const std::vector<std::string>& request_ids,
+      int32_t num_speculative_tokens) const;
+
   void clear(const std::vector<int32_t>& embedding_ids);
 
  private: