metrale_grammar/matcher/
batch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// BatchGrammarMatcher — batched bitmask fill / token acceptance.
4//
5// Port of `class BatchGrammarMatcher` from `cpp/grammar_matcher.cc`.
6// The C++ version uses a hand-rolled `ThreadPool`; this port uses the
7// crate's existing `rayon` dependency for the parallel fill path.
8//
9// SIMPLIFICATION vs C++
10// ---------------------
11// The C++ `BatchFillNextTokenBitmask` lets `indices` remap matcher `i`
12// to an arbitrary tensor slice. This port keeps that capability for the
13// sequential path; the parallel path requires the *natural* mapping
14// (`matchers[i]` -> slice `i`) so the per-thread mutable slices are
15// provably disjoint without `unsafe`. With an explicit `indices`
16// permutation the fill falls back to sequential, which is always
17// correct.
18
19use rayon::prelude::*;
20
21use super::bitmask::bitmask_size;
22use super::fill::FillError;
23use super::matcher::GrammarMatcher;
24
25/// A batched driver over many [`GrammarMatcher`]s — fills one packed
26/// bitmask tensor for a whole decode batch in parallel.
27#[derive(Debug, Clone)]
28pub struct BatchGrammarMatcher {
29    /// Upper bound on worker threads; `1` forces the sequential path.
30    max_threads: usize,
31}
32
33impl BatchGrammarMatcher {
34    /// Construct a batched matcher.
35    ///
36    /// `max_threads`: `None` picks `available_parallelism / 2` (the C++
37    /// `"auto"` default); `Some(n)` caps at `n` (must be >= 1).
38    #[must_use]
39    pub fn new(max_threads: Option<usize>) -> Self {
40        let max_threads = match max_threads {
41            Some(n) => {
42                assert!(n >= 1, "max_threads must be >= 1");
43                n
44            }
45            None => std::thread::available_parallelism()
46                .map(|n| (n.get() / 2).max(1))
47                .unwrap_or(1),
48        };
49        Self { max_threads }
50    }
51
52    /// The configured worker-thread cap.
53    #[must_use]
54    pub fn max_threads(&self) -> usize {
55        self.max_threads
56    }
57
58    /// Fill `bitmask` for every matcher in `matchers`.
59    ///
60    /// `bitmask` is one flat buffer of `matchers.len() *
61    /// bitmask_size(vocab)` words; matcher `i` writes slice `i`. Every
62    /// matcher must share the same vocabulary size. Port of
63    /// `BatchFillNextTokenBitmask` (natural-index form).
64    ///
65    /// Returns one `Result` per matcher (same order).
66    pub fn fill_next_token_bitmask(
67        &self,
68        matchers: &mut [GrammarMatcher],
69        bitmask: &mut [i32],
70        debug: bool,
71    ) -> Vec<Result<bool, FillError>> {
72        if matchers.is_empty() {
73            return Vec::new();
74        }
75        let vocab = matchers[0].tokenizer_info().vocab_size();
76        let words = bitmask_size(vocab);
77        assert!(
78            bitmask.len() >= matchers.len() * words,
79            "bitmask buffer too small for the batch"
80        );
81
82        if self.max_threads <= 1 || matchers.len() == 1 {
83            return matchers
84                .iter_mut()
85                .zip(bitmask.chunks_mut(words))
86                .map(|(m, slice)| m.fill_next_token_bitmask(slice, 0, debug))
87                .collect();
88        }
89
90        // Parallel path: matcher `i` <-> disjoint slice `i`.
91        matchers
92            .par_iter_mut()
93            .zip(bitmask.par_chunks_mut(words))
94            .map(|(m, slice)| m.fill_next_token_bitmask(slice, 0, debug))
95            .collect()
96    }
97
98    /// Accept one token per matcher. `token_ids[i]` is fed to
99    /// `matchers[i]`. Port of `BatchAcceptToken`.
100    pub fn accept_token(
101        matchers: &mut [GrammarMatcher],
102        token_ids: &[i32],
103        debug: bool,
104    ) -> Vec<bool> {
105        assert_eq!(
106            matchers.len(),
107            token_ids.len(),
108            "matchers and token_ids must have equal length"
109        );
110        matchers
111            .iter_mut()
112            .zip(token_ids)
113            .map(|(m, &t)| m.accept_token(t, debug))
114            .collect()
115    }
116
117    /// Accept one string per matcher. Port of `BatchAcceptString`.
118    pub fn accept_string(
119        matchers: &mut [GrammarMatcher],
120        input_strs: &[&str],
121        debug: bool,
122    ) -> Vec<bool> {
123        assert_eq!(
124            matchers.len(),
125            input_strs.len(),
126            "matchers and input_strs must have equal length"
127        );
128        matchers
129            .iter_mut()
130            .zip(input_strs)
131            .map(|(m, &s)| m.accept_string(s, debug))
132            .collect()
133    }
134}
metrale_grammar/matcher/batch.rs

metrale_grammar/matcher/
batch.rs
