diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index d92d7a9298..5646bcfade 100644 --- a/xllm/core/common/global_flags.h +++ b/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); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index 66ea3e8b5c..11786cda05 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -53,6 +53,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") diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index eb374a1c0d..e4090f0292 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -101,6 +101,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; diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index aa2dc8d74d..658b079e47 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -578,6 +578,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()) diff --git a/xllm/core/framework/config/speculative_config.cpp b/xllm/core/framework/config/speculative_config.cpp index d1c203e91c..66bdb12762 100644 --- a/xllm/core/framework/config/speculative_config.cpp +++ b/xllm/core/framework/config/speculative_config.cpp @@ -78,6 +78,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() { @@ -95,6 +101,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) { @@ -145,6 +152,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() { diff --git a/xllm/core/framework/config/speculative_config.h b/xllm/core/framework/config/speculative_config.h index 38108f07e3..801ffd87de 100644 --- a/xllm/core/framework/config/speculative_config.h +++ b/xllm/core/framework/config/speculative_config.h @@ -94,7 +94,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; } @@ -125,6 +126,13 @@ class SpeculativeConfig final { PROPERTY(bool, enable_adaptive_speculative_decode) = false; 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; }; } // namespace xllm diff --git a/xllm/core/framework/speculative/embedding_cache.cpp b/xllm/core/framework/speculative/embedding_cache.cpp index ec389a7f9a..544949cbf4 100644 --- a/xllm/core/framework/speculative/embedding_cache.cpp +++ b/xllm/core/framework/speculative/embedding_cache.cpp @@ -121,7 +121,8 @@ void EmbeddingCache::write_target_context( const std::vector& 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(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(); @@ -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 EmbeddingCache::read_accepted_prefix_lengths( return accepted_prefix_lengths; } +EmbeddingCache::LaggedConfidence EmbeddingCache::read_lagged_confidence( + const std::vector& ids, + const std::vector& 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(ids.size()); + LaggedConfidence result; + result.valid.assign(static_cast(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(i)] = true; + } + return result; +} + void EmbeddingCache::clear(const std::vector& ids) { for (int32_t id : ids) { DecodeState& tail = mutable_tail(id); diff --git a/xllm/core/framework/speculative/embedding_cache.h b/xllm/core/framework/speculative/embedding_cache.h index 2bcad8d246..ed26dd6e78 100644 --- a/xllm/core/framework/speculative/embedding_cache.h +++ b/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& embedding_ids, const std::vector& 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& embedding_ids, const std::vector& 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 valid; + }; + LaggedConfidence read_lagged_confidence( + const std::vector& embedding_ids, + const std::vector& request_ids, + int32_t num_speculative_tokens) const; + void clear(const std::vector& embedding_ids); private: diff --git a/xllm/core/runtime/dflash_worker_impl.cpp b/xllm/core/runtime/dflash_worker_impl.cpp index 1a52f721c1..eaf1c43df5 100644 --- a/xllm/core/runtime/dflash_worker_impl.cpp +++ b/xllm/core/runtime/dflash_worker_impl.cpp @@ -716,13 +716,50 @@ std::optional DFlashWorkerImpl::step_decode( << "DFlash decode target state count mismatch"; update_decode_step_input(input, last_states); - DraftBlock draft_block = run_decode_draft(input, validate_input); + // Lag confidence: decide the prune from the PREVIOUS step's confidence BEFORE + // launching this step's draft. The decision is pure host (the lagged + // confidence D2H completed last step) and has no data dependency on this + // step's draft, so computing it here lets run_decode_draft build the pruned + // varlen validate batch inside the draft-overlap window instead of rebuilding + // it on the critical path in run_validate. Empty vector when the flag is off. + std::vector lagged_prefix_lengths; + if (options_.enable_lag_confidence()) { + lagged_prefix_lengths = decide_lagged_prefix_lengths(input); + } + DraftBlock draft_block = + run_decode_draft(input, validate_input, lagged_prefix_lengths); return run_validate(input, draft_block, validate_input); } +void DFlashWorkerImpl::prepare_overlap_validate_input( + const ForwardInput& input, + ForwardInput& validate_input, + const std::vector& lagged_prefix_lengths, + DraftBlock& draft_block) { + const int32_t batch_size = input.input_params.meta.num_sequences; + std::vector per_seq_val_tokens; + int32_t max_val_tokens = 0; + const bool did_prune = prefix_lengths_to_val_tokens( + lagged_prefix_lengths, batch_size, &per_seq_val_tokens, &max_val_tokens); + if (did_prune) { + // Build the pruned varlen batch here, overlapping the in-flight draft. + // run_validate consumes this directly and skips its own rebuild. + prepare_validate_inputs(input, validate_input, per_seq_val_tokens); + draft_block.varlen_prebuilt = true; + draft_block.per_seq_val_tokens = std::move(per_seq_val_tokens); + draft_block.max_val_tokens = max_val_tokens; + } else { + // No prune (flag off, first step, or every seq full width): dense batch, + // exactly as before. run_validate takes the legacy path. + prepare_validate_inputs(input, validate_input); + draft_block.varlen_prebuilt = false; + } +} + DFlashWorkerImpl::DraftBlock DFlashWorkerImpl::run_decode_draft( const ForwardInput& input, - ForwardInput& validate_input) { + ForwardInput& validate_input, + const std::vector& lagged_prefix_lengths) { Timer timer; ForwardInput query_input; @@ -736,9 +773,16 @@ DFlashWorkerImpl::DraftBlock DFlashWorkerImpl::run_decode_draft( // launch above returns immediately, so building validate_input here (it only // reads the original input; draft tokens are injected later in // fill_validate_input_from_draft_outputs) runs on the host while the draft - // computes on device, instead of delaying the draft launch. - prepare_validate_inputs(input, validate_input); - // Keep draft tokens consistent across tensor-parallel ranks. + // computes on device, instead of delaying the draft launch. Under lag + // confidence the lagged decision lets us build the pruned varlen batch here, + // moving the whole rebuild off run_validate's critical path. + DraftBlock draft_block; + prepare_overlap_validate_input( + input, validate_input, lagged_prefix_lengths, draft_block); + // Unify the draft next_tokens across the tensor-parallel group before + // process_draft_sample_output() compresses the probs into the cache, so every + // rank caches the same selected draft prob under schedule-overlap. No-op for + // a single rank. maybe_broadcast_spec_tokens(draft_output.sample_output.next_tokens); process_draft_sample_output(draft_output.sample_output); COUNTER_ADD(speculative_execution_latency_seconds_draft, @@ -752,7 +796,6 @@ DFlashWorkerImpl::DraftBlock DFlashWorkerImpl::run_decode_draft( CHECK_EQ(num_draft_tokens % num_speculative_tokens, 0) << "DFlash draft token count mismatch."; const int32_t batch_size = num_draft_tokens / num_speculative_tokens; - DraftBlock draft_block; draft_block.proposal = DraftProposal(draft_output.sample_output.next_tokens.view( {batch_size, num_speculative_tokens})); @@ -885,6 +928,48 @@ void DFlashWorkerImpl::fill_validate_input_from_draft_outputs_varlen( record_metadata_ready_event(compute_stream, validate_input); } +bool DFlashWorkerImpl::prefix_lengths_to_val_tokens( + const std::vector& prefix_lengths, + int32_t batch_size, + std::vector* per_seq_val_tokens, + int32_t* max_val_tokens) const { + const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + const int32_t default_val_tokens = num_speculative_tokens + 1; + per_seq_val_tokens->clear(); + if (prefix_lengths.empty()) { + *max_val_tokens = default_val_tokens; + return false; + } + // Note: we intentionally do NOT mask draft probs beyond each seq's + // prefix_len. The varlen validate path only sends prefix_lengths[i] draft + // tokens per seq to target; and apply_pruned_prefix_lengths downstream + // overwrites all pruned rejection-sampler outputs (via cut_mask + drop mask). + // So the sampler's decision on pruned draft slots is irrelevant to the + // emitted tokens. + per_seq_val_tokens->resize(batch_size); + bool did_prune = false; + int32_t max_width = 0; + for (int32_t i = 0; i < batch_size; ++i) { + const int32_t p = std::clamp(prefix_lengths[static_cast(i)], + /*min=*/0, + /*max=*/num_speculative_tokens); + // Per-seq validate width = accepted-draft-count + 1 bonus. When the + // controller decides prefix=0 (don't speculate this step), the seq still + // must verify its bonus token, so the minimum is 1 slot — not 2. A previous + // floor to 2 forced a phantom "draft slot" at position 0 that leaked + // whatever the sampler emitted there past the intended prefix, showing up + // as duplicate/garbage tokens in adaptive output. + const int32_t width = p + 1; + (*per_seq_val_tokens)[static_cast(i)] = width; + max_width = std::max(max_width, width); + if (width < default_val_tokens) { + did_prune = true; + } + } + *max_val_tokens = max_width; + return did_prune; +} + std::optional DFlashWorkerImpl::run_validate( const ForwardInput& input, const DraftBlock& draft_block_in, @@ -908,36 +993,33 @@ std::optional DFlashWorkerImpl::run_validate( const int32_t batch_size = input.input_params.meta.num_sequences; DraftBlock draft_block = draft_block_in; - std::vector prefix_lengths = - compute_adaptive_prefix_lengths(draft_block, input); std::vector per_seq_val_tokens; bool did_prune = false; int32_t max_val_tokens = default_val_tokens; - if (!prefix_lengths.empty()) { - // Pruned outputs are overwritten later, so draft_probs stay untouched. - per_seq_val_tokens.resize(batch_size); - max_val_tokens = 0; - for (int32_t i = 0; i < batch_size; ++i) { - int32_t p = std::clamp(prefix_lengths[static_cast(i)], - /*min=*/0, - /*max=*/num_speculative_tokens); - // Per-seq validate width = accepted-draft-count + 1 bonus. When the - // controller decides prefix=0 (don't speculate this step), the seq - // still must verify its bonus token, so the minimum is 1 slot — not - // 2. A previous floor to 2 forced a phantom "draft slot" at position - // 0 that leaked whatever the sampler emitted there past the intended - // prefix, showing up as duplicate/garbage tokens in adaptive output. - const int32_t width = p + 1; - per_seq_val_tokens[static_cast(i)] = width; - max_val_tokens = std::max(max_val_tokens, width); - if (width < default_val_tokens) { - did_prune = true; - } + if (draft_block.varlen_prebuilt) { + // Lag confidence: run_decode_draft already decided the prune and built the + // pruned varlen validate_input in the draft-overlap window. Consume the + // decision directly; only the draft-token fill (which needs this step's + // draft output) remains on the critical path below. + per_seq_val_tokens = draft_block.per_seq_val_tokens; + did_prune = !per_seq_val_tokens.empty(); + max_val_tokens = + did_prune ? draft_block.max_val_tokens : default_val_tokens; + } else { + // Legacy this-step path: decide now (blocking on this step's draft output) + // and rebuild the varlen batch here, on the critical path. + std::vector prefix_lengths = + options_.enable_lag_confidence() + ? draft_block.lagged_prefix_lengths + : compute_adaptive_prefix_lengths(draft_block, input); + did_prune = prefix_lengths_to_val_tokens( + prefix_lengths, batch_size, &per_seq_val_tokens, &max_val_tokens); + if (did_prune) { + prepare_validate_inputs(input, validate_input, per_seq_val_tokens); } } if (did_prune) { - apply_per_seq_varlen_prune(input, validate_input, per_seq_val_tokens); fill_validate_input_from_draft_outputs_varlen( draft_block, validate_input, *compute_stream_, per_seq_val_tokens); } else { @@ -991,16 +1073,16 @@ std::optional DFlashWorkerImpl::run_validate( // dense rejection sampling. if (did_prune) { // effective_prefix[i] mirrors the controller's decision (0-based, clamped - // to [0, num_speculative_tokens]). Since per_seq_val_tokens[i] is exactly - // prefix_lengths[i] + 1 now (bonus slot only when prefix=0), we could - // equivalently write per_seq_val_tokens[i] - 1; keeping the raw - // prefix_lengths read here documents the semantic and stays robust if the - // width calculation grows another guard later. + // to [0, num_speculative_tokens]). per_seq_val_tokens[i] is exactly that + // clamped prefix + 1 (bonus slot only when prefix=0), so recover it as + // per_seq_val_tokens[i] - 1. Deriving from per_seq_val_tokens (rather than + // the raw prefix_lengths) keeps this correct on both the lag-prebuilt path + // — where the decision was made in run_decode_draft and only + // per_seq_val_tokens survives — and the legacy this-step path. std::vector effective_prefix(batch_size); for (int32_t i = 0; i < batch_size; ++i) { - int32_t p = prefix_lengths[static_cast(i)]; effective_prefix[static_cast(i)] = - std::clamp(p, 0, options_.num_speculative_tokens()); + per_seq_val_tokens[static_cast(i)] - 1; } adaptive_pruning::PrunedPrefixMasks masks = adaptive_pruning::build_pruned_prefix_masks( @@ -1035,7 +1117,27 @@ std::optional DFlashWorkerImpl::run_validate( // as rejections. Zero extra device sync — we're already on CPU. record_validate_metrics( val_output, did_prune ? per_seq_val_tokens : std::vector{}); - write_target_context_to_cache(input, val_output); + // Under lag confidence, stash THIS step's confidence into the cache so the + // next step's decision can consume it (lag-1). The store D2H runs here, past + // the validate sync above, off the critical path. Source mirrors the + // controller's: ConfidenceHead output when present, else proposal probs. Only + // a [batch, num_speculative_tokens] tensor is stored — DFlash's N+1-wide + // proposal probs have no confidence head and the wrong width, so they are + // dropped (undefined) and that step's slots simply carry no lagged + // confidence (read falls back to full width), never tripping the write's + // width CHECK. + torch::Tensor confidence_to_store; + if (options_.enable_lag_confidence()) { + const torch::Tensor confidence_src = + draft_block.confidence_probs.defined() + ? draft_block.confidence_probs + : draft_block.proposal.draft_probs().value_or(torch::Tensor()); + if (confidence_src.defined() && confidence_src.dim() == 2 && + confidence_src.size(1) == options_.num_speculative_tokens()) { + confidence_to_store = confidence_src; + } + } + write_target_context_to_cache(input, val_output, confidence_to_store); if (!enable_schedule_overlap() && !driver_ && !dp_driver_) { return std::nullopt; @@ -1199,13 +1301,22 @@ void DFlashWorkerImpl::update_decode_step_input( input.device_tensors_ready = false; } -void DFlashWorkerImpl::prepare_validate_inputs(const ForwardInput& input, - ForwardInput& validate_input) { +void DFlashWorkerImpl::prepare_validate_inputs( + const ForwardInput& input, + ForwardInput& validate_input, + const std::vector& per_seq_val_tokens) { c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); ForwardInput prepared_input = input; prepared_input.metadata_ready_event.reset(); - SpeculativeWorkerImpl::prepare_validate_inputs(prepared_input, - validate_input); + if (per_seq_val_tokens.empty()) { + SpeculativeWorkerImpl::prepare_validate_inputs(prepared_input, + validate_input); + } else { + // Pruned varlen batch: Σ per_seq_val_tokens[i] tokens instead of the dense + // [batch, N+1] layout, via the base SpeculativeWorkerImpl per-seq builder. + SpeculativeWorkerImpl::prepare_validate_inputs( + prepared_input, validate_input, per_seq_val_tokens); + } validate_input.input_params.embedding.input_embedding = torch::Tensor(); record_metadata_ready_event(*prepare_stream_, validate_input); } @@ -1352,7 +1463,8 @@ void DFlashWorkerImpl::write_context_kv( void DFlashWorkerImpl::write_target_context_to_cache( const ForwardInput& input, - const SampleOutput& validate_output) { + const SampleOutput& validate_output, + const torch::Tensor& confidence) { const torch::Tensor& accepted_embeddings = validate_output.embeddings; CHECK(accepted_embeddings.defined()) << "DFlash validate target embeddings are undefined."; @@ -1423,7 +1535,8 @@ void DFlashWorkerImpl::write_target_context_to_cache( input.input_params.embedding.request_ids, validate_output.next_tokens, validate_output.embeddings, - options_.num_speculative_tokens()); + options_.num_speculative_tokens(), + confidence); } // ----------------------------------------------------------------------------- @@ -1471,6 +1584,12 @@ std::vector DFlashWorkerImpl::compute_adaptive_prefix_lengths( return {}; } + return compute_prefix_lengths_from_probs(probs_for_controller, input); +} + +std::vector DFlashWorkerImpl::compute_prefix_lengths_from_probs( + const torch::Tensor& probs_for_controller, + const ForwardInput& input) { const int32_t batch_size = input.input_params.meta.num_sequences; std::vector per_seq_kv_lens(static_cast(batch_size), 0.0); const Slice kv_seq_lens = @@ -1488,21 +1607,45 @@ std::vector DFlashWorkerImpl::compute_adaptive_prefix_lengths( return prefix_lengths; } -void DFlashWorkerImpl::apply_per_seq_varlen_prune( - const ForwardInput& input, - ForwardInput& validate_input, - const std::vector& per_seq_val_tokens) { - const int32_t num_sequences = input.input_params.meta.num_sequences; - CHECK_EQ(static_cast(per_seq_val_tokens.size()), num_sequences); - c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); - ForwardInput prepared_input = input; - prepared_input.metadata_ready_event.reset(); - ForwardInput new_validate; - SpeculativeWorkerImpl::prepare_validate_inputs( - prepared_input, new_validate, per_seq_val_tokens); - new_validate.input_params.embedding.input_embedding = torch::Tensor(); - record_metadata_ready_event(*prepare_stream_, new_validate); - validate_input = std::move(new_validate); +std::vector DFlashWorkerImpl::decide_lagged_prefix_lengths( + const ForwardInput& input) { + const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + if (adaptive_spec_controller_ == nullptr || + !adaptive_spec_controller_->enabled() || embedding_cache_ == nullptr) { + return {}; + } + const auto& embedding = input.input_params.embedding; + if (embedding.embedding_ids.empty()) { + return {}; + } + const int32_t batch_size = input.input_params.meta.num_sequences; + if (static_cast(embedding.embedding_ids.size()) != batch_size) { + return {}; + } + + // Prune from the PREVIOUS step's confidence (lag-1): its D2H already + // completed at last step's write, so this decision is pure host and no longer + // waits on this step's draft forward — it overlaps the in-flight draft. + EmbeddingCache::LaggedConfidence lagged = + embedding_cache_->read_lagged_confidence(embedding.embedding_ids, + embedding.request_ids, + num_speculative_tokens); + std::vector prefix_lengths = + compute_prefix_lengths_from_probs(lagged.confidence, input); + if (prefix_lengths.empty()) { + return {}; + } + // Freshness fallback: a request with no usable lagged confidence (first + // decode step, or slot recycled by a new request) must NOT be pruned on a + // zero-filled row — force full width, the cost-model-independent analog of + // SGLang's survival=1. + CHECK_EQ(static_cast(prefix_lengths.size()), batch_size); + for (int32_t i = 0; i < batch_size; ++i) { + if (!lagged.valid[static_cast(i)]) { + prefix_lengths[static_cast(i)] = num_speculative_tokens; + } + } + return prefix_lengths; } void DFlashWorkerImpl::record_validate_metrics( diff --git a/xllm/core/runtime/dflash_worker_impl.h b/xllm/core/runtime/dflash_worker_impl.h index 00fd8d3097..1f101ab7c4 100644 --- a/xllm/core/runtime/dflash_worker_impl.h +++ b/xllm/core/runtime/dflash_worker_impl.h @@ -101,14 +101,36 @@ class DFlashWorkerImpl : public SpeculativeWorkerImpl { DraftProposal proposal; // DSpark ConfidenceHead output for adaptive pruning; empty otherwise. torch::Tensor confidence_probs; + // Adaptive lag-confidence pruning decision, computed in step_decode from + // the PREVIOUS step's confidence (overlapping this step's draft forward) + // and consumed by run_validate. Per-seq validate prefix length; empty when + // lag confidence is off. Not produced by run_decode_draft. + std::vector lagged_prefix_lengths; + // Set when run_decode_draft already built the pruned varlen validate batch + // in the draft-overlap window (lag confidence on + a pruning decision). + // run_validate then skips its own metadata rebuild and consumes these + // directly. varlen_prebuilt=false means the dense batch is in + // validate_input and run_validate takes the legacy this-step decision + + // rebuild path. + bool varlen_prebuilt = false; + std::vector per_seq_val_tokens; + int32_t max_val_tokens = 0; // No-sync draft inputs must outlive validation's stream sync. std::vector> retained_inputs; }; // virtual: DSpark overrides the draft sampling (parallel block sample -> // one forward + sequential Markov-head sampling loop). - virtual DraftBlock run_decode_draft(const ForwardInput& input, - ForwardInput& validate_input); + // + // lagged_prefix_lengths (lag confidence only): the pre-draft prune decision. + // When it prunes, run_decode_draft builds the pruned varlen validate batch in + // the draft-overlap window (setting DraftBlock.varlen_prebuilt) so the whole + // rebuild overlaps the in-flight draft instead of sitting on run_validate's + // critical path. Empty => build the dense batch as before. + virtual DraftBlock run_decode_draft( + const ForwardInput& input, + ForwardInput& validate_input, + const std::vector& lagged_prefix_lengths = {}); // Block layout hook: false (DFlash) -> query_width N+1, slot 0 is the // un-selected anchor; true (DSpark) -> query_width N, every position predicts @@ -117,13 +139,31 @@ class DFlashWorkerImpl : public SpeculativeWorkerImpl { // stays here and a subclass flips one bit. virtual bool sample_from_anchor() const { return false; } + // Build the validate batch in the draft-overlap window and record the prune + // decision into `draft_block`. Both DFlash and DSpark run_decode_draft call + // this in place of the bare prepare_validate_inputs, so the lag overlap-build + // covers both algorithms. When lagged_prefix_lengths prunes, builds the + // pruned varlen batch (sets draft_block.varlen_prebuilt / per_seq_val_tokens + // / max_val_tokens); otherwise builds the dense batch (varlen_prebuilt=false) + // exactly as before. + void prepare_overlap_validate_input( + const ForwardInput& input, + ForwardInput& validate_input, + const std::vector& lagged_prefix_lengths, + DraftBlock& draft_block); + // Shared with subclasses (DSpark): build the N/N+1-wide draft query block and // the target validate input. A DSpark override of run_decode_draft calls both // before its draft forward. + // + // per_seq_val_tokens: when non-empty, build a *pruned varlen* validate batch + // (Σ per_seq_val_tokens[i] tokens) instead of the dense [batch, N+1] batch. void prepare_query_inputs(const ForwardInput& input, ForwardInput& query_input); - void prepare_validate_inputs(const ForwardInput& input, - ForwardInput& validate_input); + void prepare_validate_inputs( + const ForwardInput& input, + ForwardInput& validate_input, + const std::vector& per_seq_val_tokens = {}); private: bool draft_use_block_parallel_rows() const { @@ -179,15 +219,29 @@ class DFlashWorkerImpl : public SpeculativeWorkerImpl { const DraftBlock& draft_block, const ForwardInput& input); - // Zero out draft probs beyond each sequence's prefix_len so the rejection - // Per-seq varlen prune: rebuild validate_input as a true varlen - // [Σ per_seq_val_tokens[i], ...] batch so target forward only spends - // compute on tokens each seq's prefix_len actually needs. Reuses the base - // SpeculativeWorkerImpl per-seq builder. - void apply_per_seq_varlen_prune( - const ForwardInput& input, - ForwardInput& validate_input, - const std::vector& per_seq_val_tokens); + // Core of the adaptive decision, shared by the this-step and lag-confidence + // paths: build per-seq kv lengths and run the controller on the given + // [batch, num_speculative_tokens] probs. The cost model is reused verbatim. + std::vector compute_prefix_lengths_from_probs( + const torch::Tensor& probs_for_controller, + const ForwardInput& input); + + // Lag-confidence decision (enable_lag_confidence): prune from the PREVIOUS + // step's confidence read from the embedding cache, so the decision does not + // data-depend on this step's draft forward. Requests with no fresh lagged + // confidence (first step / recycled slot) fall back to full width. Empty when + // adaptive is off. + std::vector decide_lagged_prefix_lengths(const ForwardInput& input); + + // Convert a per-seq prefix_len vector into per-seq validate widths + // (prefix_len + 1 bonus). Writes per_seq_val_tokens (empty when + // prefix_lengths is empty) and max_val_tokens; returns did_prune (true iff + // any seq's width is below the full N+1). Shared by the lag overlap-build + // path (run_decode_draft) and the legacy this-step path (run_validate). + bool prefix_lengths_to_val_tokens(const std::vector& prefix_lengths, + int32_t batch_size, + std::vector* per_seq_val_tokens, + int32_t* max_val_tokens) const; // Record precise (draft, accepted) counters. Padded -1 slots at positions // past per_seq_val_tokens[i]-1 are excluded — the count only walks each @@ -213,8 +267,10 @@ class DFlashWorkerImpl : public SpeculativeWorkerImpl { const torch::Tensor& positions_device, const torch::Tensor& new_cache_slots_device); - void write_target_context_to_cache(const ForwardInput& input, - const SampleOutput& validate_output); + void write_target_context_to_cache( + const ForwardInput& input, + const SampleOutput& validate_output, + const torch::Tensor& confidence = torch::Tensor()); protected: std::unique_ptr draft_impl_; diff --git a/xllm/core/runtime/dspark_worker_impl.cpp b/xllm/core/runtime/dspark_worker_impl.cpp index 571864cbbc..5afbdd1d16 100644 --- a/xllm/core/runtime/dspark_worker_impl.cpp +++ b/xllm/core/runtime/dspark_worker_impl.cpp @@ -38,7 +38,8 @@ DSparkWorkerImpl::DSparkWorkerImpl(const ParallelArgs& parallel_args, DSparkWorkerImpl::DraftBlock DSparkWorkerImpl::run_decode_draft( const ForwardInput& input, - ForwardInput& validate_input) { + ForwardInput& validate_input, + const std::vector& lagged_prefix_lengths) { Timer timer; // Same input build as DFlash, but sample_from_anchor()==true makes the query @@ -84,8 +85,12 @@ DSparkWorkerImpl::DraftBlock DSparkWorkerImpl::run_decode_draft( << "DSpark draft forward must return logits."; // Match DFlash's host/device overlap: validation input construction only // reads the original input, so prepare it after the asynchronous draft - // launch instead of delaying that launch. - prepare_validate_inputs(input, validate_input); + // launch instead of delaying that launch. Under lag confidence the lagged + // decision lets us build the pruned varlen batch here, overlapping the + // in-flight draft and keeping the rebuild off run_validate's critical path. + DraftBlock draft_block; + prepare_overlap_validate_input( + input, validate_input, lagged_prefix_lengths, draft_block); const int64_t num_rows = draft_output->logits.size(/*dim=*/0); CHECK_EQ(num_rows % num_speculative_tokens, 0) @@ -124,7 +129,6 @@ DSparkWorkerImpl::DraftBlock DSparkWorkerImpl::run_decode_draft( base_logits, last_hidden, anchor_token_ids, sampling_params_on_device); }(); - DraftBlock draft_block; draft_block.proposal = std::move(sampled.proposal); draft_block.confidence_probs = std::move(sampled.confidence_probs); draft_block.retained_inputs = take_retained_inputs(*draft_output); diff --git a/xllm/core/runtime/dspark_worker_impl.h b/xllm/core/runtime/dspark_worker_impl.h index 4eeda0c93e..65c54ad82f 100644 --- a/xllm/core/runtime/dspark_worker_impl.h +++ b/xllm/core/runtime/dspark_worker_impl.h @@ -45,8 +45,10 @@ class DSparkWorkerImpl final : public DFlashWorkerImpl { bool sample_from_anchor() const override { return true; } // Build the DSpark query and produce the complete proposal block. - DraftBlock run_decode_draft(const ForwardInput& input, - ForwardInput& validate_input) override; + DraftBlock run_decode_draft( + const ForwardInput& input, + ForwardInput& validate_input, + const std::vector& lagged_prefix_lengths = {}) override; private: // sample_block's two outputs: the sampled draft proposal and, when a diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 281968d6db..e7908a874a 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -92,6 +92,8 @@ struct Options { PROPERTY(double, adaptive_speculative_min_gain) = 0.0; + PROPERTY(bool, enable_lag_confidence) = false; + // enable speculative decode PROPERTY(bool, enable_speculative_decode) = false; diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index df69d84918..7f9b13fe3a 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -174,6 +174,7 @@ Options create_options(const std::string& instance_name, bool is_local) { speculative_config.enable_adaptive_speculative_decode()) .adaptive_speculative_min_gain( speculative_config.adaptive_speculative_min_gain()) + .enable_lag_confidence(speculative_config.enable_lag_confidence()) .num_request_handling_threads( service_config.num_request_handling_threads()) .communication_backend(parallel_config.communication_backend())