metrale_model_engine/model/trait_impl/
impl_device_feed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `impl ModelDeviceFeed for TransformerModel`, mostly delegating to `<method>_dispatch`
4//! helpers in the sibling modules.
5//!
6//! Owner: model-engine.
7//! Invariants: the ones in `trait_impl/mod.rs`.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::DevicePtr;
11
12use crate::model::types::TransformerModel;
13use crate::traits::{FeedSource, ModelDeviceFeed, RowMask, SequenceState};
14
15impl ModelDeviceFeed for TransformerModel {
16    // 2026-09-25: Device token feed. The `_dispatch` helpers down to
17    // `fed_step_settled_dispatch` are defined in `feed.rs`.
18    fn supports_device_token_feed(&self) -> bool {
19        self.supports_device_token_feed_dispatch()
20    }
21
22    fn decode_batch_fed(
23        &self,
24        sources: &[FeedSource],
25        seqs: &mut [&mut SequenceState],
26        stream: u64,
27    ) -> Result<DevicePtr> {
28        self.gdn_carry_flush_pending()?;
29        self.decode_batch_fed_dispatch(sources, seqs, stream)
30    }
31
32    fn argmax_batch_to_feed(
33        &self,
34        logits: DevicePtr,
35        masks: &[RowMask],
36        dst: *mut u32,
37        event: u64,
38        stream: u64,
39    ) -> Result<()> {
40        self.argmax_batch_to_feed_dispatch(logits, masks, dst, event, stream)
41    }
42
43    fn reserve_decode_block(&self, seq: &mut SequenceState) -> Result<usize> {
44        self.reserve_decode_block_dispatch(seq)
45    }
46
47    fn release_decode_blocks(&self, seq: &mut SequenceState, blocks: usize) -> Result<()> {
48        self.release_decode_blocks_dispatch(seq, blocks)
49    }
50
51    fn fed_step_settled(&self) {
52        self.fed_step_settled_dispatch()
53    }
54
55    fn event_query(&self, event: u64) -> Result<bool> {
56        self.gpu.event_query(event)
57    }
58
59    fn event_synchronize(&self, event: u64) -> Result<()> {
60        self.gpu.event_synchronize(event)
61    }
62
63    fn destroy_event(&self, event: u64) -> Result<()> {
64        self.gpu.destroy_event(event)
65    }
66
67    fn copy_d2h_async(&self, src: DevicePtr, dst: &mut [u8], stream: u64) -> Result<()> {
68        self.gpu.copy_d2h_async(src, dst, stream)
69    }
70
71    fn alloc_host_pinned(&self, bytes: usize) -> Result<*mut u8> {
72        self.gpu.alloc_host_pinned(bytes)
73    }
74
75    fn free_host_pinned(&self, ptr: *mut u8, bytes: usize) -> Result<()> {
76        self.gpu.free_host_pinned(ptr, bytes)
77    }
78}