metrale_model_engine/model/trait_impl/
impl_device_feed.rs1use 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 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}