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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
//! バッチ推論(`Sequential::predict_batches`)と推論フェーズ計測の公開面
//! (イシュー #2192・#2582・親 #2131・#2581)。
//!
//! `compat::Sequential::predict_batches` が `DataLoader` を全バッチ反復して
//! 推論する際の内訳(データロード・tape 構築・forward・デバイス転送の累計
//! 時間・呼び出し回数)を、本モジュールの [`get_phase_metrics`]・
//! [`reset_phase_metrics`]・[`PhaseMetrics`] で読む。公開形は
//! `docs/facade-predict-batches-phase-metrics-decision.md` §8.4 の確定形
//! ((a)〜(f))に従い、実体は非公開の `batch` モジュールに置いて本ファイルの
//! `pub use` でフラットに公開する。
//!
//! **集計単位はスレッド単位**: 累計は thread-local で、「呼び出しスレッドが
//! `predict_batches` で計測した累計」を意味する(他スレッドの計測は含まない。
//! ホットパスにロックが要らずテストが分離できるため)。呼び出し前後の差分は
//! [`PhaseMetrics::since`] で取る。[`InferencePhase::DeviceTransfer`] は CPU
//! 固定経路では常に `calls == 0`(予約)。
//!
//! # 使用例
//!
//! ```
//! use fandhe_ai::Tensor;
//! use fandhe_ai::compat::Sequential;
//! use fandhe_ai::data::{DataLoader, DataLoaderConfig, TensorDataset};
//! use fandhe_ai::inference::{InferencePhase, get_phase_metrics, reset_phase_metrics};
//!
//! let model = Sequential::new().add_linear(2, 1, 7).unwrap();
//! let x = Tensor::<f32>::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]).unwrap();
//! let loader = DataLoader::new(
//! TensorDataset::new(x).unwrap(),
//! DataLoaderConfig::new(2),
//! )
//! .unwrap();
//!
//! reset_phase_metrics();
//! let before = get_phase_metrics();
//! let outputs = model.predict_batches(&loader).unwrap();
//! let delta = get_phase_metrics().since(&before);
//!
//! assert_eq!(outputs.len(), 2);
//! assert_eq!(delta.batches(), 2);
//! assert_eq!(delta.samples(), 3);
//! assert_eq!(delta.phase(InferencePhase::Forward).calls(), 2);
//! assert_eq!(delta.phase(InferencePhase::DeviceTransfer).calls(), 0);
//! assert_eq!(delta.total().calls(), 1);
//! ```
//!
//! # 自己回帰生成 `generate`(イシュー #2575・親 #2499)
//!
//! `fandhe_ai_autodiff::generate` の自己回帰ループを [`AutoregressiveModel`]・
//! [`GenerateConfig`]・[`SamplingStrategy`]・[`generate`] の 4 名で純再エクスポート
//! する(`nn::rnn`/`nn::kv_cache` と同型。facade 独自の newtype・`Tape::generate`・
//! `Sequential::generate` は作らない)。公開形の正は `docs/facade-generate-decision.md`
//! §13.2・§17(リポジトリ所有者本人の承認コメント issuecomment-6033824965)。
//! エラー型は既存の [`crate::AutodiffError`] を流用する。
//!
//! - 入出力は token id 列(`Tensor<i32>`)のみ(トークナイザは対象外)。EOS による
//! 早期停止はなく、常に `max_length` まで生成する。`top_k > vocab` は拒否する。
//! - [`GenerateConfig`] の検証(`validate`)は非公開のまま [`generate`] の入口で必ず
//! 呼ばれる(pub フィールドを書き換えて矛盾させた設定は `Err`)。
//! - 乱数は独立した `Generator`(xorshift64*)のみを使い、グローバル RNG は消費しない。
//! 暗号学的に安全ではないため、生成 token をセキュリティ用途に使わないこと。
//!
//! **KV 付きモデルの組み方**: `forward_step` の `caches[i]` は、[`crate::tape`] で作った `Tape` の
//! `stateful_attention_forward_with_cache`(#3096・設計記録 §19)で進める。`nn::kv_cache::StatefulAttention`
//! を層ごとに持ち(`num_kv_layers()` は層数を返す)、`forward_step` では `tape.stateful_attention_forward_with_cache(&sa, &x, &mut caches[i])`
//! を呼ぶ。この形ならキャッシュの reset・巻き戻しを `generate` 系が担い、1 つのモデルを複数要求で
//! 共有しても状態は混ざらない。`RefCell<StatefulAttention>` を内部に持つ旧来の形
//! (`stateful_attention_forward`。渡された `caches` を使わず `num_kv_layers()` は 0)も動作するが
//! 推奨しない(prefill 前の reset が利用者の責務になる)。
//!
//! ```
//! use fandhe_ai::Tensor;
//! use fandhe_ai::nn::kv_cache::KvCache;
//! use fandhe_ai::inference::{AutoregressiveModel, GenerateConfig, SamplingStrategy, generate};
//! use fandhe_ai::AutodiffError;
//!
//! /// 直前トークンだけで logits が決まる表引きモデル(状態なし)。
//! struct Table;
//!
//! impl AutoregressiveModel for Table {
//! fn num_kv_layers(&self) -> usize {
//! 0
//! }
//! fn forward_step(
//! &self,
//! new_ids: &Tensor<i32>,
//! _caches: &mut [KvCache],
//! ) -> Result<Tensor<f32>, AutodiffError> {
//! let (b, l) = (new_ids.shape()[0], new_ids.shape()[1]);
//! let ids = new_ids.as_slice().unwrap();
//! let mut out = vec![0.0f32; b * l * 4];
//! for (i, &id) in ids.iter().enumerate() {
//! // 次トークンは (id + 1) % 4 が最大。
//! out[i * 4 + ((id as usize + 1) % 4)] = 1.0;
//! }
//! Ok(Tensor::new(out, &[b, l, 4]).unwrap())
//! }
//! }
//!
//! let prompt = Tensor::<i32>::new(vec![0], &[1]).unwrap();
//! let cfg = GenerateConfig::new(5, SamplingStrategy::Greedy);
//! let out = generate(&Table, &prompt, &cfg).unwrap();
//! assert_eq!(out.shape(), &[5]);
//! assert_eq!(out.as_slice().unwrap(), &[0, 1, 2, 3, 0]);
//! ```
//!
//! # speculative decoding(greedy)と連続バッチング第 1 段階(イシュー #2934・親 #2932)
//!
//! `fandhe_ai_autodiff::generate` の [`generate_speculative`]・[`SpeculativeConfig`]、
//! `generate::scheduler` の [`BatchScheduler`]・[`RequestId`]・[`SchedulerLimits`] の 5 名を
//! 純再エクスポートする(facade 独自の newtype・`Tape`/`Sequential` のメソッド・別名は作らない。
//! `fandhe-ai =0.10.0` に対しては追加のみ)。公開形の正は
//! `docs/facade-speculative-decoding-batching-design.md` §17(リポジトリ所有者本人の承認コメント
//! issuecomment-6067263650 の項 1。実装記録は §18)。
//!
//! - [`generate_speculative`]: **B = 1 に限る**。Greedy は token 列一致、TopK/Temperature は
//! speculative sampling(#3101。同一シードでも `generate` と token 列は一致しない。保証は
//! 構成上の分布保存で、統計検定は契約外。設計記録 §20)。入力は `[T]`/`[1, T]`。受理統計は返さない。先読み長 `k` は残り長 `R` に応じて
//! 丸められる(`R >= 2` は `min(k, R - 1)`、`R == 1` は 1 step decode)。`num_kv_layers() == 0` の
//! モデルは `Err` で拒否する(fail-closed)。
//! - [`BatchScheduler`]: [`SchedulerLimits`] の 3 上限(同時実行数・待ち行列長・`max_length`)は
//! 必須で 0 は拒否する。要求単位の失敗は他の要求から切り離して `take_failed` に積む。完了条件は
//! `max_length` 到達のみ(EOS 停止なし)。呼び出しスレッド上で同期実行し、スレッドも I/O も持たない。
//! - [`BatchScheduler::step_batched`](#3106。設計記録 §21): `step` の opt-in 版で、decode 中の要求が
//! 2 件以上のとき 1 回の `forward_step`(`[B, 1]`・左詰め padding の `KvCache`)で束ねる。`step` は
//! 束ねず token 列の完全一致のまま。判定は単独 `generate` と統一複合判定(相対誤差 1e-3 未満 または
//! 絶対誤差 1e-5 未満)で、token 列の一致は Greedy かつ margin が十分な入力に限る(TopK/Temperature
//! では契約にしない)。束ねの失敗はその step に限り B = 1 へフォールバックして要求単位で切り離す
//! (forward が 1 回余分にかかる)。`step` と交互に呼べる。性能保証はない。対象モデルは、attention を
//! すべて `forward_with_cache` 系の入口で通し、行ごとの位置を `KvCache::valid_len` から決め、
//! attention 以外を行ごとに独立な演算に限る形(型でも実行時でも検出できないため利用者が保証する)。
//! - **対象モデルの契約**: 対象は、生成状態を渡された `caches` だけに持つモデル(`stateful_attention_forward_with_cache`
//! で `caches[i]` を進める形を含む)と、状態を持たないモデル。「内部状態保持型」(生成状態を、渡された `caches` の外に持つ型。`RefCell<StatefulAttention>`
//! 等)は `num_kv_layers()` の値によらず**対象外**で、出力を保証しない。型でも実行時でも検出できない
//! ため、利用者が保証する(設計記録 §17.5・§19.6)。
//! - 乱数は独立した `Generator`(xorshift64*)のみで非暗号。生成 token をセキュリティ用途に使わないこと。
//!
//! **KV 付きモデルの渡し方**: [`generate_speculative`] には、`stateful_attention_forward_with_cache` で
//! `caches[i]` を進める形のモデルを渡す(`num_kv_layers() > 0`。棄却時の巻き戻しは内部で行う。設計記録 §19.6)。
//! 下の doctest は、公開面だけで組んだ表引き KV モデルによる成功例(1 つ目の speculative 例。
//! 出力は単独 `generate` と token 列が一致する)と、fail-closed 契約の例(2 つ目)を示す。
//! 広い掃引は結合テスト `crates/facade/tests/speculative_batching_facade.rs` が担う(#3098)。
//!
//! ```
//! use fandhe_ai::inference::{
//! AutoregressiveModel, GenerateConfig, SamplingStrategy, SpeculativeConfig, generate,
//! generate_speculative,
//! };
//! use fandhe_ai::nn::kv_cache::{KvCache, MultiheadAttentionConfig, StatefulAttention};
//! use fandhe_ai::{AutodiffError, Tensor};
//!
//! const V: usize = 7;
//!
//! /// `caches[0]` だけを進める表引き KV モデル。logits は (直前 token, 位置) だけで決まる one-hot。
//! /// 位置は `caches[0].seq_len()` から取るため、棄却後の巻き戻しが誤れば出力がずれる。
//! struct Lm {
//! sa: StatefulAttention,
//! mul: usize,
//! }
//!
//! impl AutoregressiveModel for Lm {
//! fn num_kv_layers(&self) -> usize {
//! 1
//! }
//! fn forward_step(
//! &self,
//! new_ids: &Tensor<i32>,
//! caches: &mut [KvCache],
//! ) -> Result<Tensor<f32>, AutodiffError> {
//! let pos0 = caches[0].seq_len();
//! let tape = fandhe_ai::tape();
//! let table: Vec<f32> = (0..V * 4).map(|i| (i % 5) as f32 * 0.1).collect();
//! let emb = tape.var(&Tensor::new(table, &[V, 4]).unwrap());
//! let x = emb.embedding(new_ids, None)?;
//! tape.stateful_attention_forward_with_cache(&self.sa, &x, &mut caches[0])?;
//! let l = new_ids.shape()[new_ids.shape().len() - 1];
//! let mut out = vec![0.0f32; l * V];
//! for (i, &tok) in new_ids.as_slice().unwrap().iter().enumerate() {
//! out[i * V + (self.mul * tok as usize + pos0 + i + 1) % V] = 1.0;
//! }
//! Ok(Tensor::new(out, &[1, l, V]).unwrap())
//! }
//! }
//!
//! let make = |mul, seed| Lm {
//! sa: StatefulAttention::from_config(&MultiheadAttentionConfig::new(4, 2), seed).unwrap(),
//! mul,
//! };
//! let target = make(3, 1);
//! // 規則の違う draft(棄却が起きる)。
//! let draft = make(2, 2);
//!
//! let prompt = Tensor::<i32>::new(vec![0], &[1]).unwrap();
//! let cfg = GenerateConfig::new(6, SamplingStrategy::Greedy);
//! let out = generate_speculative(&target, &draft, &prompt, &cfg, &SpeculativeConfig::new(2)).unwrap();
//! // 出力は target 単独の `generate` と token 列が一致する。
//! let alone = generate(&target, &prompt, &cfg).unwrap();
//! assert_eq!(out.as_slice().unwrap(), alone.as_slice().unwrap());
//! assert_eq!(out.as_slice().unwrap().len(), 6);
//! ```
//!
//! ```
//! use fandhe_ai::inference::{
//! AutoregressiveModel, BatchScheduler, GenerateConfig, SamplingStrategy, SchedulerLimits,
//! generate,
//! };
//! use fandhe_ai::nn::kv_cache::KvCache;
//! use fandhe_ai::{AutodiffError, Tensor};
//!
//! /// 直前トークンだけで logits が決まる表引きモデル(状態なし)。
//! struct Table;
//!
//! impl AutoregressiveModel for Table {
//! fn num_kv_layers(&self) -> usize {
//! 0
//! }
//! fn forward_step(
//! &self,
//! new_ids: &Tensor<i32>,
//! _caches: &mut [KvCache],
//! ) -> Result<Tensor<f32>, AutodiffError> {
//! let (b, l) = (new_ids.shape()[0], new_ids.shape()[1]);
//! let ids = new_ids.as_slice().unwrap();
//! let mut out = vec![0.0f32; b * l * 4];
//! for (i, &id) in ids.iter().enumerate() {
//! out[i * 4 + ((id as usize + 1) % 4)] = 1.0;
//! }
//! Ok(Tensor::new(out, &[b, l, 4]).unwrap())
//! }
//! }
//!
//! let limits = SchedulerLimits::new(2, 4, 8).unwrap();
//! let mut scheduler = BatchScheduler::new(limits);
//! let cfg = GenerateConfig::new(5, SamplingStrategy::Greedy);
//! let p0 = Tensor::<i32>::new(vec![0], &[1]).unwrap();
//! let p1 = Tensor::<i32>::new(vec![2], &[1]).unwrap();
//! let id0 = scheduler.submit(&p0, &cfg).unwrap();
//! let id1 = scheduler.submit(&p1, &cfg).unwrap();
//! assert_ne!(id0, id1);
//!
//! let mut finished = Vec::new();
//! for _ in 0..16 {
//! scheduler.step(&Table).unwrap();
//! finished.extend(scheduler.take_finished());
//! if scheduler.queued_len() + scheduler.active_len() == 0 {
//! break;
//! }
//! }
//! assert!(scheduler.take_failed().is_empty());
//! assert_eq!(finished.len(), 2);
//! // 各要求の出力は単独の `generate` と token 列が一致する。
//! for (id, out) in &finished {
//! let prompt = if *id == id0 { &p0 } else { &p1 };
//! let alone = generate(&Table, prompt, &cfg).unwrap();
//! assert_eq!(out.as_slice().unwrap(), alone.as_slice().unwrap());
//! }
//! ```
//!
//! 束ねる経路(`step_batched`)の利用例。状態なしのモデルでは単独 `generate` と token 列が一致する。
//!
//! ```
//! use fandhe_ai::inference::{
//! AutoregressiveModel, BatchScheduler, GenerateConfig, SamplingStrategy, SchedulerLimits,
//! generate,
//! };
//! use fandhe_ai::nn::kv_cache::KvCache;
//! use fandhe_ai::{AutodiffError, Tensor};
//!
//! /// 直前トークンだけで logits が決まる表引きモデル(状態なし)。
//! struct Table;
//!
//! impl AutoregressiveModel for Table {
//! fn num_kv_layers(&self) -> usize {
//! 0
//! }
//! fn forward_step(
//! &self,
//! new_ids: &Tensor<i32>,
//! _caches: &mut [KvCache],
//! ) -> Result<Tensor<f32>, AutodiffError> {
//! let (b, l) = (new_ids.shape()[0], new_ids.shape()[1]);
//! let ids = new_ids.as_slice().unwrap();
//! let mut out = vec![0.0f32; b * l * 4];
//! for (i, &id) in ids.iter().enumerate() {
//! out[i * 4 + ((id as usize + 1) % 4)] = 1.0;
//! }
//! Ok(Tensor::new(out, &[b, l, 4]).unwrap())
//! }
//! }
//!
//! let limits = SchedulerLimits::new(2, 4, 8).unwrap();
//! let mut scheduler = BatchScheduler::new(limits);
//! let cfg = GenerateConfig::new(5, SamplingStrategy::Greedy);
//! let p0 = Tensor::<i32>::new(vec![0], &[1]).unwrap();
//! let p1 = Tensor::<i32>::new(vec![2], &[1]).unwrap();
//! let id0 = scheduler.submit(&p0, &cfg).unwrap();
//! let _id1 = scheduler.submit(&p1, &cfg).unwrap();
//!
//! let mut finished = Vec::new();
//! for _ in 0..16 {
//! scheduler.step_batched(&Table).unwrap();
//! finished.extend(scheduler.take_finished());
//! if scheduler.queued_len() + scheduler.active_len() == 0 {
//! break;
//! }
//! }
//! assert!(scheduler.take_failed().is_empty());
//! assert_eq!(finished.len(), 2);
//! for (id, out) in &finished {
//! let prompt = if *id == id0 { &p0 } else { &p1 };
//! let alone = generate(&Table, prompt, &cfg).unwrap();
//! assert_eq!(out.as_slice().unwrap(), alone.as_slice().unwrap());
//! }
//! ```
//!
//! ```
//! use fandhe_ai::inference::{
//! AutoregressiveModel, GenerateConfig, SamplingStrategy, SpeculativeConfig,
//! generate_speculative,
//! };
//! use fandhe_ai::nn::kv_cache::KvCache;
//! use fandhe_ai::{AutodiffError, Tensor};
//!
//! /// 状態なしのモデル(`num_kv_layers() == 0`)。
//! struct Table;
//!
//! impl AutoregressiveModel for Table {
//! fn num_kv_layers(&self) -> usize {
//! 0
//! }
//! fn forward_step(
//! &self,
//! new_ids: &Tensor<i32>,
//! _caches: &mut [KvCache],
//! ) -> Result<Tensor<f32>, AutodiffError> {
//! let (b, l) = (new_ids.shape()[0], new_ids.shape()[1]);
//! Ok(Tensor::new(vec![0.0f32; b * l * 4], &[b, l, 4]).unwrap())
//! }
//! }
//!
//! let prompt = Tensor::<i32>::new(vec![0], &[1]).unwrap();
//! let spec = SpeculativeConfig::new(2);
//!
//! // KV を使わない(`num_kv_layers() == 0` の)モデルは fail-closed で拒否される。
//! let greedy = GenerateConfig::new(5, SamplingStrategy::Greedy);
//! let r = generate_speculative(&Table, &Table, &prompt, &greedy, &spec);
//! assert!(matches!(r, Err(AutodiffError::InvalidArgument(_))));
//!
//! // サンプリング戦略でも `num_kv_layers() == 0` は拒否される。
//! let topk = GenerateConfig::new(5, SamplingStrategy::TopK(2));
//! let r = generate_speculative(&Table, &Table, &prompt, &topk, &spec);
//! assert!(matches!(r, Err(AutodiffError::InvalidArgument(_))));
//! ```
pub
pub use ;
// 自己回帰生成(イシュー #2575)。autodiff の形をそのまま再エクスポートする(承認形。
// `api_surface.rs::GENERATE_APPROVED_REEXPORT` が完全一致を検査するため、分割・別名・
// 別ファイルへの移動をしない)。
pub use ;
// 連続バッチングのスケジューラ(第 1 段階。イシュー #2934・親 #2932)。autodiff の形を
// そのまま再エクスポートする(承認形。承認根拠: リポジトリ所有者本人の承認コメント
// issuecomment-6067263650 の項 1、公開形の正は `docs/facade-speculative-decoding-batching-design.md`
// §17.2。`api_surface.rs::SCHEDULER_APPROVED_REEXPORT` が完全一致を検査するため、分割・
// 別名・モジュール自体の再エクスポート・別ファイルへの移動をしない)。
pub use ;
// greedy 版 speculative decoding(B = 1。イシュー #2934・親 #2932)。承認根拠・検査は
// 上の連続バッチング文と同じ(`api_surface.rs::SPECULATIVE_APPROVED_REEXPORT`)。
pub use ;