feat(metrics): spec-decode acceptance counters for the DFlash draft path

[scatyf3]

Description

Fixes #604 , add acceptance length to /metrics via snapshot.

Type of Change

New feature (non-breaking change which adds functionality)

SpecDecodeCounters

First, we maintain the SpecDecodeCounters, where this counter is mirror from vllm vllm/rust/src/engine-core-client/src/protocol/stats.rs at 8e61b646e2d157f9b93451fa048f9c8530c8a67b · vllm-project/vllm · GitHub

However, we use fix length array for num_accepted_tokens_per_pos instead of vector, this is because LoadSnapshot need copy trait. Thus we hardcode a max MAX_SPEC_TOKENS = 32, which is hard to exceed it in real practice.

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct SpecDecodeCounters {
    /// Configured max draft tokens proposed per verify step (`K`), and the
    /// number of leading [`Self::num_accepted_tokens_per_pos`] entries that
    /// carry meaning. Never exceeds [`MAX_SPEC_TOKENS`].
    pub num_spec_tokens: u64,
    /// Draft proposals — one per request per verify step.
    pub num_drafts: u64,
    /// Draft tokens proposed for verification (pre-acceptance) in total.
    pub num_draft_tokens: u64,
    /// Draft tokens accepted in total, excluding the bonus token.
    pub num_accepted_tokens: u64,
    /// Accepted-token count indexed by draft position: `[i]` is how often the
    /// `i`-th draft was accepted. Only `[..num_spec_tokens]` is meaningful; the
    /// tail stays zero so the array width can be a constant.
    pub num_accepted_tokens_per_pos: [u64; MAX_SPEC_TOKENS],
}

we set up an observe draft function to record sd status after verify

// openinfer-engine/src/engine.rs
impl SpecDecodeCounters {
    /// Fold one request's verify outcome into the totals
    pub fn observe_draft(&mut self, num_draft_tokens: usize, num_accepted: usize) {
        self.num_drafts += 1;
        self.num_draft_tokens += num_draft_tokens as u64;
        self.num_accepted_tokens += num_accepted as u64;
        let tallied = num_accepted.min(self.num_spec_tokens as usize);
        for slot in self.num_accepted_tokens_per_pos.iter_mut().take(tallied) {
            *slot += 1;
        }
    }
}

// openinfer-qwen3/src/executor/spec.rs
impl Qwen3Executor {
	pub(super) fn execute_speculative_verify_impl(
	        &mut self,
        plan: VerifyPlan<'_>,
    ) -> Result<VerifyResult> {
	    // verify
	    // record
        if let Some(counters) = self.spec_decode_counters.as_mut() {
            for (req, req_result) in plan.requests.iter().zip(&result.requests) {
                let num_draft_tokens = req.as_slice().len().saturating_sub(1);
                counters.observe_draft(num_draft_tokens, req_result.matched_draft_tokens);
            }
        }
}

we test counter's correctness via spec_counters_observe_draft_tallies_positions and spec_counters_clamp_oversized_k

convert and publish sd status via frontend bridge

we Carry SpecDecodeCounters on the LoadSnapshot and fill spec_decoding_stats in the bridge.

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct LoadSnapshot {
    pub kv_used_blocks: u64,
    pub kv_total_blocks: u64,
    /// Requests currently occupying a decode/prefill slot.
    pub num_running_reqs: u64,
    /// Requests admitted but not yet running (KV pressure, prefetch wait).
    pub num_waiting_reqs: u64,
    /// Cumulative spec-decode counters, or `None` when no draft model is loaded.
    pub spec_decode: Option<SpecDecodeCounters>,
}

we calculate delta in frontend bridge, convert our SpecDecodeCounters to vllm frontend format SpecDecodingStats and publish them:

fn spec_decode_delta(last: &SpecDecodeCounters, cur: &SpecDecodeCounters) -> SpecDecodingStats {
    let width = (cur.num_spec_tokens as usize).min(MAX_SPEC_TOKENS);
    let num_accepted_tokens_per_pos = cur.num_accepted_tokens_per_pos[..width]
        .iter()
        .zip(&last.num_accepted_tokens_per_pos)
        .map(|(cur_pos, last_pos)| cur_pos.saturating_sub(*last_pos))
        .collect();
    SpecDecodingStats {
        num_spec_tokens: cur.num_spec_tokens,
        num_drafts: cur.num_drafts.saturating_sub(last.num_drafts),
        num_draft_tokens: cur.num_draft_tokens.saturating_sub(last.num_draft_tokens),
        num_accepted_tokens: cur
            .num_accepted_tokens
            .saturating_sub(last.num_accepted_tokens),
        num_accepted_tokens_per_pos,
    }
}

async fn publish_scheduler_stats(
    engine_index: u32,
    mut load_rx: watch::Receiver<LoadSnapshot>,
    output_tx: mpsc::UnboundedSender<EngineCoreOutputs>,
    shutdown: CancellationToken,
) -> Result<()> {
    let mut last_spec = SpecDecodeCounters::default();
    loop {
        let spec_decoding_stats = if let Some(cur) = &snapshot.spec_decode {
            let delta = spec_decode_delta(&last_spec, cur);
            last_spec = *cur;
            // Intervals with no verify step are the common case — prefill,
            // idle, and plain decode all publish without drafting. Reporting
            // one would divide by a zero `num_drafts` in the frontend's
            // acceptance-rate log, so leave the field `None` there, as vLLM's
            // own scheduler does. Nothing is lost by dropping it: every counter
            // moves only inside `observe_draft`, so a zero `num_drafts` delta
            // means nothing else moved either.
            (delta.num_drafts > 0).then_some(delta)
        } else {
            // No drafter (or one that just went away): forget the totals so a
            // drafter loaded later starts its diff from zero instead of being
            // saturated away against a stale high-water mark.
            last_spec = SpecDecodeCounters::default();
            None
        };
		let stats = SchedulerStats {
            num_running_reqs: snapshot.num_running_reqs,
            num_waiting_reqs: snapshot.num_waiting_reqs,
            kv_cache_usage: if snapshot.kv_total_blocks == 0 {
                0.0
            } else {
                snapshot.kv_used_blocks as f64 / snapshot.kv_total_blocks as f64
            },
            spec_decoding_stats,
            ..SchedulerStats::default()
        };
    }

here, correctness is checked by spec_delta_telescopes_to_cumulative

e2e test

To finalize our counter's correctness test, we use the measure metric from docs/models/qwen3/dspark-integration.md as reference. where they has detailed data from acceptance length to per position histgram:

config	rounds	mean accepted draft	zero-accept	full-7	hist 0..7
DSpark	19,294	2.52	29.2%	17.4%	[5636, 3942, 2394, 1549, 1042, 742, 639, 3350]
DFlash	21,214	2.30	32.2%	13.9%	[6838, 4340, 2685, 1648, 1071, 962, 731, 2939]

document's mesurement protocol is: 5090 GPU / CUDA 13.1，target Qwen3-4B，vllm-bench --temperature 0 --ignore-eos，draft dspark_qwen3_4b_block7 vs dflash_qwen3_4b_block7, markov_rank=0 , use chat(sharegpt) + poem(sonnet) + rand(random) + code(speed-bench), c1/c4/c8.

I use A6000+CUDA 12.8 with default seed=0, other remains the same. The modified test script is in my branch (not in this pr)

1. tools/bench/run_spec_accept_sweep.sh
2. tools/bench/spec_accept_metrics.py verify-log args

our result is shown in this table.

config	rounds	mean accepted draft	zero-accept	full-7	hist 0..7
DSpark	21,256	2.76	27.5%	20.6%	[5836, 3872, 2635, 1641, 1235, 967, 686, 4384]
DFlash	23,808	2.36	30.3%	14.3%	[7223, 4894, 3180, 1997, 1314, 1023, 763, 3414]

Due to original document do not report prompt subsampling seed and the difference from GPU, result is small difference but the trends remain same. The /metrics report and the dflash_lane.rs debug log agree exactly — verify-log compares rounds, accepted tokens, and the histogram bin-by-bin over the same server lifetime, and every bin is equal on both drafters. We can assert that our spec decoding counter implementation is correct.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 26423fd860

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Keep Qwen3.5 LoadSnapshot literals buildable

Adding spec_decode as a required LoadSnapshot field breaks the Qwen3.5 feature path: openinfer-qwen35/src/scheduler.rs:838 still constructs LoadSnapshot { kv_used_blocks, kv_total_blocks, num_running_reqs, num_waiting_reqs } without either spec_decode: None or ..LoadSnapshot::default(), so cargo ... --features qwen35 fails with a missing-field error and the supported Qwen3.5 engine cannot build. Please update that scheduler publisher (and any other exhaustive literals) when extending this shared struct.

AGENTS.md reference: AGENTS.md:L9-L14

Useful? React with 👍 / 👎.

[xiaguan]

Thanks for adding this plumbing. The core transport choice looks right: the scheduler publishes cumulative counters so a coalescing watch cannot lose increments, while the frontend converts them to interval deltas before calling the vLLM metrics API. The manual real-model correlation in the PR description is also useful evidence.

I found several items that need to be addressed before merge:

1. Please fix the existing automated P1 in the Qwen3.5 scheduler. openinfer-qwen35/src/scheduler.rs:838 still constructs LoadSnapshot without the new spec_decode field. That package is not in the current CI matrix, but this is a required-field Rust struct literal and will fail when Qwen3.5 is compiled. The automated review is correct and has not been addressed yet.

2. Please restore green CI. CPU Clippy currently fails because LoadSnapshot grew to 328 bytes and openinfer-sim/tests/frontend_e2e.rs:185 passes it by value (large_types_passed_by_value). The DCO sign-off check is also failing.

3. Please give the metrics types a module boundary and reduce the commentary. I do not think the placement is the contributor's fault: LoadSnapshot already lived in engine.rs, and there was no existing engine metrics module to extend. However, this PR adds about 128 lines there and takes engine.rs from 940 to 1068 lines. Please introduce an openinfer-engine/src/metrics.rs (with the necessary re-exports) for LoadSnapshot, SpecDecodeCounters, and their focused tests. While moving it, shorten the comments to the invariants a maintainer actually needs. In particular, the long checkpoint/current-workspace/CI narrative around MAX_SPEC_TOKENS and the repeated idle-interval narration obscure a fairly small contract.

4. Please tighten the tests instead of testing arithmetic several times. The position-prefix tally test is valuable: [2, 1, 0] is a non-obvious representation and should stay. The async publisher test is also valuable because it exercises first publish, idle omission, and resume. The extra tail-zero/sum/monotonic assertions and most of their narration are redundant. spec_delta_telescopes_to_cumulative mostly re-tests subtraction and currently does not even cover per-position deltas; please fold the meaningful coalescing/per-position assertion into the async publisher test and remove the arithmetic-only test. A real /metrics E2E would be better, but I am fine tracking that as a follow-up rather than blocking this PR because the PR already includes a manual real-model comparison.

5. Do not silently change the configured K in the public snapshot. SpecDecodeCounters::new stores min(K, 32) in a field documented as the configured num_spec_tokens. For a wider checkpoint the totals remain exact, but the metric reports a different K and silently presents a truncated acceptance curve as complete. The simplest contract is to reject unsupported K during speculative-model loading; alternatively, preserve the actual K and model truncation explicitly instead of overwriting it.

Once these are addressed, the remaining executor → scheduler → bridge → Prometheus data path looks sound to me.

[scatyf3]

Thank you for your review~ I'll fix my PR by the end of this week at the latest

[scatyf3]

Hi, I already fix the PR you mentioned

1. misc:
   1. add spec_decode to qwen3.5 path
   2. fix cpu clippy by pass loadsnapshot as reference instead of pass it directly
   3. fix the DCO sign-off check
2. reconstruct:
   move LoadSnapshot and SpecDecodeCounters to openinfer-engine/src/metrics.rs
3. shorten the comments t
4. test
   1. remove  extra tail-zero/sum/monotonic assertions
   2. add cpu sim e2e test
5. remove silent change for K = min(K, 32), now engine will reject spec decoding request when K>32

[xiaguan]

LGTM — all review items are addressed and CI is green. Please rebase onto main and we're good to go.

[scatyf3]

sure, I'll do that