metrale_model_engine/model/trait_impl/
impl_logits.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `impl ModelLogits 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::ModelLogits;
14
15impl ModelLogits for TransformerModel {
16    fn vocab_size(&self) -> usize {
17        self.vocab_size_dispatch()
18    }
19
20    fn copy_logits_to_host(&self, logits_ptr: DevicePtr, dst: &mut [u8]) -> Result<()> {
21        self.copy_logits_to_host_dispatch(logits_ptr, dst)
22    }
23
24    fn logits_ptr_is_fp32(&self, logits_ptr: DevicePtr) -> bool {
25        self.logits_ptr_is_fp32_dispatch(logits_ptr)
26    }
27
28    fn logits_buffer_ptr(&self) -> DevicePtr {
29        self.logits_buffer_ptr_dispatch()
30    }
31
32    fn argmax_on_device(&self, logits_ptr: DevicePtr, _stream: u64) -> Result<u32> {
33        self.argmax_on_device_dispatch(logits_ptr, _stream)
34    }
35
36    fn argmax_batch(&self, logits_ptr: DevicePtr, n: usize, _stream: u64) -> Result<Vec<u32>> {
37        self.argmax_batch_dispatch(logits_ptr, n, _stream)
38    }
39
40    fn hidden_after_norm(&self) -> DevicePtr {
41        self.hidden_after_norm_dispatch()
42    }
43
44    fn decode_logits_fp32(&self) -> bool {
45        self.decode_logits_fp32_dispatch()
46    }
47
48    fn decode_logits_ptr(&self) -> DevicePtr {
49        self.decode_logits_ptr_dispatch()
50    }
51}
metrale_model_engine/model/trait_impl/impl_logits.rs

metrale_model_engine/model/trait_impl/
impl_logits.rs
