kopitiam_runtime/
traits.rs

1//! The Model Runtime boundary: [`Model`] and [`Backend`].
2//!
3//! Everything above this line ([`crate::generate`], and eventually
4//! `kopitiam-ai`'s real `ModelAdapter` — see this crate's parent epic) is
5//! written against these two traits, not against [`crate::model::QwenModel`]
6//! directly, the same way `kopitiam-workflow` is written against
7//! `kopitiam_ai::ModelAdapter` rather than any one concrete adapter. That
8//! mirroring is deliberate: it is the same "own the boundary, not the
9//! implementation" shape CLAUDE.md's Architecture section asks every layer
10//! of this platform to follow.
11
12use kopitiam_core::{Device, Result};
13use kopitiam_tensor::Tensor;
14
15use crate::kv_cache::KvCache;
16
17/// A causal (decoder-only) language model that can run a forward pass and
18/// produce next-token logits.
19///
20/// # Why this trait exists with exactly one implementation
21///
22/// [`crate::model::QwenModel`] is the only [`Model`] today. That is not
23/// reason enough on its own to add a trait — CLAUDE.md is explicit that
24/// unused abstraction is a cost, not a virtue — but this one earns its
25/// keep on two counts: [`crate::generate`] is written against `Model`, not
26/// `QwenModel`, so it does not have to change when a second architecture
27/// (a non-GQA LLaMA checkpoint, say, or a future non-transformer
28/// architecture) arrives; and it marks precisely the seam a second
29/// architecture would implement against, the same role
30/// [`kopitiam_core::Device`]'s single `Cpu` variant plays for backends (see
31/// that type's docs, which this trait's shape deliberately echoes).
32pub trait Model {
33    /// Runs a forward pass over `token_ids` — new tokens only (a multi-token
34    /// prompt on the first call, ordinarily one token per call thereafter)
35    /// — using and extending `cache`. Returns logits shaped
36    /// `[token_ids.len(), vocab_size]`: one row of unnormalized
37    /// per-vocabulary scores for every input position, in the same order as
38    /// `token_ids`.
39    fn forward(&self, token_ids: &[u32], cache: &mut KvCache) -> Result<Tensor>;
40
41    /// The vocabulary size logits' last dimension is sized to.
42    fn vocab_size(&self) -> usize;
43
44    /// The context window `cache` should be constructed with via
45    /// [`KvCache::new`] to run this model without hitting
46    /// [`kopitiam_core::Error::IndexOutOfBounds`] from
47    /// [`KvCache::append`](crate::kv_cache::KvCache::append).
48    fn max_context(&self) -> usize;
49
50    /// Builds a fresh, empty [`KvCache`] sized correctly for this model
51    /// (layer count and context window). This is the only correct way to
52    /// construct a cache for a given model — [`KvCache::new`] takes those
53    /// two numbers as bare `usize`s and trusts the caller to get them
54    /// right, which is fine for `KvCache`'s own unit tests but exactly the
55    /// kind of easy-to-mismatch call [`crate::generate::generate`] (generic
56    /// over any [`Model`], not just [`crate::model::QwenModel`]) should
57    /// never have to get right by hand.
58    fn new_cache(&self) -> KvCache;
59}
60
61/// The execution backend a [`Model`] runs its tensor ops on.
62///
63/// # Why this trait exists with exactly one implementation
64///
65/// Today every [`Model`] impl computes directly against
66/// `kopitiam_tensor::Tensor`'s plain-CPU `f32` kernels — there is no
67/// dispatch through this trait yet, and adding one purely to satisfy this
68/// task brief's ask for a `Backend` seam, without a second backend that
69/// would use it, is exactly the unnecessary abstraction CLAUDE.md warns
70/// against. So this trait is deliberately minimal: it names the seam (what
71/// a caller would ask a backend for — its [`Device`]) without inventing a
72/// dispatch mechanism no code calls yet. A real second backend (a
73/// SIMD-tuned kernel set, a threaded scheduler — see the parent epic's
74/// Phase 2/3) is expected to grow this trait's surface *and* wire `Model`
75/// impls to route their tensor ops through it; until then, [`CpuBackend`]
76/// exists so [`crate::model::QwenModel`] has something concrete to report
77/// through [`Model`]'s eventual `device()` accessor rather than hardcoding
78/// [`Device::Cpu`] as a bare literal at every call site.
79pub trait Backend {
80    /// Where this backend's tensors live and compute.
81    fn device(&self) -> Device;
82}
83
84/// The only [`Backend`] the Kopitiam Runtime implements: plain CPU
85/// execution via `kopitiam-tensor`'s `f32` kernels. See [`Backend`]'s docs
86/// for why this is a real (if minimal) type rather than a bare `Device`
87/// literal.
88#[derive(Debug, Clone, Copy, Default)]
89pub struct CpuBackend;
90
91impl Backend for CpuBackend {
92    fn device(&self) -> Device {
93        Device::Cpu
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn cpu_backend_reports_the_cpu_device() {
103        assert_eq!(CpuBackend.device(), Device::Cpu);
104    }
105}
kopitiam_runtime/traits.rs

kopitiam_runtime/
traits.rs
