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

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

//! 2026-09-25: Post-`</think>` symmetric mask for `</think>` and `<think>`.
//!
//! While `seq.think_ended` is set, masks the close token
//! (`ctx.think_end_token`) and the open token (`seq.think_start_token`) to
//! `-inf`, so the model can neither repeat `</think>` nor re-open `<think>`.
//! 2026-10-14: Off when the model's `thinking_steering` is off: vLLM masks
//! neither, so a repeated `</think>` stays content and `<think>` re-opens
//! reasoning (`decode_logits_step/per_token.rs`).
//!
//! Owner: scheduler.
//! Invariants: none beyond the types.

use super::{LogitsContext, LogitsProcessor, ProcessorOutcome};
use crate::scheduler::ActiveSeq;

pub struct PostCloseThinkMask;

impl LogitsProcessor for PostCloseThinkMask {
    fn apply(
        &self,
        logits: &mut [f32],
        a: &mut ActiveSeq,
        ctx: &LogitsContext,
    ) -> ProcessorOutcome {
        if ctx.watchdog.thinking_steering && a.think_ended {
            if let Some(end_tok) = ctx.think_end_token {
                let end_idx = end_tok as usize;
                if end_idx < logits.len() {
                    logits[end_idx] = f32::NEG_INFINITY;
                }
            }
            if let Some(start_tok) = a.think_start_token {
                let start_idx = start_tok as usize;
                if start_idx < logits.len() {
                    logits[start_idx] = f32::NEG_INFINITY;
                }
            }
        }
        ProcessorOutcome::Continue
    }

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