met 1.0.0-beta-preview

Metrale Engine: pure Rust LLM inference server (HTTP API, scheduler, TUI, CLI).
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

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

//! 2026-09-25: Forced `</think>` injection once `force_end_thinking` is armed.
//!
//! Inside thinking, the stage asks
//! [`crate::scheduler::confidence::should_inject_think_end`] with the
//! sequence's `force_end_thinking` and `in_code_fence` and two values computed
//! here:
//! - `at_sentence_boundary`: the previous output token is set in
//!   `ctx.boundary_mask`;
//! - `defer_hard_override`: thinking has reached `THINK_DEFER_BUDGET_FACTOR`
//!   times the thinking budget (`THINK_DEFER_ABS_CEILING` with no budget), or
//!   `seq.sentence_defer_count` has reached `MAX_SENTENCE_DEFER_TOKENS`.
//!
//! When the gate says inject, every logit is set to `-inf` and `</think>` to
//! `0.0`, so `</think>` is the only finite logit. When `force_end_thinking` is
//! set but the injection does not happen, `seq.sentence_defer_count` is
//! incremented, which bounds the deferral at `MAX_SENTENCE_DEFER_TOKENS` steps.
//!
//! 2026-10-14: When the model's `thinking_steering` is off, the only thing
//! that arms `force_end_thinking` is a client thinking budget, and the close
//! is injected on the next step with no deferral, as vLLM's thinking-budget
//! processor does.
//!
//! Owner: scheduler.
//! Invariants: none beyond the types.

use super::{LogitsContext, LogitsProcessor, ProcessorOutcome};
use crate::scheduler::ActiveSeq;
use crate::scheduler::confidence::{
    MAX_SENTENCE_DEFER_TOKENS, THINK_DEFER_ABS_CEILING, THINK_DEFER_BUDGET_FACTOR,
    should_inject_think_end,
};

pub struct ForcedThinkEndInjector;

impl LogitsProcessor for ForcedThinkEndInjector {
    fn apply(
        &self,
        logits: &mut [f32],
        a: &mut ActiveSeq,
        ctx: &LogitsContext,
    ) -> ProcessorOutcome {
        let at_sentence_boundary = a
            .output_tokens
            .last()
            .copied()
            .and_then(|prev_tok| {
                ctx.boundary_mask
                    .clone()
                    .as_deref()
                    .and_then(|m| m.get(prev_tok as usize).copied())
            })
            .unwrap_or(false);
        let defer_hard_override = match a.thinking_budget {
            Some(b) => a.thinking_tokens >= b.saturating_mul(THINK_DEFER_BUDGET_FACTOR),
            None => a.thinking_tokens >= THINK_DEFER_ABS_CEILING,
        } || a.sentence_defer_count >= MAX_SENTENCE_DEFER_TOKENS
            || !ctx.watchdog.thinking_steering;
        if a.inside_thinking
            && should_inject_think_end(
                a.force_end_thinking,
                a.in_code_fence,
                at_sentence_boundary,
                defer_hard_override,
            )
            && let Some(end_tok) = ctx.think_end_token
        {
            let end_idx = end_tok as usize;
            if end_idx < logits.len() {
                for logit in logits.iter_mut() {
                    *logit = f32::NEG_INFINITY;
                }
                logits[end_idx] = 0.0;
            }
        } else if a.inside_thinking && a.force_end_thinking {
            // 2026-09-25: Armed but not injected this step: count it toward
            // the MAX_SENTENCE_DEFER_TOKENS override.
            a.sentence_defer_count = a.sentence_defer_count.saturating_add(1);
        }
        ProcessorOutcome::Continue
    }

    fn name(&self) -> &'static str {
        "forced_think_end_injector"
    }
}