met 1.0.0-beta-preview

// SPDX-License-Identifier: MIT OR Apache-2.0

//! 2026-09-25: single-sequence DFlash verify of a variable-length draft
//! block.
//!
//! Owner: scheduler.
//! Invariants:
//! - Multi-rank (2026-10-08): the workers run the same K-row verify forward,
//!   announced by `EP_CMD_VERIFY_KGAMMA` before it, and receive the committed
//!   row count before any token is emitted (an emit can finish the sequence).
//!   The wire shape is in `metrale_model_layers::speculative::kgamma_wire`.
//!   A single rank sends nothing.
//! - No logprobs: tokens are emitted with `None`.

use super::*;

/// 2026-09-25: verify `[last_token, drafts..]` with
/// `model.decode_verify_dflash`. Drafts are accepted up to the first one
/// that differs from its row's pick; the pick at that row is emitted as the
/// bonus token and the later drafts are dropped. Picks go through
/// `verify_pick_all_with_pipeline` unless `dflash_masked_verify` is off on a
/// DFlash serve.
pub fn step_verify_dflash(
    model: &dyn Model,
    a: &mut ActiveSeq,
    sched: &crate::scheduler::sched_ctx::SchedCtx,
    drafts: &[u32],
    num_drafts: usize,
    verify_ctx: &crate::scheduler::logit_processors::LogitsContext,
    dflash_verify_raw_argmax: bool,
) {
    if let Err(e) = model.sync_secondary() {
        tracing::error!("sync_secondary: {e:#}");
        a.finished = true;
        return;
    }

    let mut tokens = Vec::with_capacity(drafts.len() + 1);
    tokens.push(a.last_token);
    tokens.extend_from_slice(drafts);

    let ep = model.is_ep();
    if ep && let Err(e) = ep_send_verify_kgamma(model, a.seq.slot_idx as u32, &tokens) {
        tracing::error!("EP broadcast verify_kgamma: {e:#}");
        super::lifecycle::fail_sequence(a, format!("EP broadcast verify_kgamma: {e:#}"));
        return;
    }

    // 2026-09-25: `METRALE_DFLASH_STEP_TIMING=1` logs the verify and propose
    // walls separately at the end of the step.
    let step_timing = sched.levers.dflash_step_timing;
    let t_verify = sched.io.clock.now();
    let verified_argmax = match model.decode_verify_dflash(&tokens, &mut a.seq, 0) {
        Ok(v) => v,
        Err(e) => {
            tracing::error!("decode_verify_dflash: {e:#}");
            // 2026-10-08: The workers wait for a verdict after their verify; send the anchor
            // row alone so none takes the next command for it.
            if ep && let Err(e2) = model.ep_broadcast_cmd(1) {
                tracing::error!("EP broadcast verify_kgamma verdict after a failed verify: {e2:#}");
            }
            a.finished = true;
            return;
        }
    };
    let verify_ms = if step_timing {
        sched
            .io
            .clock
            .now()
            .saturating_duration_since(t_verify)
            .as_secs_f64()
            * 1000.0
    } else {
        0.0
    };
    a.last_token_time = sched.io.clock.now();

    // 2026-09-25: on a DFlash serve with `dflash_masked_verify` off, judge on
    // the raw argmax, with no masks or penalties.
    let verified = if dflash_verify_raw_argmax && !sched.levers.dflash_masked_verify {
        verified_argmax
    } else {
        crate::scheduler::verify_pipeline_helper::verify_pick_all_with_pipeline(
            model,
            &verified_argmax,
            a,
            verify_ctx,
            0,
        )
    };

    // 2026-09-25: `drafts[i]` is accepted while it equals `verified[i]`, the
    // pick for the token after `tokens[i]`, and only while a bonus row
    // `verified[i + 1]` exists.
    let mut num_accepted = 0usize;
    for i in 0..drafts.len() {
        if i + 1 >= verified.len() {
            break;
        }
        if drafts[i] == verified[i] {
            num_accepted += 1;
        } else {
            break;
        }
    }

    // 2026-10-08: The workers commit the anchor row plus the accepted drafts.
    if ep && let Err(e) = model.ep_broadcast_cmd((num_accepted + 1) as u32) {
        tracing::error!("EP broadcast verify_kgamma verdict: {e:#}");
        super::lifecycle::fail_sequence(a, format!("EP broadcast verify_kgamma verdict: {e:#}"));
        return;
    }

    // 2026-09-25: with `METRALE_DFLASH_ADAPTIVE=1`, a low mean over the
    // accept window suspends this sequence's speculation (`adaptive_spec`).
    crate::scheduler::adaptive_spec::record_verify(a, num_accepted, sched);
    // 2026-09-25: the gamma resolver's accept signal: was the first draft
    // accepted.
    sched.dflash_rung.observe_step(num_accepted >= 1);
    // 2026-10-10: The single-stream adaptive draft count's rates (`--dflash-adaptive-k`).
    sched
        .dflash_rung
        .adaptive_observe(&mut a.spec_adapt.k, drafts.len(), num_accepted);

    // 2026-09-25: the verify forward advanced `seq_len` and `seq.tokens` by
    // `tokens.len()`. Keep the pre-verify prefix, the accepted drafts and the
    // bonus position (`pre_verify_len + num_accepted + 1`) and drop the rest.
    // `emit_token` writes `output_tokens` only, so the bonus stays where the
    // verify put it in `seq.tokens`.
    let pre_verify_len = a.seq.seq_len.saturating_sub(tokens.len());
    let target_seq_len = pre_verify_len + num_accepted + 1;
    let to_drop = a.seq.seq_len.saturating_sub(target_seq_len);
    if to_drop > 0 {
        a.seq.seq_len = target_seq_len;
        let pop_n = to_drop.min(a.seq.tokens.len());
        for _ in 0..pop_n {
            a.seq.tokens.pop();
        }
    }

    tracing::debug!(
        "CTX_VERIFY slot={} pre_verify_len={} na={} k={}",
        a.seq.slot_idx,
        pre_verify_len,
        num_accepted,
        drafts.len() + 1,
    );
    commit_verified_ctx(model, &mut a.seq, num_accepted, pre_verify_len, sched);

    for i in 0..num_accepted {
        emit_token(a, drafts[i], None, sched);
        if a.finished {
            return;
        }
    }

    // 2026-09-25: the bonus is `verified[num_accepted]`: the correction at
    // the first mismatch, or the token after a full accept.
    let bonus_idx = num_accepted;
    if bonus_idx < verified.len() {
        let bonus = verified[bonus_idx];
        emit_token(a, bonus, None, sched);
        if a.finished {
            return;
        }
        a.last_token = bonus;
    }

    sched.io.tel.count_spec_verify(
        "dflash",
        if num_accepted == drafts.len() {
            "accept_all"
        } else {
            "accept_partial"
        },
    );
    sched.io.tel.spec_verified(drafts.len(), num_accepted);

    tracing::info!(
        "DFLASH K=γ verify: γ={} accepted={}/{} ({:.0}%) seq_len={}",
        drafts.len(),
        num_accepted,
        drafts.len(),
        100.0 * (num_accepted as f64) / (drafts.len() as f64),
        a.seq.seq_len,
    );

    // 2026-09-25: commit the SSM state to rows `0..=num_accepted` of the
    // `drafts.len() + 1` verified rows; `commit_accepted_prefix` treats a
    // full accept as a no-op. A failed commit finishes the sequence.
    let k_verify = drafts.len() + 1;
    let total_accepted = num_accepted + 1;
    if let Err(e) = model.commit_accepted_prefix(&mut a.seq, total_accepted, k_verify) {
        tracing::error!("commit_accepted_prefix (dflash): {e:#}");
        a.finished = true;
        return;
    }

    let bonus_token_idx = total_accepted.saturating_sub(1);
    if let Err(e) = model.save_hidden_for_mtp(bonus_token_idx, 0) {
        tracing::error!("save_hidden_for_mtp (dflash): {e:#}");
    }

    if let Err(e) = model.trim_proposer_state(&mut a.seq, num_accepted, 0) {
        tracing::error!("trim_proposer_state: {e:#}");
    }

    // 2026-09-25: no propose while adaptive speculation has suspended this
    // sequence; with no drafts, `step_mtp` bootstraps it next step.
    let _mtp_grammar_mask = mtp_grammar_mask_for(a);
    let t_propose = sched.io.clock.now();
    // 2026-10-10: The next step's drafts: the adaptive draft count's choice on a single
    // stream (0 suspends speculation until the re-probe), else this step's count.
    let next_drafts = match sched
        .dflash_rung
        .adaptive_drafts(&a.spec_adapt.k, num_drafts)
    {
        Some(0) => {
            crate::scheduler::adaptive_spec::suspend(a, sched);
            0
        }
        Some(k) => k,
        None => num_drafts,
    };
    let mut proposed = false;
    if next_drafts > 0 && crate::scheduler::adaptive_spec::spec_allowed(a, sched) {
        proposed = true;
        match model.run_mtp_propose_multi(
            a.last_token,
            a.seq.seq_len,
            next_drafts,
            &mut a.seq,
            0,
            _mtp_grammar_mask.as_deref(),
        ) {
            Ok(d) if !d.is_empty() => a.pending_drafts = d,
            Ok(_) => {}
            Err(e) => tracing::error!("run_mtp_propose_multi (dflash): {e:#}"),
        }
    }
    // 2026-10-10: A step that proposed is a whole speculative step (verify, then propose):
    // its wall feeds the adaptive cost scale. The clock is read only when the controller is
    // installed.
    if proposed && sched.dflash_rung.adaptive_k_on() {
        let step_ms = sched
            .io
            .clock
            .now()
            .saturating_duration_since(t_verify)
            .as_secs_f64()
            * 1000.0;
        sched.dflash_rung.adaptive_step_ms(drafts.len(), step_ms);
    }
    if step_timing {
        let propose_ms = sched
            .io
            .clock
            .now()
            .saturating_duration_since(t_propose)
            .as_secs_f64()
            * 1000.0;
        tracing::info!(
            "DFLASH STEP_TIMING: verify={:.1}ms propose={:.1}ms (K={}, accepted={})",
            verify_ms,
            propose_ms,
            tokens.len(),
            num_accepted,
        );
    }
}

/// 2026-10-09: Give the drafter the verified rows `0..=num_accepted` of a single-sequence
/// verify on a DFlash serve, at RoPE base `pre_verify_len` (the anchor's position). Every verify
/// width calls it (K=γ here; K=2, K=3 and K=4 in their steps), so the context stays one row per
/// position whichever width ran. With `dflash_unified_ctx` (on unless
/// `METRALE_DFLASH_UNIFIED_CTX=0`) `commit_ctx` commits them; otherwise, when
/// `dflash_eagle_fix` is on, `dflash_eagle_kgamma_append` appends them. A failure of either is
/// logged and the step continues: the drafter only loses context.
pub(super) fn commit_verified_ctx(
    model: &dyn Model,
    seq: &mut SequenceState,
    num_accepted: usize,
    pre_verify_len: usize,
    sched: &crate::scheduler::sched_ctx::SchedCtx,
) {
    if sched.levers.dflash_unified_ctx {
        if let Err(e) = model.commit_ctx(seq, num_accepted + 1, pre_verify_len, 0) {
            tracing::error!("commit_ctx (verify): {e:#}");
        }
    } else if sched.levers.dflash_eagle_fix
        && let Err(e) = model.dflash_eagle_kgamma_append(seq, num_accepted, pre_verify_len)
    {
        tracing::error!("dflash_eagle_kgamma_append: {e:#}");
    }
}

/// 2026-10-08: Announce a K=γ verify of `tokens` in slot `slot` to the worker ranks: the
/// `(seq_id, cmd)` preamble, K, then the tokens in one bulk broadcast. Refuses a width the
/// worker would refuse, before sending anything.
fn ep_send_verify_kgamma(model: &dyn Model, slot: u32, tokens: &[u32]) -> anyhow::Result<()> {
    let k = metrale_model_layers::speculative::kgamma_width(tokens.len() as u32)?;
    model.ep_broadcast_cmd_for_seq(
        slot,
        metrale_model_layers::speculative::EP_CMD_VERIFY_KGAMMA,
    )?;
    model.ep_broadcast_cmd(k as u32)?;
    model.ep_broadcast_tokens(tokens)?;
    Ok(())
}