sokuto_runtime/engine/
cache.rs

1//! # 中間特徴量キャッシュおよびゼロコピー連携モジュール
2//!
3//! 長文 `State` の中間特徴量キャッシュ、
4//! `xxHash3` 高速キー生成、`moka` 並行キャッシュストレージ、
5//! `ort` の `IoBinding` によるゼロコピーバインド、
6//! L3キャッシュ局所性最適化（`FP16`/`INT8` 圧縮）、
7//! 64バイトメモリアライメント、および `Tokio` ワーカースレッド競合隔離を提供する。
8
9use std::alloc::{Layout, alloc_zeroed, dealloc};
10use std::fmt;
11use std::ops::{Deref, DerefMut};
12use std::path::Path;
13use std::ptr::NonNull;
14use std::sync::atomic::{AtomicUsize, Ordering};
15use std::sync::{Arc, Mutex};
16use std::time::Duration;
17
18use half::f16;
19use moka::sync::Cache;
20use ort::memory::{Allocator, MemoryInfo};
21use ort::session::Session;
22use ort::value::{Tensor, Value};
23use sokuto_core::contract::calibration::CalibrationConfig;
24use sokuto_core::contract::model_spec::{
25    ModelInputDimensions, TENSOR_ATTENTION_MASK, TENSOR_INPUT_IDS, TENSOR_LOGITS,
26    TENSOR_OP_INDICES, TENSOR_QUERY_IDS, TENSOR_QUERY_MASK, TENSOR_STATE_HIDDEN_STATES,
27    TENSOR_STATE_MASK,
28};
29use sokuto_core::decision::{evaluate_question, evaluate_question_with_buf};
30use sokuto_core::gating::GatingConfig;
31use sokuto_core::schema::{Answer, Question};
32use xxhash_rust::xxh3::xxh3_64;
33
34use crate::engine::config::{ExecutionProvider, SessionConfig};
35use crate::engine::gating::{apply_gating_to_answer_with_energy, calculate_choice_energy};
36use crate::engine::session::create_session_builder;
37use crate::error::{Result, RuntimeError};
38use crate::tokenizer::{JevTokenizer, TokenizedQuestion};
39
40/// 64バイト境界メモリアライメント判定関数。
41///
42/// `AVX-512` や `AVX2` のベクトル化命令において、未アライメントアクセスによる
43/// セグメンテーション違反や性能劣化を防ぐためにアライメント状態を検証する。
44pub fn is_aligned_64<T>(ptr: *const T) -> bool {
45    (ptr as usize).is_multiple_of(64)
46}
47
48/// デュアルセッション推論における推奨デフォルト最大チャンクサイズ。
49///
50/// 長文 `State` 特徴量のリピート展開による一時的なヒープメモリ確保スパイク(OOM Killer リスク) を抑制するため、
51/// 4 サンプル程度を推奨上限とする。
52pub const DEFAULT_DUAL_MAX_BATCH_CHUNK_SIZE: usize = 4;
53
54/// 64バイト境界にアライメントされた連続ヒープメモリバッファ。
55///
56/// `SIMD` カーネルの高速演算要件（64バイトアライメント）を満たす安全な連続バッファを提供する。
57pub struct AlignedBuffer64<T: Copy> {
58    ptr: NonNull<T>,
59    len: usize,
60    capacity: usize,
61    layout: Layout,
62}
63
64impl<T: Copy> AlignedBuffer64<T> {
65    /// 指定された要素数で64バイトアライメントされたゼロ初期化バッファを生成する。
66    ///
67    /// # 引数
68    /// - `len`: 要素数。
69    pub fn with_capacity_zeroed(len: usize) -> Self {
70        if len == 0 {
71            let layout = Layout::from_size_align(0, 64).unwrap_or(Layout::new::<u8>());
72            return Self {
73                ptr: NonNull::dangling(),
74                len: 0,
75                capacity: 0,
76                layout,
77            };
78        }
79
80        let size = len
81            .checked_mul(std::mem::size_of::<T>())
82            .expect("メモリ確保サイズがオーバーフローしました。");
83        let layout = Layout::from_size_align(size, 64)
84            .expect("64バイトアライメントレイアウトの作成に失敗しました。");
85
86        // SAFETY: 有効なレイアウトと非ゼロサイズを検証済みであり、直後にゼロ初期化メモリとして確保する。
87        let raw_ptr = unsafe { alloc_zeroed(layout) as *mut T };
88        let ptr = NonNull::new(raw_ptr).expect("アライメントメモリの割り当てに失敗しました。");
89
90        Self {
91            ptr,
92            len,
93            capacity: len,
94            layout,
95        }
96    }
97
98    /// スライスから64バイトアライメントされたバッファを複製生成する。
99    ///
100    /// # 引数
101    /// - `slice`: 複製元スライス。
102    pub fn from_slice(slice: &[T]) -> Self {
103        let buf = Self::with_capacity_zeroed(slice.len());
104        if !slice.is_empty() {
105            // SAFETY: `buf` の長さは `slice.len()` と完全に一致しており、重なり合わない有効なポインタである。
106            unsafe {
107                std::ptr::copy_nonoverlapping(slice.as_ptr(), buf.ptr.as_ptr(), slice.len());
108            }
109        }
110        buf
111    }
112
113    /// バッファの要素数を取得する。
114    pub fn len(&self) -> usize {
115        self.len
116    }
117
118    /// バッファが空であるかを判定する。
119    pub fn is_empty(&self) -> bool {
120        self.len == 0
121    }
122
123    /// 内部生ポインタが64バイト境界に整列しているかを検証する。
124    pub fn is_aligned_64(&self) -> bool {
125        is_aligned_64(self.ptr.as_ptr())
126    }
127}
128
129impl<T: Copy> Deref for AlignedBuffer64<T> {
130    type Target = [T];
131
132    fn deref(&self) -> &Self::Target {
133        if self.len == 0 {
134            &[]
135        } else {
136            // SAFETY: 有効な割り当て領域と要素数を保持している。
137            unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
138        }
139    }
140}
141
142impl<T: Copy> DerefMut for AlignedBuffer64<T> {
143    fn deref_mut(&mut self) -> &mut Self::Target {
144        if self.len == 0 {
145            &mut []
146        } else {
147            // SAFETY: 排他的借用下で有効な割り当て領域と要素数を保持している。
148            unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
149        }
150    }
151}
152
153impl<T: Copy> Drop for AlignedBuffer64<T> {
154    fn drop(&mut self) {
155        if self.capacity > 0 && self.layout.size() > 0 {
156            // SAFETY: 構築時に確保した同一のポインタとレイアウトで解放する。
157            unsafe {
158                dealloc(self.ptr.as_ptr() as *mut u8, self.layout);
159            }
160        }
161    }
162}
163
164// SAFETY: 保持する型 `T` が `Send` であれば、スレッド間で安全に所有権を移動可能である。
165unsafe impl<T: Copy + Send> Send for AlignedBuffer64<T> {}
166// SAFETY: 保持する型 `T` が `Sync` であれば、スレッド間で安全に参照を共有可能である。
167unsafe impl<T: Copy + Sync> Sync for AlignedBuffer64<T> {}
168
169impl<T: Copy + fmt::Debug> fmt::Debug for AlignedBuffer64<T> {
170    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171        f.debug_list().entries(self.deref().iter()).finish()
172    }
173}
174
175/// 中間特徴量キャッシュのメモリ圧縮モード。
176///
177/// 8kトークンの隠れ状態（`FP32` で25.16MB）によるL3キャッシュ破壊を防ぎ、
178/// メモリバス帯域の飽和を抑制するための圧縮方式を指定する。
179#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
180pub enum CacheCompressionMode {
181    /// 圧縮なし（`FP32` の `Value` を直接キャッシュ保持し、ゼロコピーで直接バインド）。
182    #[default]
183    None,
184    /// `FP16` (`half::f16`) 半精度浮動小数点圧縮（50%の容量削減、誤差 <= 2e-5）。
185    Fp16,
186    /// トークン単位対称 `INT8` 量子化（75%の容量削減、誤差 <= 8e-4）。
187    Int8Symmetric,
188}
189
190/// `f32` スライスを `half::f16` へ圧縮する。
191///
192/// # 引数
193/// - `src`: 入力 `f32` スライス。
194pub fn compress_f32_to_fp16(src: &[f32]) -> Vec<f16> {
195    src.iter().map(|&x| f16::from_f32(x)).collect()
196}
197
198/// `half::f16` スライスを `f32` スライスへ展開する。
199///
200/// # 引数
201/// - `src`: 入力 `half::f16` スライス。
202/// - `dst`: 出力先 `f32` スライス。
203pub fn decompress_fp16_to_f32(src: &[f16], dst: &mut [f32]) {
204    assert_eq!(src.len(), dst.len(), "展開元と先の要素数が一致しません。");
205    for (s, d) in src.iter().zip(dst.iter_mut()) {
206        *d = s.to_f32();
207    }
208}
209
210/// `f32` スライスをトークン単位対称 `INT8` へ圧縮する。
211///
212/// 各トークンベクトル（長さ `hidden_size`）ごとに最大絶対値 $M = \max |x|$ を算出し、
213/// スケール $s = M / 127.0$ を用いて対称量子化を行う。
214///
215/// # 引数
216/// - `src`: 入力 `f32` スライス。
217/// - `hidden_size`: 隠れ層次元数。
218pub fn compress_f32_to_int8_symmetric(src: &[f32], hidden_size: usize) -> (Vec<i8>, Vec<f32>) {
219    assert!(hidden_size > 0, "hidden_sizeは1以上である必要があります。");
220    assert_eq!(
221        src.len() % hidden_size,
222        0,
223        "データ長はhidden_sizeの倍数である必要があります。"
224    );
225
226    let num_tokens = src.len() / hidden_size;
227    let mut quantized = Vec::with_capacity(src.len());
228    let mut scales = Vec::with_capacity(num_tokens);
229
230    for token_slice in src.chunks_exact(hidden_size) {
231        let mut max_abs = 0.0f32;
232        for &val in token_slice {
233            let abs_val = val.abs();
234            if abs_val > max_abs {
235                max_abs = abs_val;
236            }
237        }
238
239        let scale = if max_abs > 1e-12 {
240            max_abs / 127.0
241        } else {
242            1.0
243        };
244        scales.push(scale);
245
246        let inv_scale = 1.0 / scale;
247        for &val in token_slice {
248            let q = (val * inv_scale).round().clamp(-128.0, 127.0) as i8;
249            quantized.push(q);
250        }
251    }
252
253    (quantized, scales)
254}
255
256/// トークン単位対称 `INT8` スライスを `f32` スライスへ展開する。
257///
258/// # 引数
259/// - `src`: 量子化データスライス。
260/// - `scales`: 各トークンのスケールスライス。
261/// - `hidden_size`: 隠れ層次元数。
262/// - `dst`: 出力先 `f32` スライス。
263pub fn decompress_int8_symmetric_to_f32(
264    src: &[i8],
265    scales: &[f32],
266    hidden_size: usize,
267    dst: &mut [f32],
268) {
269    assert_eq!(src.len(), dst.len(), "展開元と先の要素数が一致しません。");
270    assert_eq!(
271        src.len() / hidden_size,
272        scales.len(),
273        "トークン数とスケール数が一致しません。"
274    );
275
276    for (token_idx, (q_chunk, &scale)) in
277        src.chunks_exact(hidden_size).zip(scales.iter()).enumerate()
278    {
279        let offset = token_idx * hidden_size;
280        for (i, &q) in q_chunk.iter().enumerate() {
281            dst[offset + i] = (q as f32) * scale;
282        }
283    }
284}
285
286/// キャッシュ内に保持されるテンソル実体ペイロード。
287pub enum CachedTensorPayload {
288    /// 非圧縮の `Value`。直接 `IoBinding::bind_input` へゼロコピーで受け渡し可能。
289    Raw(Arc<Value>),
290    /// `FP16` 圧縮済みテンソルデータ。
291    Fp16(Arc<Vec<f16>>),
292    /// トークン単位対称 `INT8` 量子化済みテンソルデータ。
293    Int8Symmetric {
294        /// 量子化された `i8` 配列。
295        quantized: Arc<Vec<i8>>,
296        /// 各トークンのスケール値配列。
297        scales: Arc<Vec<f32>>,
298    },
299}
300
301impl fmt::Debug for CachedTensorPayload {
302    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
303        match self {
304            Self::Raw(_) => write!(f, "CachedTensorPayload::Raw(Value)"),
305            Self::Fp16(data) => write!(f, "CachedTensorPayload::Fp16(len={})", data.len()),
306            Self::Int8Symmetric { quantized, scales } => write!(
307                f,
308                "CachedTensorPayload::Int8Symmetric(q_len={}, scales_len={})",
309                quantized.len(),
310                scales.len()
311            ),
312        }
313    }
314}
315
316/// キャッシュに格納される `State` 中間特徴量エントリ。
317#[derive(Debug)]
318pub struct CachedStateFeatures {
319    /// `xxHash3` (64-bit) によるキー。
320    pub hash_key: u64,
321    /// テンソル形状 (`[batch_size, seq_len, hidden_size]`)。
322    pub shape: Vec<usize>,
323    /// テンソル実体ペイロード。
324    pub payload: CachedTensorPayload,
325    /// テンソルが占有する実メモリバイト数（`moka` の `weigher` に使用）。
326    pub byte_size: usize,
327    /// 採用された圧縮モード。
328    pub compression_mode: CacheCompressionMode,
329    /// `State` の有効トークンアテンションマスク列 (`[batch_size, seq_len]`)。
330    pub state_mask: Vec<i64>,
331    /// キャッシュされた State 原文テキスト（または先頭抜粋）。
332    /// State-Offloading 時に System 2 へのエスカレーションプロンプト生成で活用される。
333    pub state_text: Option<String>,
334}
335
336impl CachedStateFeatures {
337    /// 非圧縮の生 `Value` からキャッシュエントリを生成する。
338    ///
339    /// # 引数
340    /// - `hash_key`: 64-bitハッシュキー。
341    /// - `shape`: テンソル形状。
342    /// - `ort_value`: ONNX Runtimeの出力 `Value`。
343    /// - `byte_size`: 占有バイト数。
344    /// - `state_mask`: `State` 有効トークンマスク列。空の場合は全トークン有効 (`1`) として補完される。
345    pub fn new_raw(
346        hash_key: u64,
347        shape: Vec<usize>,
348        ort_value: Value,
349        byte_size: usize,
350        state_mask: Vec<i64>,
351    ) -> Self {
352        Self::new_raw_with_text(hash_key, shape, ort_value, byte_size, state_mask, None)
353    }
354
355    /// 非圧縮の生 `Value` と State 原文テキストからキャッシュエントリを生成する。
356    pub fn new_raw_with_text(
357        hash_key: u64,
358        shape: Vec<usize>,
359        ort_value: Value,
360        byte_size: usize,
361        state_mask: Vec<i64>,
362        state_text: Option<String>,
363    ) -> Self {
364        let total_tokens = if shape.len() >= 2 {
365            shape[0] * shape[1]
366        } else {
367            0
368        };
369        let mut state_mask = state_mask;
370        if state_mask.is_empty() && total_tokens > 0 {
371            state_mask.resize(total_tokens, 1i64);
372        }
373
374        Self {
375            hash_key,
376            shape,
377            payload: CachedTensorPayload::Raw(Arc::new(ort_value)),
378            byte_size,
379            compression_mode: CacheCompressionMode::None,
380            state_mask,
381            state_text,
382        }
383    }
384
385    /// `f32` スライスを指定された圧縮モードで変換してキャッシュエントリを生成する。
386    ///
387    /// # 引数
388    /// - `hash_key`: 64-bitハッシュキー。
389    /// - `shape`: テンソル形状 (`[batch_size, seq_len, hidden_size]`)。
390    /// - `raw_slice`: 展開済み `f32` スライス。
391    /// - `mode`: 圧縮モード。
392    /// - `state_mask`: `State` 有効トークンマスク列。空の場合は全トークン有効 (`1`) として補完される。
393    pub fn new_compressed(
394        hash_key: u64,
395        shape: Vec<usize>,
396        raw_slice: &[f32],
397        mode: CacheCompressionMode,
398        state_mask: Vec<i64>,
399    ) -> Result<Self> {
400        Self::new_compressed_with_text(hash_key, shape, raw_slice, mode, state_mask, None)
401    }
402
403    /// `f32` スライスと State 原文テキストを指定された圧縮モードで変換してキャッシュエントリを生成する。
404    pub fn new_compressed_with_text(
405        hash_key: u64,
406        shape: Vec<usize>,
407        raw_slice: &[f32],
408        mode: CacheCompressionMode,
409        state_mask: Vec<i64>,
410        state_text: Option<String>,
411    ) -> Result<Self> {
412        let hidden_size = if shape.len() >= 3 {
413            shape[2]
414        } else {
415            return Err(RuntimeError::ModelContractViolation(format!(
416                "State特徴量は3次元テンソルである必要があります: {shape:?}。"
417            )));
418        };
419
420        let total_tokens = if shape.len() >= 2 {
421            shape[0] * shape[1]
422        } else {
423            0
424        };
425        let mut state_mask = state_mask;
426        if state_mask.is_empty() && total_tokens > 0 {
427            state_mask.resize(total_tokens, 1i64);
428        }
429
430        match mode {
431            CacheCompressionMode::None => {
432                let allocator = Allocator::default();
433                let mut tensor =
434                    Tensor::<f32>::new(&allocator, shape.clone()).map_err(RuntimeError::Ort)?;
435                let (_, dst) = tensor
436                    .try_extract_tensor_mut::<f32>()
437                    .map_err(RuntimeError::Ort)?;
438                dst.copy_from_slice(raw_slice);
439                let byte_size = std::mem::size_of_val(raw_slice);
440                Ok(Self {
441                    hash_key,
442                    shape,
443                    payload: CachedTensorPayload::Raw(Arc::new(tensor.into_dyn())),
444                    byte_size,
445                    compression_mode: mode,
446                    state_mask,
447                    state_text,
448                })
449            }
450            CacheCompressionMode::Fp16 => {
451                let compressed = compress_f32_to_fp16(raw_slice);
452                let byte_size = std::mem::size_of_val(compressed.as_slice());
453                Ok(Self {
454                    hash_key,
455                    shape,
456                    payload: CachedTensorPayload::Fp16(Arc::new(compressed)),
457                    byte_size,
458                    compression_mode: mode,
459                    state_mask,
460                    state_text,
461                })
462            }
463            CacheCompressionMode::Int8Symmetric => {
464                let (quantized, scales) = compress_f32_to_int8_symmetric(raw_slice, hidden_size);
465                let byte_size = quantized.len() * std::mem::size_of::<i8>()
466                    + scales.len() * std::mem::size_of::<f32>();
467                Ok(Self {
468                    hash_key,
469                    shape,
470                    payload: CachedTensorPayload::Int8Symmetric {
471                        quantized: Arc::new(quantized),
472                        scales: Arc::new(scales),
473                    },
474                    byte_size,
475                    compression_mode: mode,
476                    state_mask,
477                    state_text,
478                })
479            }
480        }
481    }
482
483    /// State 原文テキストを設定するビルダーメソッド。
484    pub fn with_state_text(mut self, text: Option<String>) -> Self {
485        self.state_text = text;
486        self
487    }
488
489    /// 推論入力としてバインド可能な `Arc<Value>` を取得する。
490    ///
491    /// 非圧縮時は保持している `Arc<Value>` をそのままゼロコピーで返却し、
492    /// 圧縮時はオンザフライで ONNX Runtime のアロケータメモリへ直接デコードを展開してテンソルを生成する。
493    pub fn as_ort_value(&self) -> Result<Arc<Value>> {
494        match &self.payload {
495            CachedTensorPayload::Raw(val) => Ok(Arc::clone(val)),
496            CachedTensorPayload::Fp16(compressed) => {
497                let allocator = Allocator::default();
498                let mut tensor = Tensor::<f32>::new(&allocator, self.shape.clone())
499                    .map_err(RuntimeError::Ort)?;
500                let (_, dst) = tensor
501                    .try_extract_tensor_mut::<f32>()
502                    .map_err(RuntimeError::Ort)?;
503                decompress_fp16_to_f32(compressed, dst);
504                Ok(Arc::new(tensor.into_dyn()))
505            }
506            CachedTensorPayload::Int8Symmetric { quantized, scales } => {
507                let hidden_size = self.hidden_size();
508                let allocator = Allocator::default();
509                let mut tensor = Tensor::<f32>::new(&allocator, self.shape.clone())
510                    .map_err(RuntimeError::Ort)?;
511                let (_, dst) = tensor
512                    .try_extract_tensor_mut::<f32>()
513                    .map_err(RuntimeError::Ort)?;
514                decompress_int8_symmetric_to_f32(quantized, scales, hidden_size, dst);
515                Ok(Arc::new(tensor.into_dyn()))
516            }
517        }
518    }
519
520    /// 指定されたバッチサイズに合わせてバインド可能な `Arc<Value>` を取得する。
521    ///
522    /// `target_batch == 1` かつ自身のバッチサイズが 1 の場合はゼロコピー返却する。
523    /// `target_batch > 1` の場合は、単一サンプルの特徴量データを `target_batch` 回リピートした
524    /// テンソル (`[target_batch, seq_len, hidden_size]`) を ONNX Runtime アロケータメモリへ直接構築して返却する。
525    /// 中間ベクタのヒープ確保と再アロケーションを回避し、64バイトアライメントを維持する。
526    pub fn as_ort_value_for_batch(&self, target_batch: usize) -> Result<Arc<Value>> {
527        if target_batch == 0 {
528            return Err(RuntimeError::InvalidTensorData(
529                "target_batch は 1 以上である必要があります。".to_string(),
530            ));
531        }
532
533        if target_batch == 1 && self.batch_size() == 1 {
534            return self.as_ort_value();
535        }
536
537        let seq_len = self.seq_len();
538        let hidden_size = self.hidden_size();
539        let single_len = seq_len * hidden_size;
540        let target_shape = vec![target_batch, seq_len, hidden_size];
541
542        let allocator = Allocator::default();
543        let mut tensor = Tensor::<f32>::new(&allocator, target_shape).map_err(RuntimeError::Ort)?;
544        let (_, dst) = tensor
545            .try_extract_tensor_mut::<f32>()
546            .map_err(RuntimeError::Ort)?;
547
548        match &self.payload {
549            CachedTensorPayload::Raw(val) => {
550                let (_, single_slice) =
551                    val.try_extract_tensor::<f32>().map_err(RuntimeError::Ort)?;
552                for chunk in dst.chunks_exact_mut(single_len) {
553                    chunk.copy_from_slice(single_slice);
554                }
555            }
556            CachedTensorPayload::Fp16(compressed) => {
557                let (first_chunk, rest) = dst.split_at_mut(single_len);
558                decompress_fp16_to_f32(compressed, first_chunk);
559                for chunk in rest.chunks_exact_mut(single_len) {
560                    chunk.copy_from_slice(first_chunk);
561                }
562            }
563            CachedTensorPayload::Int8Symmetric { quantized, scales } => {
564                let (first_chunk, rest) = dst.split_at_mut(single_len);
565                decompress_int8_symmetric_to_f32(quantized, scales, hidden_size, first_chunk);
566                for chunk in rest.chunks_exact_mut(single_len) {
567                    chunk.copy_from_slice(first_chunk);
568                }
569            }
570        }
571
572        Ok(Arc::new(tensor.into_dyn()))
573    }
574
575    /// バッチサイズを取得する。
576    pub fn batch_size(&self) -> usize {
577        self.shape.first().copied().unwrap_or(1)
578    }
579
580    /// 系列長（トークン数）を取得する。
581    pub fn seq_len(&self) -> usize {
582        if self.shape.len() >= 2 {
583            self.shape[1]
584        } else {
585            0
586        }
587    }
588
589    /// 隠れ層次元数を取得する。
590    pub fn hidden_size(&self) -> usize {
591        if self.shape.len() >= 3 {
592            self.shape[2]
593        } else {
594            0
595        }
596    }
597
598    /// 指定されたバッチサイズに合わせて複製された `state_mask` を取得する。
599    ///
600    /// キャッシュされた単一 `State` のマスク列 (`[state_seq_len]`) を `target_batch` 回リピートした
601    /// テンソル用フラットベクタ (`[target_batch * state_seq_len]`) を生成する。
602    pub fn state_mask_for_batch(&self, target_batch: usize) -> Vec<i64> {
603        let seq_len = self.seq_len();
604        let mut repeated = Vec::with_capacity(target_batch * seq_len);
605        let mask_slice = if self.state_mask.len() == seq_len {
606            &self.state_mask[..]
607        } else if self.state_mask.len() >= seq_len {
608            &self.state_mask[..seq_len]
609        } else {
610            &[]
611        };
612
613        if mask_slice.len() == seq_len {
614            for _ in 0..target_batch {
615                repeated.extend_from_slice(mask_slice);
616            }
617        } else {
618            let fallback = vec![1i64; seq_len];
619            for _ in 0..target_batch {
620                repeated.extend_from_slice(&fallback);
621            }
622        }
623        repeated
624    }
625}
626
627/// 並行 `State` 特徴量キャッシュマネージャー。
628///
629/// `xxHash3` による高速キー計算と、`moka` の TinyLFU/LRU 並行キャッシュにより、
630/// 複数スレッドからの同時参照をロックフリーに処理し、バイト重み付き退避を行う。
631pub struct StateCacheManager {
632    cache: Cache<u64, Arc<CachedStateFeatures>>,
633    compression_mode: CacheCompressionMode,
634}
635
636impl StateCacheManager {
637    /// 新規キャッシュマネージャーを初期化する。
638    ///
639    /// # 引数
640    /// - `max_capacity_bytes`: 最大常駐メモリ容量（バイト数）。
641    /// - `ttl_secs`: エントリの有効期限（秒数）。
642    /// - `compression_mode`: キャッシュ圧縮モード。
643    pub fn new(
644        max_capacity_bytes: u64,
645        ttl_secs: u64,
646        compression_mode: CacheCompressionMode,
647    ) -> Self {
648        let cache = Cache::builder()
649            .weigher(|_key, value: &Arc<CachedStateFeatures>| {
650                value.byte_size.min(u32::MAX as usize) as u32
651            })
652            .max_capacity(max_capacity_bytes)
653            .time_to_live(Duration::from_secs(ttl_secs))
654            .build();
655
656        Self {
657            cache,
658            compression_mode,
659        }
660    }
661
662    /// `State` テキストから `xxHash3` (64-bit) ハッシュキーを高速算出する。
663    ///
664    /// `AVX2`、`AVX-512`、`ARM NEON` 命令により、8kトークンのテキストでも1μs未満で完了する。
665    ///
666    /// # 引数
667    /// - `state_text`: `State` 原文テキスト。
668    pub fn compute_key(state_text: &str) -> u64 {
669        xxh3_64(state_text.as_bytes())
670    }
671
672    /// 指定されたキーの特徴量を取得する。
673    ///
674    /// # 引数
675    /// - `key`: `xxHash3` キー。
676    pub fn get(&self, key: u64) -> Option<Arc<CachedStateFeatures>> {
677        self.cache.get(&key)
678    }
679
680    /// 特徴量をキャッシュに登録する。
681    ///
682    /// # 引数
683    /// - `key`: `xxHash3` キー。
684    /// - `features`: キャッシュエントリ。
685    pub fn insert(&self, key: u64, features: Arc<CachedStateFeatures>) {
686        self.cache.insert(key, features);
687    }
688
689    /// 指定されたキーが存在するかを判定する。
690    pub fn contains_key(&self, key: u64) -> bool {
691        self.cache.contains_key(&key)
692    }
693
694    /// 指定されたキーのエントリを無効化（削除）する。
695    pub fn invalidate(&self, key: u64) {
696        self.cache.invalidate(&key);
697    }
698
699    /// 全キャッシュエントリを破棄する。
700    pub fn clear(&self) {
701        self.cache.invalidate_all();
702    }
703
704    /// 現在保持しているキャッシュエントリ数を取得する。
705    pub fn entry_count(&self) -> u64 {
706        self.cache.entry_count()
707    }
708
709    /// キャッシュ内の総メモリ占有バイト数を取得する。
710    pub fn weighted_size(&self) -> u64 {
711        self.cache.weighted_size()
712    }
713
714    /// 設定されている圧縮モードを取得する。
715    pub fn compression_mode(&self) -> CacheCompressionMode {
716        self.compression_mode
717    }
718
719    /// 保留中のキャッシュメンテナンスタスク（退避、カウント更新等）を直ちに同期実行する。
720    pub fn run_pending_tasks(&self) {
721        self.cache.run_pending_tasks();
722    }
723}
724
725/// 2セッション分割型推論エンジン (`DualSessionEngine`)。
726///
727/// `state_session` で `State` をエンコードし、得られた隠れ状態をキャッシュした上で、
728/// `query_session` の入力スロットへ `IoBinding` でゼロコピーバインドして超低遅延推論を実行する。
729///
730/// 内部にセッションプール（`query_sessions`、`state_sessions`）を保持し、
731/// 複数スレッドからの並行評価要求に対してロック競合を回避し、並列推論を実行する。
732pub struct DualSessionEngine {
733    state_sessions: Vec<Mutex<Session>>,
734    query_sessions: Vec<Mutex<Session>>,
735    next_state_idx: AtomicUsize,
736    next_query_idx: AtomicUsize,
737    cache_manager: Arc<StateCacheManager>,
738    active_provider: ExecutionProvider,
739}
740
741impl DualSessionEngine {
742    /// ファイルパスからデュアルモデルを読み込み、エンジンを初期化する。
743    ///
744    /// # 引数
745    /// - `state_model_path`: `state_encoder.onnx` ファイルパス。
746    /// - `query_model_path`: `query_evaluator.onnx` ファイルパス。
747    /// - `config`: セッション構築設定。
748    /// - `cache_manager`: キャッシュマネージャー。
749    pub fn new(
750        state_model_path: impl AsRef<Path>,
751        query_model_path: impl AsRef<Path>,
752        config: SessionConfig,
753        cache_manager: Arc<StateCacheManager>,
754    ) -> Result<Self> {
755        let pool_size = config.pool_size.max(1);
756        let mut state_sessions = Vec::with_capacity(pool_size);
757        let mut query_sessions = Vec::with_capacity(pool_size);
758        let mut active_provider = ExecutionProvider::CPU;
759
760        for i in 0..pool_size {
761            let (mut state_builder, provider) = create_session_builder(&config, pool_size)?;
762            if i == 0 {
763                active_provider = provider;
764            }
765            let s_session = state_builder
766                .commit_from_file(state_model_path.as_ref())
767                .map_err(RuntimeError::Ort)?;
768            state_sessions.push(Mutex::new(s_session));
769
770            let (mut query_builder, _) = create_session_builder(&config, pool_size)?;
771            let q_session = query_builder
772                .commit_from_file(query_model_path.as_ref())
773                .map_err(RuntimeError::Ort)?;
774            query_sessions.push(Mutex::new(q_session));
775        }
776
777        Ok(Self {
778            state_sessions,
779            query_sessions,
780            next_state_idx: AtomicUsize::new(0),
781            next_query_idx: AtomicUsize::new(0),
782            cache_manager,
783            active_provider,
784        })
785    }
786
787    /// メモリ上のバイト列からデュアルモデルを読み込み、エンジンを初期化する。
788    ///
789    /// # 引数
790    /// - `state_model_bytes`: `state_encoder.onnx` バイナリ。
791    /// - `query_model_bytes`: `query_evaluator.onnx` バイナリ。
792    /// - `config`: セッション構築設定。
793    /// - `cache_manager`: キャッシュマネージャー。
794    pub fn from_bytes(
795        state_model_bytes: &[u8],
796        query_model_bytes: &[u8],
797        config: SessionConfig,
798        cache_manager: Arc<StateCacheManager>,
799    ) -> Result<Self> {
800        let pool_size = config.pool_size.max(1);
801        let mut state_sessions = Vec::with_capacity(pool_size);
802        let mut query_sessions = Vec::with_capacity(pool_size);
803        let mut active_provider = ExecutionProvider::CPU;
804
805        for i in 0..pool_size {
806            let (mut state_builder, provider) = create_session_builder(&config, pool_size)?;
807            if i == 0 {
808                active_provider = provider;
809            }
810            let s_session = state_builder
811                .commit_from_memory(state_model_bytes)
812                .map_err(RuntimeError::Ort)?;
813            state_sessions.push(Mutex::new(s_session));
814
815            let (mut query_builder, _) = create_session_builder(&config, pool_size)?;
816            let q_session = query_builder
817                .commit_from_memory(query_model_bytes)
818                .map_err(RuntimeError::Ort)?;
819            query_sessions.push(Mutex::new(q_session));
820        }
821
822        Ok(Self {
823            state_sessions,
824            query_sessions,
825            next_state_idx: AtomicUsize::new(0),
826            next_query_idx: AtomicUsize::new(0),
827            cache_manager,
828            active_provider,
829        })
830    }
831
832    /// 構築済みのセッションからエンジンインスタンスを生成する。
833    pub fn from_sessions(
834        state_session: Session,
835        query_session: Session,
836        active_provider: ExecutionProvider,
837        cache_manager: Arc<StateCacheManager>,
838    ) -> Self {
839        Self {
840            state_sessions: vec![Mutex::new(state_session)],
841            query_sessions: vec![Mutex::new(query_session)],
842            next_state_idx: AtomicUsize::new(0),
843            next_query_idx: AtomicUsize::new(0),
844            cache_manager,
845            active_provider,
846        }
847    }
848
849    /// セッションプールから利用可能な `Query` セッションを取得する。
850    ///
851    /// ロックされていない空きセッションを優先探索し、全セッションが使用中の場合はブロック待機する。
852    fn acquire_query_session(&self) -> Result<std::sync::MutexGuard<'_, Session>> {
853        let pool_len = self.query_sessions.len();
854        let start_idx = self.next_query_idx.fetch_add(1, Ordering::Relaxed) % pool_len;
855        for offset in 0..pool_len {
856            let idx = (start_idx + offset) % pool_len;
857            if let Ok(guard) = self.query_sessions[idx].try_lock() {
858                return Ok(guard);
859            }
860        }
861        self.query_sessions[start_idx].lock().map_err(|e| {
862            RuntimeError::ModelContractViolation(format!("Query セッションロック獲得失敗: {e}。"))
863        })
864    }
865
866    /// セッションプールから利用可能な `State` セッションを取得する。
867    ///
868    /// ロックされていない空きセッションを優先探索し、全セッションが使用中の場合はブロック待機する。
869    fn acquire_state_session(&self) -> Result<std::sync::MutexGuard<'_, Session>> {
870        let pool_len = self.state_sessions.len();
871        let start_idx = self.next_state_idx.fetch_add(1, Ordering::Relaxed) % pool_len;
872        for offset in 0..pool_len {
873            let idx = (start_idx + offset) % pool_len;
874            if let Ok(guard) = self.state_sessions[idx].try_lock() {
875                return Ok(guard);
876            }
877        }
878        self.state_sessions[start_idx].lock().map_err(|e| {
879            RuntimeError::ModelContractViolation(format!("State セッションロック獲得失敗: {e}。"))
880        })
881    }
882
883    /// 保持しているキャッシュマネージャーへの参照を取得する。
884    pub fn cache_manager(&self) -> &Arc<StateCacheManager> {
885        &self.cache_manager
886    }
887
888    /// 稼働中の Execution Provider を取得する。
889    pub fn active_provider(&self) -> ExecutionProvider {
890        self.active_provider
891    }
892
893    /// `State` をエンコードしてキャッシュへ格納し、キャッシュエントリを返す。
894    ///
895    /// 既にキャッシュに存在する場合は推論を省略し、即座にキャッシュエントリを返す。
896    ///
897    /// # 引数
898    /// - `state_text`: キャッシュキー生成用の `State` 原文テキスト。
899    /// - `input_ids`: トークンID列 (`[batch_size, seq_len]`)。
900    /// - `attention_mask`: アテンションマスク列 (`[batch_size, seq_len]`)。
901    /// - `dims`: `[batch_size, seq_len]` の次元配列。
902    pub fn encode_and_cache_state(
903        &self,
904        state_text: &str,
905        input_ids: &[i64],
906        attention_mask: &[i64],
907        dims: [usize; 2],
908    ) -> Result<Arc<CachedStateFeatures>> {
909        let key = StateCacheManager::compute_key(state_text);
910        if let Some(cached) = self.cache_manager.get(key) {
911            return Ok(cached);
912        }
913
914        let [batch_size, seq_len] = dims;
915        if input_ids.len() != batch_size * seq_len {
916            return Err(RuntimeError::InvalidTensorData(format!(
917                "input_ids 長さ ({}) が期待値 ({batch_size} * {seq_len} = {}) と一致しません。",
918                input_ids.len(),
919                batch_size * seq_len
920            )));
921        }
922        if attention_mask.len() != batch_size * seq_len {
923            return Err(RuntimeError::InvalidTensorData(format!(
924                "attention_mask 長さ ({}) が期待値 ({batch_size} * {seq_len} = {}) と一致しません。",
925                attention_mask.len(),
926                batch_size * seq_len
927            )));
928        }
929
930        let mut session = self.acquire_state_session()?;
931
932        let mut binding = session.create_binding().map_err(RuntimeError::Ort)?;
933        let mem_info = MemoryInfo::default();
934
935        let ids_val = Value::from_array(([batch_size, seq_len], input_ids.to_vec()))
936            .map_err(RuntimeError::Ort)?;
937        let mask_val = Value::from_array(([batch_size, seq_len], attention_mask.to_vec()))
938            .map_err(RuntimeError::Ort)?;
939
940        binding
941            .bind_input(TENSOR_INPUT_IDS, &ids_val)
942            .map_err(RuntimeError::Ort)?;
943        binding
944            .bind_input(TENSOR_ATTENTION_MASK, &mask_val)
945            .map_err(RuntimeError::Ort)?;
946        binding
947            .bind_output_to_device(TENSOR_STATE_HIDDEN_STATES, &mem_info)
948            .map_err(RuntimeError::Ort)?;
949
950        let mut outputs = session.run_binding(&binding).map_err(RuntimeError::Ort)?;
951        let state_val = outputs.remove(TENSOR_STATE_HIDDEN_STATES).ok_or_else(|| {
952            RuntimeError::ModelContractViolation(format!(
953                "出力テンソル '{TENSOR_STATE_HIDDEN_STATES}' が存在しません。"
954            ))
955        })?;
956
957        let (shape, flat_data) = state_val
958            .try_extract_tensor::<f32>()
959            .map_err(RuntimeError::Ort)?;
960        let shape_vec: Vec<usize> = shape.iter().map(|&x| x as usize).collect();
961
962        let compression = self.cache_manager.compression_mode();
963        let state_text_owned = Some(state_text.to_string());
964        let cached_entry = if compression == CacheCompressionMode::None {
965            let total_elements: usize = shape_vec.iter().product();
966            let byte_size = total_elements * std::mem::size_of::<f32>();
967            Arc::new(CachedStateFeatures::new_raw_with_text(
968                key,
969                shape_vec,
970                state_val,
971                byte_size,
972                attention_mask.to_vec(),
973                state_text_owned,
974            ))
975        } else {
976            Arc::new(CachedStateFeatures::new_compressed_with_text(
977                key,
978                shape_vec,
979                flat_data,
980                compression,
981                attention_mask.to_vec(),
982                state_text_owned,
983            )?)
984        };
985
986        self.cache_manager.insert(key, Arc::clone(&cached_entry));
987        Ok(cached_entry)
988    }
989
990    /// キャッシュされた `State` 特徴量をゼロコピーバインドして質問を推論する。
991    ///
992    /// # 引数
993    /// - `cached_state`: キャッシュされた `State` 特徴量。
994    /// - `state_mask`: `State` マスク列 (`[batch_size, state_seq_len]`)。
995    /// - `query_ids`: 質問トークンID列 (`[batch_size, query_seq_len]`)。
996    /// - `query_mask`: 質問アテンションマスク列 (`[batch_size, query_seq_len]`)。
997    /// - `op_indices`: 候補アンカーインデックス (`[batch_size, num_options]`)。
998    /// - `query_dims`: 質問側の入力次元情報。
999    ///
1000    /// # 戻り値
1001    /// 各バッチサンプルのロジット配列 (`Vec<Vec<f64>>`)。
1002    pub fn evaluate_query_with_cache(
1003        &self,
1004        cached_state: &CachedStateFeatures,
1005        state_mask: &[i64],
1006        query_ids: &[i64],
1007        query_mask: &[i64],
1008        op_indices: &[i64],
1009        query_dims: ModelInputDimensions,
1010    ) -> Result<Vec<Vec<f64>>> {
1011        let state_val = cached_state.as_ort_value_for_batch(query_dims.batch_size)?;
1012        self.evaluate_query_with_state_value(
1013            state_val.as_ref(),
1014            cached_state.seq_len(),
1015            state_mask,
1016            query_ids,
1017            query_mask,
1018            op_indices,
1019            query_dims,
1020        )
1021    }
1022
1023    /// 構築済みの `State` テンソル `Value` を直接バインドして質問を推論する。
1024    ///
1025    /// 複数チャンクのバッチ推論において、同一チャンク長の `state_val` を再利用して
1026    /// リピートテンソルの再生成アロケーションをスキップするために使用する。
1027    #[allow(clippy::too_many_arguments)]
1028    pub fn evaluate_query_with_state_value(
1029        &self,
1030        state_val: &Value,
1031        state_seq_len: usize,
1032        state_mask: &[i64],
1033        query_ids: &[i64],
1034        query_mask: &[i64],
1035        op_indices: &[i64],
1036        query_dims: ModelInputDimensions,
1037    ) -> Result<Vec<Vec<f64>>> {
1038        let batch_size = query_dims.batch_size;
1039        let query_seq_len = query_dims.sequence_length;
1040        let num_options = query_dims.num_options;
1041
1042        // 1. 動的軸の事前検証 (Fast-Fail)。
1043        if state_mask.len() != batch_size * state_seq_len {
1044            return Err(RuntimeError::CacheDimensionMismatch {
1045                expected: batch_size * state_seq_len,
1046                actual: state_mask.len(),
1047            });
1048        }
1049        if query_ids.len() != batch_size * query_seq_len {
1050            return Err(RuntimeError::InvalidTensorData(format!(
1051                "query_ids 長さ ({}) が期待値 ({batch_size} * {query_seq_len} = {}) と一致しません。",
1052                query_ids.len(),
1053                batch_size * query_seq_len
1054            )));
1055        }
1056        if query_mask.len() != batch_size * query_seq_len {
1057            return Err(RuntimeError::InvalidTensorData(format!(
1058                "query_mask 長さ ({}) が期待値 ({batch_size} * {query_seq_len} = {}) と一致しません。",
1059                query_mask.len(),
1060                batch_size * query_seq_len
1061            )));
1062        }
1063        if op_indices.len() != batch_size * num_options {
1064            return Err(RuntimeError::InvalidTensorData(format!(
1065                "op_indices 長さ ({}) が期待値 ({batch_size} * {num_options} = {}) と一致しません。",
1066                op_indices.len(),
1067                batch_size * num_options
1068            )));
1069        }
1070
1071        // 2. `op_indices` の範囲外検証（クエリ系列長基準、パディング-1は許容）。
1072        for (i, &idx) in op_indices.iter().enumerate() {
1073            if idx != -1 && (idx < 0 || idx >= query_seq_len as i64) {
1074                return Err(RuntimeError::InvalidTensorData(format!(
1075                    "op_indices[{i}] = {idx} が有効範囲 (-1 または 0..{query_seq_len}) の範囲外です。"
1076                )));
1077            }
1078        }
1079
1080        let mut session = self.acquire_query_session()?;
1081
1082        let mut binding = session.create_binding().map_err(RuntimeError::Ort)?;
1083        let mem_info = MemoryInfo::default();
1084
1085        // 3. キャッシュ特徴量のゼロコピー / バッチ適合バインド。
1086        binding
1087            .bind_input(TENSOR_STATE_HIDDEN_STATES, state_val)
1088            .map_err(RuntimeError::Ort)?;
1089
1090        // 4. `Query` 側入力テンソルのバインド。
1091        let s_mask_val = Value::from_array(([batch_size, state_seq_len], state_mask.to_vec()))
1092            .map_err(RuntimeError::Ort)?;
1093        let q_ids_val = Value::from_array(([batch_size, query_seq_len], query_ids.to_vec()))
1094            .map_err(RuntimeError::Ort)?;
1095        let q_mask_val = Value::from_array(([batch_size, query_seq_len], query_mask.to_vec()))
1096            .map_err(RuntimeError::Ort)?;
1097        let op_idx_val = Value::from_array(([batch_size, num_options], op_indices.to_vec()))
1098            .map_err(RuntimeError::Ort)?;
1099
1100        binding
1101            .bind_input(TENSOR_STATE_MASK, &s_mask_val)
1102            .map_err(RuntimeError::Ort)?;
1103        binding
1104            .bind_input(TENSOR_QUERY_IDS, &q_ids_val)
1105            .map_err(RuntimeError::Ort)?;
1106        binding
1107            .bind_input(TENSOR_QUERY_MASK, &q_mask_val)
1108            .map_err(RuntimeError::Ort)?;
1109        binding
1110            .bind_input(TENSOR_OP_INDICES, &op_idx_val)
1111            .map_err(RuntimeError::Ort)?;
1112
1113        binding
1114            .bind_output_to_device(TENSOR_LOGITS, &mem_info)
1115            .map_err(RuntimeError::Ort)?;
1116
1117        // 5. 推論実行。
1118        let outputs = session.run_binding(&binding).map_err(RuntimeError::Ort)?;
1119        let logits_out = outputs.get(TENSOR_LOGITS).ok_or_else(|| {
1120            RuntimeError::ModelContractViolation(format!(
1121                "出力テンソル '{TENSOR_LOGITS}' が存在しません。"
1122            ))
1123        })?;
1124
1125        let (shape, logits_f32) = logits_out
1126            .try_extract_tensor::<f32>()
1127            .map_err(RuntimeError::Ort)?;
1128
1129        if shape.len() != 2 {
1130            return Err(RuntimeError::ModelContractViolation(format!(
1131                "出力テンソル '{TENSOR_LOGITS}' は 2 次元テンソル [batch, num_options] である必要があります: {shape:?}。"
1132            )));
1133        }
1134        let shape_stack = [shape[0] as usize, shape[1] as usize];
1135        query_dims.validate_output_shape(&shape_stack)?;
1136
1137        // 6. 出力値の有限性 (`NaN`/`Inf`) 検証および `f32` から `f64` へのキャスト。
1138        let mut batched_logits = Vec::with_capacity(batch_size);
1139        for chunk in logits_f32.chunks_exact(num_options) {
1140            let mut sample_logits = Vec::with_capacity(num_options);
1141            for &val in chunk {
1142                if !val.is_finite() {
1143                    return Err(RuntimeError::InvalidTensorData(format!(
1144                        "logits に非有限値 (NaN または Inf) が検出されました: {val}。"
1145                    )));
1146                }
1147                sample_logits.push(val as f64);
1148            }
1149            batched_logits.push(sample_logits);
1150        }
1151
1152        Ok(batched_logits)
1153    }
1154
1155    /// `Tokio` ワーカースレッドの枯渇を防ぐため、`tokio::task::spawn_blocking` で隔離して推論を実行する。
1156    ///
1157    /// # 引数
1158    /// - `engine`: デュアルセッションエンジン。
1159    /// - `cached_state`: キャッシュされた `State` 特徴量。
1160    /// - `state_mask`: `State` マスク列。
1161    /// - `query_ids`: 質問トークンID列。
1162    /// - `query_mask`: 質問アテンションマスク列。
1163    /// - `op_indices`: 候補アンカーインデックス。
1164    /// - `query_dims`: 質問側の入力次元情報。
1165    pub fn evaluate_query_async(
1166        engine: Arc<Self>,
1167        cached_state: Arc<CachedStateFeatures>,
1168        state_mask: Vec<i64>,
1169        query_ids: Vec<i64>,
1170        query_mask: Vec<i64>,
1171        op_indices: Vec<i64>,
1172        query_dims: ModelInputDimensions,
1173    ) -> tokio::task::JoinHandle<Result<Vec<Vec<f64>>>> {
1174        tokio::task::spawn_blocking(move || {
1175            engine.evaluate_query_with_cache(
1176                &cached_state,
1177                &state_mask,
1178                &query_ids,
1179                &query_mask,
1180                &op_indices,
1181                query_dims,
1182            )
1183        })
1184    }
1185
1186    /// キャッシュされた `State` 特徴量をゼロコピー参照し、複数の質問を一括またはマイクロバッチ評価する。
1187    ///
1188    /// 内部で各質問の `Query` 部分を事前トークナイズし、`chunk_size` 単位のマイクロバッチとして
1189    /// `evaluate_query_with_cache` を実行することで、L3 キャッシュ保護と並列推論を両立する。
1190    ///
1191    /// # 引数
1192    /// - `tokenizer`: 高速トークナイザー。
1193    /// - `cached_state`: キャッシュされた `State` 特徴量。
1194    /// - `state_text`: `State` テキスト (ゲーティングや OOD 判定用、任意)。
1195    /// - `questions`: 評価対象質問群のマップ。
1196    /// - `calib_config`: 事後較正設定。
1197    /// - `gating_config`: 確信度ゲーティング設定 (任意)。
1198    /// - `chunk_size`: マイクロバッチ分割サイズ (OOM 防止)。
1199    ///
1200    /// # 戻り値
1201    /// 各質問 ID に対応する決定結果 (`IndexMap<String, Answer>`)。
1202    #[allow(clippy::too_many_arguments)]
1203    pub fn evaluate_batch_questions_with_cache(
1204        &self,
1205        tokenizer: &JevTokenizer,
1206        cached_state: &CachedStateFeatures,
1207        state_text: Option<&str>,
1208        questions: &indexmap::IndexMap<String, Question>,
1209        calib_config: &CalibrationConfig,
1210        gating_config: Option<&GatingConfig>,
1211        chunk_size: usize,
1212    ) -> Result<indexmap::IndexMap<String, Answer>> {
1213        if questions.is_empty() {
1214            return Ok(indexmap::IndexMap::new());
1215        }
1216
1217        let state_seq_len = cached_state.seq_len();
1218        let single_sample_bytes =
1219            state_seq_len * cached_state.hidden_size() * std::mem::size_of::<f32>();
1220
1221        let max_safe_chunk = if single_sample_bytes > 12 * 1024 * 1024 {
1222            2
1223        } else if single_sample_bytes > 4 * 1024 * 1024 {
1224            DEFAULT_DUAL_MAX_BATCH_CHUNK_SIZE
1225        } else {
1226            DEFAULT_DUAL_MAX_BATCH_CHUNK_SIZE * 2
1227        };
1228
1229        let effective_chunk_size = chunk_size.max(1).min(max_safe_chunk);
1230        let mut answers = indexmap::IndexMap::with_capacity(questions.len());
1231        let pad_id = tokenizer.pad_token_id().unwrap_or(0) as i64;
1232
1233        let mut tokenized_items = Vec::with_capacity(questions.len());
1234        for (qid, q) in questions {
1235            let t = tokenizer.encode_query(q)?;
1236            tokenized_items.push((qid, q, t));
1237        }
1238
1239        let mut scratch_probs = Vec::new();
1240        let mut reusable_state_val: Option<(usize, Arc<Value>)> = None;
1241
1242        for chunk in tokenized_items.chunks(effective_chunk_size) {
1243            let chunk_len = chunk.len();
1244            let max_q_len = chunk
1245                .iter()
1246                .map(|(_, _, t)| t.input_ids.len())
1247                .max()
1248                .unwrap_or(1);
1249            let max_opts = chunk
1250                .iter()
1251                .map(|(_, _, t)| t.op_indices.len())
1252                .max()
1253                .unwrap_or(1);
1254
1255            let mut batch_query_ids = vec![pad_id; chunk_len * max_q_len];
1256            let mut batch_query_mask = vec![0i64; chunk_len * max_q_len];
1257            let mut batch_op_indices = vec![-1i64; chunk_len * max_opts];
1258            let batch_state_mask = cached_state.state_mask_for_batch(chunk_len);
1259
1260            for (i, (_, _, t)) in chunk.iter().enumerate() {
1261                let q_len = t.input_ids.len();
1262                let q_start = i * max_q_len;
1263                batch_query_ids[q_start..q_start + q_len].copy_from_slice(&t.input_ids);
1264                batch_query_mask[q_start..q_start + q_len].copy_from_slice(&t.attention_mask);
1265
1266                let opt_len = t.op_indices.len();
1267                let opt_start = i * max_opts;
1268                batch_op_indices[opt_start..opt_start + opt_len].copy_from_slice(&t.op_indices);
1269            }
1270
1271            let query_dims = ModelInputDimensions {
1272                batch_size: chunk_len,
1273                sequence_length: max_q_len,
1274                num_options: max_opts,
1275            };
1276
1277            let state_val = match &reusable_state_val {
1278                Some((len, val)) if *len == chunk_len => Arc::clone(val),
1279                _ => {
1280                    let val = cached_state.as_ort_value_for_batch(chunk_len)?;
1281                    reusable_state_val = Some((chunk_len, Arc::clone(&val)));
1282                    val
1283                }
1284            };
1285
1286            let logits_per_sample = self.evaluate_query_with_state_value(
1287                state_val.as_ref(),
1288                state_seq_len,
1289                &batch_state_mask,
1290                &batch_query_ids,
1291                &batch_query_mask,
1292                &batch_op_indices,
1293                query_dims,
1294            )?;
1295
1296            for (i, (qid, q, t)) in chunk.iter().enumerate() {
1297                let k_i = t.op_indices.len();
1298                let sample_logits = &logits_per_sample[i][..k_i];
1299
1300                if scratch_probs.len() < k_i {
1301                    scratch_probs.resize(k_i, 0.0);
1302                }
1303
1304                let mut answer = evaluate_question_with_buf(
1305                    q,
1306                    sample_logits,
1307                    calib_config,
1308                    &mut scratch_probs[..k_i],
1309                )?;
1310
1311                let energy = calculate_choice_energy(q, sample_logits, calib_config);
1312                let effective_state_text = match state_text {
1313                    Some(s) if !s.trim().is_empty() => Some(s),
1314                    _ => cached_state.state_text.as_deref(),
1315                };
1316                apply_gating_to_answer_with_energy(
1317                    &mut answer,
1318                    energy,
1319                    effective_state_text,
1320                    Some(qid.as_str()),
1321                    Some(q),
1322                    gating_config,
1323                    calib_config,
1324                );
1325
1326                answers.insert((*qid).clone(), answer);
1327            }
1328        }
1329
1330        Ok(answers)
1331    }
1332
1333    /// 単一の質問クエリに対する各選択肢の生ロジット列 (`Vec<f64>`) を抽出する。
1334    ///
1335    /// # 引数
1336    /// - `cached_state`: エンコード済みの中間特徴量キャッシュ。
1337    /// - `tokenized`: トークナイズ済みの質問クエリ。
1338    ///
1339    /// # 戻り値
1340    /// 該当質問の各選択肢 (`[OP]`) に対応する生ロジット列。
1341    pub fn forward_query_raw(
1342        &self,
1343        cached_state: &CachedStateFeatures,
1344        tokenized: &TokenizedQuestion,
1345    ) -> Result<Vec<f64>> {
1346        let batch_state_mask = cached_state.state_mask_for_batch(1);
1347        let query_dims = ModelInputDimensions {
1348            batch_size: 1,
1349            sequence_length: tokenized.input_ids.len(),
1350            num_options: tokenized.op_indices.len(),
1351        };
1352        let mut logits = self.evaluate_query_with_cache(
1353            cached_state,
1354            &batch_state_mask,
1355            &tokenized.input_ids,
1356            &tokenized.attention_mask,
1357            &tokenized.op_indices,
1358            query_dims,
1359        )?;
1360        let sample_logits = logits.pop().unwrap_or_default();
1361        let k = tokenized.op_indices.len();
1362        Ok(sample_logits.into_iter().take(k).collect())
1363    }
1364
1365    /// 単一の質問定義とキャッシュ済み `State` から各選択肢の未正規化生ロジット列 (`Vec<f64>`) を直接抽出する。
1366    ///
1367    /// # 引数
1368    /// - `tokenizer`: Jev 高速トークナイザー。
1369    /// - `cached_state`: エンコード済みの中間特徴量キャッシュ。
1370    /// - `question`: 質問仕様。
1371    ///
1372    /// # 戻り値
1373    /// 各選択肢の未正規化生ロジット列。
1374    pub fn forward_question_raw(
1375        &self,
1376        tokenizer: &JevTokenizer,
1377        cached_state: &CachedStateFeatures,
1378        question: &Question,
1379    ) -> Result<Vec<f64>> {
1380        let tokenized = tokenizer.encode_query(question)?;
1381        self.forward_query_raw(cached_state, &tokenized)
1382    }
1383
1384    /// 単一の質問をキャッシュ済み `State` に対して推論し、決定結果 (`Answer`) を返却する。
1385    ///
1386    /// # 引数
1387    /// - `tokenizer`: Jev 高速トークナイザー。
1388    /// - `cached_state`: エンコード済みの中間特徴量キャッシュ。
1389    /// - `state_text`: 任意で渡す文脈テキスト (ゲーティングメタデータ用)。
1390    /// - `question`: 質問仕様。
1391    /// - `calib_config`: 較正温度設定。
1392    ///
1393    /// # 戻り値
1394    /// 確率較正済みの `Answer`。
1395    pub fn predict_single_question(
1396        &self,
1397        tokenizer: &JevTokenizer,
1398        cached_state: &CachedStateFeatures,
1399        state_text: Option<&str>,
1400        question: &Question,
1401        calib_config: &CalibrationConfig,
1402    ) -> Result<Answer> {
1403        let tokenized = tokenizer.encode_query(question)?;
1404        let logits = self.forward_query_raw(cached_state, &tokenized)?;
1405        let mut answer =
1406            evaluate_question(question, &logits, calib_config).map_err(RuntimeError::Core)?;
1407
1408        let energy = calculate_choice_energy(question, &logits, calib_config);
1409        let effective_state_text = match state_text {
1410            Some(s) if !s.trim().is_empty() => Some(s),
1411            _ => cached_state.state_text.as_deref(),
1412        };
1413        apply_gating_to_answer_with_energy(
1414            &mut answer,
1415            energy,
1416            effective_state_text,
1417            None,
1418            Some(question),
1419            None,
1420            calib_config,
1421        );
1422
1423        Ok(answer)
1424    }
1425}
1426
1427#[cfg(test)]
1428mod tests {
1429    use super::*;
1430
1431    #[test]
1432    fn test_xxh3_key_computation_deterministic() {
1433        let text1 = "本契約において、甲は乙に対し対価を支払うものとする。";
1434        let text2 = "本契約において、甲は乙に対し対価を支払うものとする。";
1435        let text3 = "別のテキストです。";
1436
1437        let key1 = StateCacheManager::compute_key(text1);
1438        let key2 = StateCacheManager::compute_key(text2);
1439        let key3 = StateCacheManager::compute_key(text3);
1440
1441        assert_eq!(key1, key2);
1442        assert_ne!(key1, key3);
1443        assert_ne!(key1, 0);
1444    }
1445
1446    #[test]
1447    fn test_aligned_buffer64_alignment() {
1448        let buf = AlignedBuffer64::<f32>::with_capacity_zeroed(128);
1449        assert_eq!(buf.len(), 128);
1450        assert!(buf.is_aligned_64());
1451        assert!(is_aligned_64(buf.as_ptr()));
1452
1453        for &val in buf.iter() {
1454            assert_eq!(val, 0.0);
1455        }
1456
1457        let slice = vec![1.5f32, 2.5, 3.5, 4.5];
1458        let buf_from_slice = AlignedBuffer64::from_slice(&slice);
1459        assert_eq!(buf_from_slice.len(), 4);
1460        assert!(buf_from_slice.is_aligned_64());
1461        assert_eq!(&buf_from_slice[..], &slice[..]);
1462    }
1463
1464    #[test]
1465    fn test_fp16_compression_roundtrip_precision() {
1466        let original: Vec<f32> = (0..1000).map(|i| (i as f32) * 0.005 - 2.5).collect();
1467        let compressed = compress_f32_to_fp16(&original);
1468        assert_eq!(compressed.len(), original.len());
1469
1470        let mut restored = vec![0.0f32; original.len()];
1471        decompress_fp16_to_f32(&compressed, &mut restored);
1472
1473        let mut max_err = 0.0f32;
1474        for (&orig, &rest) in original.iter().zip(restored.iter()) {
1475            let err = (orig - rest).abs();
1476            if err > max_err {
1477                max_err = err;
1478            }
1479        }
1480
1481        // `FP16` の最大絶対誤差は 2e-3 未満（正規化範囲内で高精度）。
1482        assert!(max_err < 2e-3, "FP16 誤差が大きすぎます: {max_err}");
1483    }
1484
1485    #[test]
1486    fn test_int8_symmetric_compression_roundtrip() {
1487        let hidden_size = 64;
1488        let num_tokens = 4;
1489        let mut original = Vec::with_capacity(num_tokens * hidden_size);
1490
1491        for t in 0..num_tokens {
1492            let scale = (t + 1) as f32 * 0.8;
1493            for i in 0..hidden_size {
1494                original.push(((i as f32) / (hidden_size as f32) - 0.5) * scale);
1495            }
1496        }
1497
1498        let (quantized, scales) = compress_f32_to_int8_symmetric(&original, hidden_size);
1499        assert_eq!(quantized.len(), original.len());
1500        assert_eq!(scales.len(), num_tokens);
1501
1502        let mut restored = vec![0.0f32; original.len()];
1503        decompress_int8_symmetric_to_f32(&quantized, &scales, hidden_size, &mut restored);
1504
1505        let mut max_err = 0.0f32;
1506        for (&orig, &rest) in original.iter().zip(restored.iter()) {
1507            let err = (orig - rest).abs();
1508            if err > max_err {
1509                max_err = err;
1510            }
1511        }
1512
1513        // `INT8` 対称量子化の量子化ステップ幅は `scale / 127`。最大誤差が 0.05 未満であることを確認。
1514        assert!(max_err < 0.05, "INT8 誤差が大きすぎます: {max_err}");
1515    }
1516
1517    #[test]
1518    fn test_state_cache_manager_crud_and_weigher() {
1519        let manager = StateCacheManager::new(
1520            1024 * 1024, // 1MB
1521            60,
1522            CacheCompressionMode::Fp16,
1523        );
1524
1525        let text = "特許権侵害訴訟における証拠保全手続きの概要。";
1526        let key = StateCacheManager::compute_key(text);
1527        assert!(!manager.contains_key(key));
1528        assert!(manager.get(key).is_none());
1529
1530        let shape = vec![1, 8, 16];
1531        let raw_data = vec![0.1f32; 8 * 16];
1532        let state_mask = vec![1i64; 8];
1533        let features = Arc::new(
1534            CachedStateFeatures::new_compressed(
1535                key,
1536                shape,
1537                &raw_data,
1538                CacheCompressionMode::Fp16,
1539                state_mask,
1540            )
1541            .unwrap(),
1542        );
1543
1544        manager.insert(key, Arc::clone(&features));
1545        assert!(manager.contains_key(key));
1546        manager.run_pending_tasks();
1547
1548        let retrieved = manager.get(key).expect("キャッシュヒットするはずです。");
1549        assert_eq!(retrieved.hash_key, key);
1550        assert_eq!(retrieved.seq_len(), 8);
1551        assert_eq!(retrieved.hidden_size(), 16);
1552        assert_eq!(retrieved.state_mask, vec![1i64; 8]);
1553        assert!(manager.weighted_size() > 0);
1554
1555        manager.invalidate(key);
1556        manager.run_pending_tasks();
1557        assert!(!manager.contains_key(key));
1558    }
1559
1560    #[test]
1561    fn test_as_ort_value_for_batch_repeat() {
1562        let text = "バッチリピートテスト用のStateテキスト。";
1563        let key = StateCacheManager::compute_key(text);
1564        let seq_len = 4;
1565        let hidden_size = 8;
1566        let shape = vec![1, seq_len, hidden_size];
1567        let raw_data: Vec<f32> = (0..seq_len * hidden_size).map(|i| i as f32).collect();
1568        // パディングを含むマスク列 (先頭3トークン有効、4トークン目パディング)
1569        let custom_state_mask = vec![1i64, 1, 1, 0];
1570
1571        let features = CachedStateFeatures::new_compressed(
1572            key,
1573            shape,
1574            &raw_data,
1575            CacheCompressionMode::Fp16,
1576            custom_state_mask.clone(),
1577        )
1578        .expect("キャッシュエントリの作成に成功すること。");
1579
1580        // 単一バッチ時は [1, seq_len, hidden_size]
1581        let val_b1 = features
1582            .as_ort_value_for_batch(1)
1583            .expect("単一バッチ取得成功。");
1584        let (s1, d1) = val_b1
1585            .try_extract_tensor::<f32>()
1586            .expect("テンソル抽出成功。");
1587        assert_eq!(s1.as_ref(), &[1, seq_len as i64, hidden_size as i64]);
1588        assert_eq!(d1.len(), seq_len * hidden_size);
1589
1590        // 複数バッチ (例: batch=3) 時は [3, seq_len, hidden_size] へリピート展開
1591        let val_b3 = features
1592            .as_ort_value_for_batch(3)
1593            .expect("3バッチ取得成功。");
1594        let (s3, d3) = val_b3
1595            .try_extract_tensor::<f32>()
1596            .expect("テンソル抽出成功。");
1597        assert_eq!(s3.as_ref(), &[3, seq_len as i64, hidden_size as i64]);
1598        assert_eq!(d3.len(), 3 * seq_len * hidden_size);
1599
1600        // 1番目と2番目のスライスが元のデータと一致しているか検証
1601        for b in 0..3 {
1602            let offset = b * seq_len * hidden_size;
1603            let slice = &d3[offset..offset + seq_len * hidden_size];
1604            for (v_orig, v_rep) in d1.iter().zip(slice.iter()) {
1605                assert!((v_orig - v_rep).abs() < 1e-3);
1606            }
1607        }
1608
1609        // state_mask_for_batch のリピート展開とパディング維持の検証
1610        let mask_b3 = features.state_mask_for_batch(3);
1611        assert_eq!(mask_b3.len(), 3 * seq_len);
1612        assert_eq!(mask_b3, vec![1i64, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]);
1613
1614        // 非圧縮 (Raw) キャッシュ時の中間アロケーション回避とリピート展開の検証
1615        let raw_shape = vec![1, seq_len, hidden_size];
1616        let raw_val = Value::from_array((raw_shape.clone(), raw_data.clone()))
1617            .expect("Value の作成に成功すること。");
1618        let raw_features = CachedStateFeatures::new_raw(
1619            key,
1620            raw_shape,
1621            raw_val.into(),
1622            raw_data.len() * std::mem::size_of::<f32>(),
1623            custom_state_mask,
1624        );
1625
1626        let raw_b2 = raw_features
1627            .as_ort_value_for_batch(2)
1628            .expect("Raw バッチ取得成功。");
1629        let (s_raw, d_raw) = raw_b2
1630            .try_extract_tensor::<f32>()
1631            .expect("Raw テンソル抽出成功。");
1632        assert_eq!(s_raw.as_ref(), &[2, seq_len as i64, hidden_size as i64]);
1633        assert_eq!(d_raw.len(), 2 * seq_len * hidden_size);
1634        assert_eq!(&d_raw[..seq_len * hidden_size], &raw_data[..]);
1635        assert_eq!(&d_raw[seq_len * hidden_size..], &raw_data[..]);
1636
1637        // バッチサイズ 0 の不正指定時にエラーとなるゼロガードを検証。
1638        assert!(features.as_ort_value_for_batch(0).is_err());
1639    }
1640}
sokuto_runtime/engine/cache.rs

sokuto_runtime/engine/
cache.rs
