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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
// #4269 (workstream M of #4263): the SafeTensors CPU forward, so `apr serve`
// and `apr chat` drive SafeTensors generation through the same
// `crate::session::Session` every other architecture uses, instead of
// `AprTransformer::generate_with_cache`'s own copy of the decode loop.
//
// `StCpuForward` wraps a borrowed `AprTransformer` and an `AprKVCache` sized
// once, up front, to the model's declared context length: `AprKVCache::new`
// already pre-allocates the whole context, so (unlike `Qwen35Forward`) this
// state never needs to double. It has no batched prefill —
// `AprTransformer::forward_with_cache` is one token at a time — so
// `ArchForward::batched_prefills` is always 0 here.
//
// This file is `include!`d into `safetensors_infer.rs`'s module scope
// (matching that file's own `include!` pattern), so it reuses that file's
// `use crate::error::{RealizarError, Result};` rather than re-importing.
use crate::apr_transformer::AprKVCache;
use crate::session::ArchForward;
/// The SafeTensors CPU forward: the only ST-CPU code a verb reaches, and
/// only through a [`Session`](crate::session::Session) (#4269).
pub struct StCpuForward<'a> {
model: &'a AprTransformer,
/// The decode state; `None` until the first turn sizes it (to the whole
/// declared context, in one allocation — see the module docs).
cache: Option<AprKVCache>,
notices: Vec<String>,
}
/// A SafeTensors CPU session: the engine's
/// [`Session`](crate::session::Session) over [`StCpuForward`].
pub type StCpuSession<'a> = crate::session::Session<StCpuForward<'a>>;
impl<'a> StCpuForward<'a> {
/// Wrap `model`'s CPU forward. The decode state starts empty.
#[must_use]
pub fn new(model: &'a AprTransformer) -> Self {
Self {
model,
cache: None,
notices: vec!["Backend: CPU (SafeTensors AprTransformer forward, #4269)".to_string()],
}
}
}
impl ArchForward for StCpuForward<'_> {
fn arch(&self) -> &'static str {
"safetensors"
}
fn on_gpu(&self) -> bool {
false
}
fn context_length(&self) -> usize {
// The length `AprKVCache::new` builds to (config.rs): the declared
// context, or its own 2048 fallback when the SafeTensors
// `config.json` left it unset (0). Read off the cache once it
// exists, so this can never disagree with what was actually
// allocated.
self.cache.as_ref().map_or_else(
|| {
if self.model.config.context_length > 0 {
self.model.config.context_length
} else {
2048
}
},
AprKVCache::capacity,
)
}
fn batched_prefills(&self) -> usize {
// AprTransformer::forward_with_cache is one token at a time; there
// is no batched prefill call to count.
0
}
fn notices(&self) -> &[String] {
&self.notices
}
fn reserve(&mut self, _positions: usize) -> Result<bool> {
if self.cache.is_some() {
return Ok(false);
}
// AprKVCache::new pre-allocates the WHOLE declared context up front
// (config.rs), so one allocation covers every turn this session will
// ever see.
self.cache = Some(AprKVCache::new(&self.model.config));
Ok(true)
}
fn forward(&mut self, tokens: &[u32], start: usize) -> Result<Vec<f32>> {
let cache = self
.cache
.as_mut()
.ok_or_else(|| RealizarError::InvalidShape {
reason: "safetensors session: the CPU state was never allocated".to_string(),
})?;
if start == 0 {
cache.clear();
}
let mut logits = Vec::new();
for (pos, &token) in tokens.iter().enumerate().skip(start) {
logits = self.model.forward_with_cache(token, cache, pos)?;
}
if logits.is_empty() {
return Err(RealizarError::InvalidShape {
reason: "safetensors session: forward advanced zero positions (tokens[start..] \
was empty)"
.to_string(),
});
}
Ok(logits)
}
}
#[cfg(test)]
#[path = "safetensors_infer_session_tests.rs"]
mod safetensors_infer_session_tests;