rlx_ir/
variant.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Model execution variants — one object drives cache keys and [`DimBinding`].
6//!
7//! Mirrors the “shader components” idea from extensible shading systems: the same
8//! granularity selects **what to specialize** and **which symbolic dims to bind**.
9
10use std::collections::hash_map::DefaultHasher;
11use std::hash::{Hash, Hasher};
12
13use crate::dynamic::sym;
14use crate::shape::DimBinding;
15
16/// Coarse execution phase (prefill vs decode vs encoder).
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum ModelPhase {
19    Prefill,
20    Decode,
21    Encoder,
22    Inference,
23}
24
25/// Concrete shape bucket for compile-once / specialize-at-runtime workflows.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct ModelVariant {
28    pub batch: usize,
29    pub seq: usize,
30    pub past_seq: Option<usize>,
31    pub phase: ModelPhase,
32    /// Extra dynamic symbols beyond batch/seq/past (e.g. custom ragged axes).
33    pub extra: Vec<(u32, usize)>,
34}
35
36impl ModelVariant {
37    pub fn prefill(batch: usize, seq: usize) -> Self {
38        Self {
39            batch,
40            seq,
41            past_seq: None,
42            phase: ModelPhase::Prefill,
43            extra: Vec::new(),
44        }
45    }
46
47    /// Single-step decode: `seq` is the new token count (often 1); `past_seq` is KV length.
48    pub fn decode(batch: usize, past_seq: usize, new_tokens: usize) -> Self {
49        Self {
50            batch,
51            seq: new_tokens,
52            past_seq: Some(past_seq),
53            phase: ModelPhase::Decode,
54            extra: Vec::new(),
55        }
56    }
57
58    pub fn encoder(batch: usize, seq: usize) -> Self {
59        Self {
60            batch,
61            seq,
62            past_seq: None,
63            phase: ModelPhase::Encoder,
64            extra: Vec::new(),
65        }
66    }
67
68    pub fn with_extra(mut self, symbol: u32, size: usize) -> Self {
69        self.extra.push((symbol, size));
70        self
71    }
72
73    /// Stable cache key: phase + bound leading dims + extra symbols.
74    pub fn cache_key(&self) -> u64 {
75        let mut h = DefaultHasher::new();
76        self.phase.hash(&mut h);
77        self.batch.hash(&mut h);
78        self.seq.hash(&mut h);
79        self.past_seq.hash(&mut h);
80        for (sym, size) in &self.extra {
81            sym.hash(&mut h);
82            size.hash(&mut h);
83        }
84        h.finish()
85    }
86
87    /// Symbol bindings used by [`crate::dynamic::bind_graph`] / compile specialization.
88    pub fn dim_binding(&self) -> DimBinding {
89        let mut b = match (self.phase, self.past_seq) {
90            (ModelPhase::Decode, Some(past)) => DimBinding::batch_past_seq(self.batch, past),
91            _ => DimBinding::batch_seq(self.batch, self.seq),
92        };
93        if self.phase == ModelPhase::Decode {
94            b.set(sym::SEQ, self.seq);
95        }
96        for (sym, size) in &self.extra {
97            b.set(*sym, *size);
98        }
99        b
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn prefill_binding_sets_batch_seq() {
109        let v = ModelVariant::prefill(2, 128);
110        let b = v.dim_binding();
111        assert_eq!(b.get(sym::BATCH), Some(2));
112        assert_eq!(b.get(sym::SEQ), Some(128));
113    }
114
115    #[test]
116    fn decode_binding_sets_past_and_new_seq() {
117        let v = ModelVariant::decode(1, 64, 1);
118        let b = v.dim_binding();
119        assert_eq!(b.get(sym::BATCH), Some(1));
120        assert_eq!(b.get(sym::PAST_SEQ), Some(64));
121        assert_eq!(b.get(sym::SEQ), Some(1));
122    }
123
124    #[test]
125    fn cache_key_differs_by_phase() {
126        let a = ModelVariant::prefill(1, 8).cache_key();
127        let b = ModelVariant::decode(1, 7, 1).cache_key();
128        assert_ne!(a, b);
129    }
130}
rlx_ir/variant.rs

rlx_ir/
variant.rs
