metrale_model_engine/traits/model/
logits.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelLogits`, 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 anyhow::Result;
10use metrale_gpu_runtime::gpu::DevicePtr;
11
12/// 2026-09-26: Reading a forward pass's output: the logits buffers, device argmax and the post-norm
13/// hidden state.
14pub trait ModelLogits {
15    fn vocab_size(&self) -> usize;
16
17    /// 2026-09-25: Copy one row of logits to the host; `dst` is `vocab_size * 2` bytes for BF16
18    /// logits.
19    fn copy_logits_to_host(&self, logits_ptr: DevicePtr, dst: &mut [u8]) -> Result<()>;
20
21    /// 2026-09-25: Whether `logits_ptr` holds FP32 logits (`vocab * 4` bytes per row). Default
22    /// `false`; `TransformerModel` returns `false` too, because its `use_fp32_logits` is fixed to
23    /// `false` in `TransformerModel::new`.
24    fn logits_ptr_is_fp32(&self, _logits_ptr: DevicePtr) -> bool {
25        false
26    }
27
28    /// 2026-09-25: The device logits buffer, where a verify pass leaves its `[k, vocab]` rows.
29    fn logits_buffer_ptr(&self) -> DevicePtr;
30
31    /// 2026-09-25: Argmax of one logits row, computed on the device.
32    fn argmax_on_device(&self, logits_ptr: DevicePtr, stream: u64) -> Result<u32>;
33
34    /// 2026-09-25: Argmax of each of `n` logits rows on the device.
35    fn argmax_batch(&self, logits_ptr: DevicePtr, n: usize, stream: u64) -> Result<Vec<u32>>;
36
37    /// 2026-09-25: The post-final-norm hidden state of the last decode step.
38    fn hidden_after_norm(&self) -> DevicePtr;
39
40    /// 2026-09-25: Whether single-token decode writes FP32 logits, read from
41    /// [`Self::decode_logits_ptr`] at 4 bytes per element. Default `false`; `TransformerModel`
42    /// returns its `use_fp32_logits`, which is always `false`.
43    fn decode_logits_fp32(&self) -> bool {
44        false
45    }
46
47    /// 2026-09-25: The buffer single-token decode last wrote its logits to: FP32 when
48    /// [`Self::decode_logits_fp32`] is true, else BF16. The default panics, so a model must
49    /// override it; `TransformerModel` does.
50    fn decode_logits_ptr(&self) -> DevicePtr {
51        unreachable!(
52            "Model::decode_logits_ptr() must be overridden alongside \
53             decode_logits_fp32() — default cannot return a valid pointer."
54        )
55    }
56}
metrale_model_engine/traits/model/logits.rs

metrale_model_engine/traits/model/
logits.rs
