1use superfluid_abi::*;
2
3use crate::{Engine, EngineConfig, InMemoryRings, MockEngine, Rings};
4
5pub trait StageRings: Rings {
6 fn stage_prompt(&mut self, tokens: &[u32]) -> TokenRef;
7}
8
9impl StageRings for InMemoryRings {
10 fn stage_prompt(&mut self, tokens: &[u32]) -> TokenRef {
11 InMemoryRings::stage_prompt(self, tokens)
12 }
13}
14
15pub const DEFAULT_MAX_DECODE_LANES: u32 = 8;
16
17pub struct Harness<E: Engine = MockEngine, R: StageRings = InMemoryRings> {
18 pub engine: E,
19 pub rings: R,
20 pub next_seq: std::cell::Cell<u64>,
21 pub max_decode_lanes: u32,
22}
23
24impl Harness<MockEngine> {
25 pub fn new() -> Harness<MockEngine> {
26 Harness::with_config(EngineConfig::default())
27 }
28
29 pub fn with_config(cfg: EngineConfig) -> Harness<MockEngine> {
30 Harness::with_engine(MockEngine::new(cfg))
31 }
32}
33
34impl Default for Harness<MockEngine> {
35 fn default() -> Self {
36 Harness::new()
37 }
38}
39
40impl<E: Engine> Harness<E> {
41 pub fn with_engine(engine: E) -> Harness<E> {
42 Harness::with_engine_and_rings(engine, InMemoryRings::new())
43 }
44}
45
46impl<E: Engine, R: StageRings> Harness<E, R> {
47 pub fn with_engine_and_rings(engine: E, rings: R) -> Harness<E, R> {
48 Harness {
49 engine,
50 rings,
51 next_seq: std::cell::Cell::new(1),
52 max_decode_lanes: DEFAULT_MAX_DECODE_LANES,
53 }
54 }
55
56 pub fn with_max_decode_lanes(mut self, n: u32) -> Self {
57 self.max_decode_lanes = n;
58 self
59 }
60
61 pub fn plan(&self) -> PlanBuilder {
62 let seq = self.next_seq.get();
63 self.next_seq.set(seq + 1);
64 PlanBuilder::new(seq).max_decode_lanes(self.max_decode_lanes)
65 }
66
67 pub fn tick_ok(&mut self, pb: PlanBuilder) -> Events {
68 let (plan, arena) = pb.build();
69 let ev = self
70 .engine
71 .tick(&plan, &arena, &mut self.rings)
72 .unwrap_or_else(|s| panic!("tick rejected: {s:?}"));
73 Events::decode(ev)
74 }
75
76 pub fn tick_err(&mut self, pb: PlanBuilder) -> Status {
77 let (plan, arena) = pb.build();
78 match self.engine.tick(&plan, &arena, &mut self.rings) {
79 Ok(_) => panic!("tick unexpectedly accepted"),
80 Err(s) => s,
81 }
82 }
83
84 pub fn prompt(&mut self, tokens: &[u32]) -> TokenRef {
85 StageRings::stage_prompt(&mut self.rings, tokens)
86 }
87
88 pub fn rings_tokens(&self, r: &TokenRef) -> Vec<u32> {
89 self.rings.read_tokens(r).unwrap()
90 }
91}
92
93pub trait EngineInspect: Engine {
94 fn lane_committed(&self, lane_tag: u64) -> Option<&[u32]>;
95 fn inject_fault(&mut self, lane_tag: u64, code: u32);
96}
97
98impl EngineInspect for MockEngine {
99 fn lane_committed(&self, lane_tag: u64) -> Option<&[u32]> {
100 MockEngine::lane_committed(self, lane_tag)
101 }
102 fn inject_fault(&mut self, lane_tag: u64, code: u32) {
103 MockEngine::inject_fault(self, lane_tag, code)
104 }
105}
106
107pub struct Events {
108 pub plan_seq: u64,
109 pub tick_status: u32,
110 pub admit_results: Vec<AdmitResult>,
111 pub emits: Vec<LaneEmit>,
112 pub faults: Vec<LaneFault>,
113 pub shed_chunks: u32,
114 pub shed_entries: Vec<ShedEntry>,
115 pub op_completions: Vec<OpComplete>,
116 pub logprobs: Vec<LaneLogprob>,
117}
118
119impl Events {
120 pub fn decode(ev: &TickEvents) -> Events {
121 let bounds = ArenaBounds {
122 base: 0,
123 len: usize::MAX,
124 };
125 unsafe {
127 Events {
128 plan_seq: ev.plan_seq,
129 tick_status: ev.tick_status,
130 admit_results: read_array(&ev.admit_results, &bounds).unwrap().collect(),
131 emits: read_array(&ev.emits, &bounds).unwrap().collect(),
132 faults: read_array(&ev.faults, &bounds).unwrap().collect(),
133 shed_chunks: ev.shed.prefill_chunks_dropped,
134 shed_entries: read_array(&ev.shed.entries, &bounds).unwrap().collect(),
135 op_completions: read_array(&ev.op_completions, &bounds).unwrap().collect(),
136 logprobs: read_array(&ev.logprobs, &bounds).unwrap().collect(),
137 }
138 }
139 }
140
141 pub fn emit_for(&self, lane: u64) -> &LaneEmit {
142 self.emits
143 .iter()
144 .find(|e| e.lane_tag == lane)
145 .expect("no emit for lane")
146 }
147
148 pub fn logprobs_for(&self, lane: u64) -> Vec<&LaneLogprob> {
149 self.logprobs.iter().filter(|l| l.lane_tag == lane).collect()
150 }
151
152 pub fn admit_for(&self, lane: u64) -> &AdmitResult {
153 self.admit_results
154 .iter()
155 .find(|a| a.lane_tag == lane)
156 .expect("no admit result for lane")
157 }
158}
159
160use superfluid_abi::array::read_array;
161
162pub struct PlanBuilder {
163 plan_seq: u64,
164 flags: u32,
165 prefill_budget: u32,
166 max_decode_lanes: u32,
167 admits: Vec<LaneAdmit>,
168 commits: Vec<LaneCommit>,
169 prefills: Vec<LanePrefill>,
170 decodes: Vec<LaneDecode>,
171 retires: Vec<LaneRetire>,
172 victims: Vec<u64>,
173 evictable_classes: u64,
174 max_evict_bytes: u64,
175 protected_quota_bytes: u64,
176 on_partial_emit:
177 Option<unsafe extern "C" fn(user: *mut std::ffi::c_void, lane_tag: u64, tokens: *const u32, n_tokens: u32)>,
178 partial_emit_user: *mut std::ffi::c_void,
179}
180
181impl PlanBuilder {
182 pub fn new(plan_seq: u64) -> PlanBuilder {
183 PlanBuilder {
184 plan_seq,
185 flags: 0,
186 prefill_budget: 4096,
187 max_decode_lanes: 8,
188 admits: Vec::new(),
189 commits: Vec::new(),
190 prefills: Vec::new(),
191 decodes: Vec::new(),
192 retires: Vec::new(),
193 victims: Vec::new(),
194 evictable_classes: u64::MAX,
195 max_evict_bytes: u64::MAX,
196 protected_quota_bytes: 0,
197 on_partial_emit: None,
198 partial_emit_user: std::ptr::null_mut(),
199 }
200 }
201
202 pub fn flags(mut self, f: u32) -> Self {
203 self.flags = f;
204 self
205 }
206
207 pub fn prefill_budget(mut self, b: u32) -> Self {
208 self.prefill_budget = b;
209 self
210 }
211
212 pub fn max_decode_lanes(mut self, n: u32) -> Self {
213 self.max_decode_lanes = n;
214 self
215 }
216
217 pub fn evict_envelope(mut self, classes: u64, max_bytes: u64) -> Self {
218 self.evictable_classes = classes;
219 self.max_evict_bytes = max_bytes;
220 self
221 }
222
223 pub fn protected_quota(mut self, bytes: u64) -> Self {
224 self.protected_quota_bytes = bytes;
225 self
226 }
227
228 pub fn partial_emit(
229 mut self,
230 cb: unsafe extern "C" fn(user: *mut std::ffi::c_void, lane_tag: u64, tokens: *const u32, n_tokens: u32),
231 user: *mut std::ffi::c_void,
232 ) -> Self {
233 self.on_partial_emit = Some(cb);
234 self.partial_emit_user = user;
235 self
236 }
237
238 pub fn admit(mut self, lane: u64, prompt: TokenRef) -> Self {
239 self.admits.push(basic_admit(lane, prompt));
240 self
241 }
242
243 pub fn admit_with(
244 mut self,
245 f: impl FnOnce(LaneAdmit) -> LaneAdmit,
246 lane: u64,
247 prompt: TokenRef,
248 ) -> Self {
249 self.admits.push(f(basic_admit(lane, prompt)));
250 self
251 }
252
253 pub fn commit(mut self, lane: u64, token: u32, nonce: u64) -> Self {
254 self.commits.push(LaneCommit {
255 lane_tag: lane,
256 token_id: token,
257 _pad0: 0,
258 logits_nonce: nonce,
259 });
260 self
261 }
262
263 pub fn prefill(mut self, lane: u64, offset: u32, count: u32) -> Self {
264 self.prefills.push(LanePrefill {
265 lane_tag: lane,
266 token_offset: offset,
267 token_count: count,
268 });
269 self
270 }
271
272 pub fn decode(mut self, lane: u64, max_new: u16) -> Self {
273 self.decodes.push(LaneDecode {
274 lane_tag: lane,
275 max_new_tokens: max_new,
276 overshoot: 0,
277 _pad1: 0,
278 });
279 self
280 }
281
282 pub fn retire(mut self, lane: u64, publish: bool) -> Self {
283 self.retires.push(LaneRetire {
284 lane_tag: lane,
285 publish_to_cache: publish as u8,
286 _pad0: [0; 7],
287 });
288 self
289 }
290
291 pub fn build(self) -> (TickPlan, RecordArena) {
292 let mut arena = RecordArena::new();
293 let plan = TickPlan {
294 struct_size: std::mem::size_of::<TickPlan>() as u64,
295 plan_seq: self.plan_seq,
296 flags: self.flags,
297 _pad0: 0,
298 prefill_token_budget: self.prefill_budget,
299 max_decode_lanes: self.max_decode_lanes,
300 admits: arena.push_records(&self.admits),
301 commits: arena.push_records(&self.commits),
302 prefills: arena.push_records(&self.prefills),
303 decodes: arena.push_records(&self.decodes),
304 retires: arena.push_records(&self.retires),
305 shed_policy: ShedPolicy {
306 victim_lanes: arena.push_records(&self.victims),
307 evictable_cache_classes: self.evictable_classes,
308 protected_quota_bytes: self.protected_quota_bytes,
309 max_evict_bytes: self.max_evict_bytes,
310 },
311 on_partial_emit: self.on_partial_emit,
312 partial_emit_user: self.partial_emit_user,
313 };
314 (plan, arena)
315 }
316}
317
318pub fn register_prompt_lookup_with<E: Engine>(
319 h: &mut Harness<E>,
320 impl_hash: [u8; 32],
321 config_hash: [u8; 32],
322 kernel_caps_required: u32,
323) -> Result<(u32, Vec<ExactnessCert>), Status> {
324 let strategy_id = std::ffi::CString::new("prompt-lookup").unwrap();
325 let impl_version = std::ffi::CString::new("1.0.0").unwrap();
326 let mut arena = RecordArena::new();
327 let caps = [CapabilityReq {
328 kind_id: strategy_cap::PROPOSAL_LINEAR,
329 _pad0: 0,
330 params: Array::EMPTY,
331 }];
332 let reg = StrategyRegistration {
333 struct_size: std::mem::size_of::<StrategyRegistration>() as u64,
334 strategy_id: strategy_id.as_ptr(),
335 impl_version: impl_version.as_ptr(),
336 impl_hash,
337 config_hash,
338 artifacts: Array::EMPTY,
339 taps: Array::EMPTY,
340 capabilities: arena.push_records(&caps),
341 target_archs: Array::EMPTY,
342 kernel_caps_required,
343 _pad0: 0,
344 est_state_bytes: 0,
345 claimed_exactness: exactness::SEED_PATH_INVARIANT,
346 _pad1: [0; 3],
347 rng_contract_version: 1,
348 };
349 let grant = h.engine.strategy_register(®)?;
350 let bounds = ArenaBounds {
351 base: 0,
352 len: usize::MAX,
353 };
354 let certs: Vec<ExactnessCert> =
356 unsafe { superfluid_abi::array::read_array(&grant.certificates, &bounds) }
357 .unwrap()
358 .collect();
359 Ok((grant.strategy_slot, certs))
360}
361
362pub fn basic_admit(lane: u64, prompt: TokenRef) -> LaneAdmit {
363 LaneAdmit {
364 lane_tag: lane,
365 prompt,
366 seed_handle: 0,
367 sampling: sampling::GPU_GREEDY,
368 logit_bias_handle: 0,
369 params: SamplingParams { flags: sampling_flag::IGNORE_EOS, ..Default::default() },
370 rng_counter_base: 0,
371 grammar_handle: 0,
372 strategy_slot: 0,
373 determinism_class: determinism::BEST_EFFORT,
374 minimum_exactness: exactness::APPROXIMATE,
375 allow_approximate: 0,
376 want_logprobs: 0,
377 required_cert_id: 0,
378 host_sampler_identity: [0; 32],
379 grammar_replay: 0,
380 decode_replay: 0,
381 }
382}
383
384pub fn rows_agree(reference: &[f32], other: &[f32], tolerance: f32) -> Result<(), String> {
385 if reference.len() != other.len() {
386 return Err(format!("rows of {} and {} logits", reference.len(), other.len()));
387 }
388 let (mut at, mut diff) = (0, 0f32);
389 for (i, (&a, &b)) in reference.iter().zip(other).enumerate() {
390 if a.is_nan() || b.is_nan() {
391 return Err(format!("logit at id {i} is not a number ({a} vs {b})"));
392 }
393 if a.is_infinite() || b.is_infinite() {
394 if a == b {
395 continue;
396 }
397 return Err(format!("logit at id {i}: {a} vs {b}"));
398 }
399 let d = (a - b).abs();
400 if d > diff {
401 (at, diff) = (i, d);
402 }
403 }
404 if diff > tolerance {
405 return Err(format!(
406 "largest logit difference {diff} at id {at} ({} vs {}) exceeds {tolerance}",
407 reference[at], other[at]
408 ));
409 }
410 Ok(())
411}
412
413pub fn top2_margin(record: &LaneLogprob) -> f32 {
414 match record.n_top {
415 0 | 1 => f32::INFINITY,
416 _ => record.top_logprobs[0] - record.top_logprobs[1],
417 }
418}
419
420pub fn tokens_agree(reference: &[u32], margins: &[f32], other: &[u32], tolerance: f32) -> Result<(), String> {
421 for (k, (a, b)) in reference.iter().zip(other).enumerate() {
422 if a == b {
423 continue;
424 }
425 let margin = margins.get(k).copied().unwrap_or(f32::INFINITY);
426 if margin <= 2.0 * tolerance {
427 return Ok(());
428 }
429 return Err(format!(
430 "step {k}: {b} where the reference took {a} by a margin of {margin} (tolerance {tolerance}): {reference:?} vs {other:?}"
431 ));
432 }
433 if reference.len() != other.len() {
434 return Err(format!("{} tokens where the reference has {}: {reference:?} vs {other:?}", other.len(), reference.len()));
435 }
436 Ok(())
437}
438
439pub mod tick_contract {
440 use super::*;
441
442 pub fn admit_prefill_decode_in_one_tick<E: EngineInspect, R: StageRings>(mut h: Harness<E, R>) {
443 let p = h.prompt(&[1, 2, 3, 4]);
444 let ev = h.tick_ok(h.plan().admit(7, p).prefill(7, 0, 4).decode(7, 3));
445 assert_eq!(ev.tick_status, tick_status::OK);
446 assert_eq!(ev.admit_for(7).status, admit_status::ADMITTED);
447 let e = ev.emit_for(7);
448 assert_eq!(e.n_tokens, 3);
449 assert_eq!(e.finish, finish::NONE);
450 assert_eq!(h.engine.lane_committed(7).unwrap().len(), 7);
451 }
452
453 pub fn same_log_same_seed_same_tokens<E: Engine, R: StageRings>(make: impl Fn() -> Harness<E, R>) {
454 let run = || {
455 let mut h = make();
456 let p = h.prompt(&[10, 20, 30]);
457 let ev = h.tick_ok(h.plan().admit(1, p).prefill(1, 0, 3).decode(1, 8));
458 let e = *ev.emit_for(1);
459 h.rings_tokens(&e.token_ref)
460 };
461 assert_eq!(run(), run());
462 }
463
464 pub fn plan_seq_must_be_monotonic<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
465 let p = h.prompt(&[1]);
466 h.tick_ok(h.plan().admit(1, p).prefill(1, 0, 1));
467 let stale = PlanBuilder::new(1);
468 assert_eq!(h.tick_err(stale), Status::RejectIllegalCombination);
469 }
470
471 pub fn structural_rejections<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
472 let p = h.prompt(&[1, 2]);
473 let err = h.tick_err(h.plan().admit(1, p).admit(1, p));
474 assert_eq!(err, Status::RejectDuplicateLane);
475 let p2 = h.prompt(&[1, 2]);
476 let err = h.tick_err(h.plan().admit(2, p2).retire(2, false));
477 assert_eq!(err, Status::RejectIllegalCombination);
478 let err = h.tick_err(h.plan().commit(99, 5, 1));
479 assert_eq!(err, Status::RejectUnknownLane);
480 let err = h.tick_err(h.plan().decode(98, 1));
481 assert_eq!(err, Status::RejectUnknownLane);
482 let err = h.tick_err(h.plan().max_decode_lanes(1000));
483 assert_eq!(err, Status::RejectBudget);
484 let p3 = h.prompt(&[5]);
485 let ev = h.tick_ok(h.plan().admit(3, p3).prefill(3, 0, 1).decode(3, 1));
486 assert_eq!(ev.tick_status, tick_status::OK);
487 }
488
489 pub fn stale_ring_generation_rejected<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
490 let mut p = h.prompt(&[1, 2, 3]);
491 p.generation += 7;
492 let err = h.tick_err(h.plan().admit(1, p));
493 assert_eq!(err, Status::RejectBadRefGeneration);
494 }
495
496 pub fn drain_takes_no_new_work<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
497 let p = h.prompt(&[1]);
498 let err = h.tick_err(h.plan().flags(tick_flags::DRAIN).admit(1, p));
499 assert_eq!(err, Status::RejectIllegalCombination);
500 let p2 = h.prompt(&[1, 2]);
501 h.tick_ok(h.plan().admit(2, p2).prefill(2, 0, 2));
502 let ev = h.tick_ok(h.plan().flags(tick_flags::DRAIN).decode(2, 1));
503 assert_eq!(ev.emit_for(2).n_tokens, 1);
504 }
505
506 pub fn a_replayed_decode_tail_resumes_the_lane<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
507 const CUT: usize = 5;
508 let prompt: Vec<u32> = vec![11, 12, 13, 14, 15, 16];
509 const SEED: u64 = 1000;
510 let admitted_at = |len: usize| {
511 move |mut a: LaneAdmit| {
512 a.params.flags |= sampling_flag::IGNORE_EOS;
513 a.rng_counter_base = SEED + len as u64;
514 a
515 }
516 };
517 let n = prompt.len() as u32;
518 let plain = admitted_at(prompt.len());
519 let resumed_len = prompt.len() + CUT;
520 let replayed = move |a: LaneAdmit| {
521 let mut a = admitted_at(resumed_len)(a);
522 a.decode_replay = CUT as u32;
523 a
524 };
525
526 let p = h.prompt(&prompt);
527 let ev = h.tick_ok(h.plan().admit_with(plain, 1, p).prefill(1, 0, n).decode(1, 8));
528 let e = *ev.emit_for(1);
529 let full = h.rings_tokens(&e.token_ref);
530 assert_eq!(full.len(), 8);
531 h.tick_ok(h.plan().retire(1, false));
532
533 let mut resumed = prompt.clone();
534 resumed.extend_from_slice(&full[..CUT]);
535
536 let p = h.prompt(&resumed);
537 h.tick_err(h.plan().admit_with(replayed, 2, p).prefill(2, 0, resumed.len() as u32).decode(2, 3));
538
539 let p = h.prompt(&resumed);
540 let ev = h.tick_ok(h.plan().admit_with(replayed, 2, p).prefill(2, 0, n).decode(2, CUT as u16 + 3));
541 let e = *ev.emit_for(2);
542 assert_eq!(e.finish, finish::NONE);
543 assert_eq!(h.rings_tokens(&e.token_ref), full[CUT..], "the continuation is the lane it resumes");
544 h.tick_ok(h.plan().retire(2, false));
545
546 let p = h.prompt(&resumed);
547 let ev = h.tick_ok(h.plan().admit_with(replayed, 3, p).prefill(3, 0, n).decode(3, 3));
548 assert_eq!((ev.emit_for(3).n_tokens, ev.emit_for(3).finish), (0, finish::NONE), "inside the tail");
549 let ev = h.tick_ok(h.plan().decode(3, CUT as u16 - 3));
550 assert_eq!(ev.emit_for(3).n_tokens, 0, "the tail's last token takes the grant's last unit");
551 let ev = h.tick_ok(h.plan().decode(3, 3));
552 let e = *ev.emit_for(3);
553 assert_eq!(h.rings_tokens(&e.token_ref), full[CUT..], "paced or not, the same tokens");
554 h.tick_ok(h.plan().retire(3, false));
555
556 let p = h.prompt(&resumed);
557 let whole = move |a: LaneAdmit| {
558 let mut a = admitted_at(resumed_len)(a);
559 a.decode_replay = resumed_len as u32;
560 a
561 };
562 h.tick_err(h.plan().admit_with(whole, 4, p).decode(4, 1));
563 }
564
565 pub fn host_alternation_full_loop<E: EngineInspect, R: StageRings>(mut h: Harness<E, R>) {
566 let p = h.prompt(&[1, 2]);
567 let ev = h.tick_ok(
568 h.plan()
569 .admit_with(
570 |mut a| {
571 a.sampling = sampling::HOST;
572 a
573 },
574 1,
575 p,
576 )
577 .prefill(1, 0, 2)
578 .decode(1, 1),
579 );
580 let e = ev.emit_for(1);
581 assert_eq!(e.n_tokens, 0);
582 let nonce = e.logits_row.generation;
583 assert_ne!(nonce, 0);
584
585 let err = h.tick_err(h.plan().decode(1, 1));
586 assert_eq!(err, Status::RejectHostRules);
587
588 let err = h.tick_err(h.plan().commit(1, 42, nonce + 1).decode(1, 1));
589 assert_eq!(err, Status::RejectStaleNonce);
590
591 let ev = h.tick_ok(h.plan().commit(1, 42, nonce).decode(1, 1));
592 let committed = h.engine.lane_committed(1).unwrap();
593 assert_eq!(*committed.last().unwrap(), 42);
594 let e2 = ev.emit_for(1);
595 assert_eq!(e2.n_tokens, 0);
596 assert_ne!(e2.logits_row.generation, nonce, "fresh row, fresh nonce");
597 }
598
599 pub fn host_lane_max_new_tokens_must_be_one<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
600 let p = h.prompt(&[1]);
601 let err = h.tick_err(
602 h.plan()
603 .admit_with(
604 |mut a| {
605 a.sampling = sampling::HOST;
606 a
607 },
608 1,
609 p,
610 )
611 .prefill(1, 0, 1)
612 .decode(1, 4),
613 );
614 assert_eq!(err, Status::RejectHostRules);
615 }
616
617 pub fn commit_on_gpu_lane_rejected<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
618 let p = h.prompt(&[1]);
619 h.tick_ok(h.plan().admit(1, p).prefill(1, 0, 1));
620 let err = h.tick_err(h.plan().commit(1, 5, 1));
621 assert_eq!(err, Status::RejectIllegalCombination);
622 }
623
624 pub fn prefill_budget_sheds_and_reports<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
625 let tokens: Vec<u32> = (0..100).collect();
626 let p = h.prompt(&tokens);
627
628 let err = h.tick_err(
629 h.plan()
630 .prefill_budget(10)
631 .admit(1, p)
632 .prefill(1, 0, 100)
633 .decode(1, 1),
634 );
635 assert_eq!(err, Status::RejectIllegalCombination);
636
637 let ev = h.tick_ok(h.plan().prefill_budget(10).admit(1, p).prefill(1, 0, 100));
638 assert_eq!(ev.tick_status, tick_status::SHED);
639 assert_eq!(ev.shed_chunks, 1);
640 let entry = ev
641 .shed_entries
642 .iter()
643 .find(|e| e.kind == shed_kind::PREFILL_CHUNK)
644 .unwrap();
645 assert_eq!(entry.tokens, 90);
646 assert_eq!(entry.lane_tag, 1);
647
648 let ev = h.tick_ok(h.plan().prefill_budget(200).prefill(1, 10, 90).decode(1, 1));
649 assert_eq!(ev.tick_status, tick_status::OK);
650 assert_eq!(ev.emit_for(1).n_tokens, 1);
651 }
652
653 pub fn noncontiguous_prefill_rejected<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
654 let p = h.prompt(&[1, 2, 3, 4]);
655 h.tick_ok(h.plan().admit(1, p).prefill(1, 0, 2));
656 let err = h.tick_err(h.plan().prefill(1, 3, 1));
657 assert_eq!(err, Status::RejectIllegalCombination);
658 }
659
660 pub fn per_lane_fault_does_not_fail_tick<E: EngineInspect, R: StageRings>(mut h: Harness<E, R>) {
661 let p1 = h.prompt(&[1, 2]);
662 let p2 = h.prompt(&[3, 4]);
663 h.tick_ok(h.plan().admit(1, p1).prefill(1, 0, 2));
664 h.tick_ok(h.plan().admit(2, p2).prefill(2, 0, 2));
665 h.engine
666 .inject_fault(1, superfluid_abi::status::fault_code::GRAMMAR_OVERFLOW);
667 let ev = h.tick_ok(h.plan().decode(1, 1).decode(2, 1));
668 assert_eq!(ev.tick_status, tick_status::LANE_ERRORS);
669 assert_eq!(ev.faults.len(), 1);
670 assert_eq!(ev.faults[0].lane_tag, 1);
671 assert_eq!(
672 ev.faults[0].code,
673 superfluid_abi::status::fault_code::GRAMMAR_OVERFLOW
674 );
675 assert_eq!(ev.emit_for(2).n_tokens, 1);
676 let err = h.tick_err(h.plan().decode(1, 1));
677 assert_eq!(err, Status::RejectUnknownLane);
678 let ev = h.tick_ok(h.plan().retire(1, true).decode(2, 1));
679 assert_eq!(ev.emit_for(2).n_tokens, 1);
680 let err = h.tick_err(h.plan().retire(1, false));
681 assert_eq!(err, Status::RejectUnknownLane);
682 }
683
684 pub fn retire_publish_then_seeded_readmission<E: EngineInspect, R: StageRings>(mut h: Harness<E, R>) {
685 let paged_desc = h
686 .engine
687 .state_spaces()
688 .iter()
689 .find(|d| d.kind == space_kind::PAGED_TOKEN_KV)
690 .copied()
691 .expect("the engine advertises a paged token-KV space");
692 let space = paged_desc.space_id;
693 let unit = lcm(paged_desc.page_size_tokens.max(1) as u64, paged_desc.snapshot_interval_tokens.max(1) as u64);
694 let n = 64u64.div_ceil(unit) * unit;
695 let tokens: Vec<u32> = (100..100 + n as u32).collect();
696 let p = h.prompt(&tokens);
697 h.tick_ok(h.plan().admit(1, p).prefill(1, 0, n as u32));
698 h.tick_ok(h.plan().retire(1, true));
699
700 let mut longer = tokens.clone();
701 longer.push(100 + n as u32);
702 let m = h.engine.space_match(ALL_SPACES, &longer, &[]).unwrap();
703 let bounds = ArenaBounds {
704 base: 0,
705 len: usize::MAX,
706 };
707 let spaces: Vec<SpaceMatch> =
709 unsafe { superfluid_abi::array::read_array(&m.spaces, &bounds) }
710 .unwrap()
711 .collect();
712 let paged = spaces.iter().find(|s| s.space_id == space).expect("the paged space in the match");
713 let cands: Vec<MatchCandidate> =
715 unsafe { superfluid_abi::array::read_array(&paged.candidates, &bounds) }
716 .unwrap()
717 .collect();
718 assert!(cands.iter().any(|c| c.prefix_len == n), "candidates {:?} carry the {n}-token prefix", cands.iter().map(|c| c.prefix_len).collect::<Vec<_>>());
719
720 let handle = h
721 .engine
722 .seed_acquire(&longer, n, determinism::BEST_EFFORT)
723 .unwrap();
724
725 let p2 = h.prompt(&longer);
726 let ev = h.tick_ok(
727 h.plan()
728 .admit_with(
729 |mut a| {
730 a.seed_handle = handle;
731 a
732 },
733 2,
734 p2,
735 )
736 .prefill(2, n as u32, 1)
737 .decode(2, 1),
738 );
739 assert_eq!(ev.admit_for(2).status, admit_status::ADMITTED);
740 assert_eq!(ev.emit_for(2).n_tokens, 1);
741 assert_eq!(h.engine.lane_committed(2).unwrap().len(), n as usize + 2);
742
743 let p3 = h.prompt(&longer);
744 let err = h.tick_err(h.plan().admit_with(
745 |mut a| {
746 a.seed_handle = handle;
747 a
748 },
749 3,
750 p3,
751 ));
752 assert_eq!(err, Status::RejectStaleSeed);
753 }
754
755 fn gcd(a: u64, b: u64) -> u64 {
756 if b == 0 { a } else { gcd(b, a % b) }
757 }
758
759 fn lcm(a: u64, b: u64) -> u64 {
760 a / gcd(a, b) * b
761 }
762
763 pub fn events_echo_plan_seq<E: Engine, R: StageRings>(mut h: Harness<E, R>) {
764 let p = h.prompt(&[1]);
765 let ev = h.tick_ok(h.plan().admit(9, p).prefill(9, 0, 1));
766 assert_eq!(ev.plan_seq, 1);
767 let ev = h.tick_ok(h.plan());
768 assert_eq!(ev.plan_seq, 2);
769 }
770}