metrale_model_engine/traits/model/
device_feed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelDeviceFeed`, one of the supertraits `Model` is made of. Its methods, default
4//! bodies and docs are the ones `Model` declared before the split.
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9use super::{FeedSource, RowMask};
10use crate::traits::SequenceState;
11use anyhow::{Result, bail};
12use metrale_gpu_runtime::gpu::DevicePtr;
13
14/// 2026-09-26: The device token feed of the asynchronous scheduler router, with the event queries,
15/// async copies and pinned host memory it uses.
16pub trait ModelDeviceFeed {
17    // 2026-09-25: Device token feed: what the asynchronous scheduler router needs to run a
18    // plain decode step ahead of the host. `AsyncDeviceIo::new` refuses a model whose
19    // `supports_device_token_feed` is `false`, the default.
20
21    /// 2026-09-25: Whether the model can run a decode step whose input ids live on the device
22    /// (`decode_batch_fed`) and answer it with a masked argmax written into the feed cells
23    /// (`argmax_batch_to_feed`). Default `false`.
24    fn supports_device_token_feed(&self) -> bool {
25        false
26    }
27
28    /// 2026-09-25: The decode step over `seqs` with each row's input id read on the device from
29    /// `sources`: the forward, KV writes and block allocation of `decode_batch`, but it does not
30    /// push the token or advance `seq_len`, because the router does not know a fed row's id at
31    /// launch; the scheduler core does that once the previous step is committed. Position
32    /// `seq.seq_len` is written. Default: an error.
33    fn decode_batch_fed(
34        &self,
35        _sources: &[FeedSource],
36        _seqs: &mut [&mut SequenceState],
37        _stream: u64,
38    ) -> Result<DevicePtr> {
39        bail!("decode_batch_fed: this model has no device token feed")
40    }
41
42    /// 2026-09-25: Argmax over `masks.len()` rows of `logits`, applying each row's two-id mask
43    /// (`argmax_feed.cu`): writes the feed cells, copies the ids to the pinned `dst` with an
44    /// async D2H on `stream`, then records `event`. It never synchronises; `dst` is valid once
45    /// `event` completes. Default: an error.
46    fn argmax_batch_to_feed(
47        &self,
48        _logits: DevicePtr,
49        _masks: &[RowMask],
50        _dst: *mut u32,
51        _event: u64,
52        _stream: u64,
53    ) -> Result<()> {
54        bail!("argmax_batch_to_feed: this model has no device token feed")
55    }
56
57    /// 2026-09-25: Allocate the KV block for the next decode position (`seq.seq_len`) if the
58    /// block table lacks it; returns the blocks added (0 or 1). `Err` when the pool is empty.
59    /// Default: an error.
60    fn reserve_decode_block(&self, _seq: &mut SequenceState) -> Result<usize> {
61        bail!("reserve_decode_block: this model has no device token feed")
62    }
63
64    /// 2026-09-25: Return the last `blocks` entries of the block table, added by
65    /// `reserve_decode_block` for a discarded step, to the pool. Default: an error.
66    fn release_decode_blocks(&self, _seq: &mut SequenceState, _blocks: usize) -> Result<()> {
67        bail!("release_decode_blocks: this model has no device token feed")
68    }
69
70    /// 2026-09-25: A fed step has completed on the device, so the CUDA graph it replayed may be
71    /// evicted. Default: no-op.
72    fn fed_step_settled(&self) {}
73
74    /// 2026-09-25: Whether the work recorded against `event` has completed; never blocks.
75    /// Default `Ok(true)`.
76    fn event_query(&self, _event: u64) -> Result<bool> {
77        Ok(true)
78    }
79
80    /// 2026-09-25: Block the host until the work recorded against `event` has completed.
81    /// Default: `Ok(())`.
82    fn event_synchronize(&self, _event: u64) -> Result<()> {
83        Ok(())
84    }
85
86    fn destroy_event(&self, _event: u64) -> Result<()> {
87        Ok(())
88    }
89
90    /// 2026-09-25: Async device-to-host copy on `stream`; `dst` must stay valid and unread until
91    /// `stream` is synchronised. Default: an error.
92    fn copy_d2h_async(&self, _src: DevicePtr, _dst: &mut [u8], _stream: u64) -> Result<()> {
93        bail!("copy_d2h_async: this model has no device token feed")
94    }
95
96    /// 2026-09-25: Zeroed page-locked host memory for the router's readback ring. Default: an
97    /// error.
98    fn alloc_host_pinned(&self, _bytes: usize) -> Result<*mut u8> {
99        bail!("alloc_host_pinned: this model has no device token feed")
100    }
101
102    fn free_host_pinned(&self, _ptr: *mut u8, _bytes: usize) -> Result<()> {
103        bail!("free_host_pinned: this model has no device token feed")
104    }
105}
metrale_model_engine/traits/model/device_feed.rs

metrale_model_engine/traits/model/
device_feed.rs
