ferrox_models/
execution_plan.rs

1//! Load-time execution / memory plans (llama.cpp graph-params analogue).
2//!
3//! Selected once per model (and cached by batch geometry for decode vs
4//! prefill). Hot-path forward never re-derives architecture strings or
5//! fused-op availability.
6
7use crate::capability::{DecoderFamily, MemoryKind, QkNormStyle};
8use crate::config::{FfnActivation, RopeLayout};
9
10/// Backend fused-op availability discovered at load / first probe.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
12pub struct FusedOpCaps {
13    pub metal_flash_attn: bool,
14    pub metal_swiglu: bool,
15    pub metal_matvec: bool,
16    /// CUDA is deferred — kept for ABI stability, always false for now.
17    pub cuda_gqa: bool,
18}
19
20/// Memory layout chosen once from the architecture profile.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct MemoryPlan {
23    pub kind: MemoryKind,
24    pub swa_pattern: Option<usize>,
25    pub sliding_window: Option<usize>,
26}
27
28/// Per-model execution plan: everything the forward path needs that is
29/// constant across tokens.
30#[derive(Debug, Clone, PartialEq)]
31pub struct ExecutionPlan {
32    pub family: DecoderFamily,
33    pub rope: RopeLayout,
34    pub qk_norm: QkNormStyle,
35    pub ffn_activation: FfnActivation,
36    pub memory: MemoryPlan,
37    pub fused: FusedOpCaps,
38    pub embedding_scale: Option<f32>,
39    pub attention_scale: Option<f32>,
40    pub rope_theta_swa: Option<f32>,
41    pub attn_logit_softcap: Option<f32>,
42    pub final_logit_softcap: Option<f32>,
43}
44
45/// Cache key for decode/prefill plan reuse (batch geometry only).
46#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
47pub struct PlanGeometry {
48    pub n_tokens: usize,
49    pub n_seqs: usize,
50    pub flash_attn: bool,
51}
52
53impl ExecutionPlan {
54    /// Build from a resolved [`crate::config::ModelConfig`] + profile.
55    pub fn from_config(
56        config: &crate::config::ModelConfig,
57        family: DecoderFamily,
58        memory_kind: MemoryKind,
59        fused: FusedOpCaps,
60    ) -> Self {
61        Self {
62            family,
63            rope: config.rope_layout,
64            qk_norm: config.qk_norm_style,
65            ffn_activation: config.ffn_activation,
66            memory: MemoryPlan {
67                kind: memory_kind,
68                swa_pattern: config.swa_pattern,
69                sliding_window: config.sliding_window,
70            },
71            fused,
72            embedding_scale: config.embedding_scale,
73            attention_scale: config.attention_scale,
74            rope_theta_swa: config.rope_theta_swa,
75            attn_logit_softcap: config.attn_logit_softcap,
76            final_logit_softcap: config.final_logit_softcap,
77        }
78    }
79
80    /// Look up or insert fused-op caps for a decode/prefill geometry.
81    pub fn caps_for_geometry(
82        &self,
83        cache: &mut std::collections::HashMap<PlanGeometry, FusedOpCaps>,
84        geom: PlanGeometry,
85    ) -> FusedOpCaps {
86        *cache.entry(geom).or_insert(self.fused)
87    }
88
89    /// Probe Metal fused-op availability without changing the Llama
90    /// default path. Returns conservative caps when Metal is off.
91    pub fn probe_metal_caps() -> FusedOpCaps {
92        #[cfg(feature = "metal")]
93        {
94            let metal_on = ferrox_core::metal_dense_enabled();
95            FusedOpCaps {
96                metal_flash_attn: metal_on && ferrox_metal::attn::metal_attn_enabled(),
97                metal_swiglu: metal_on,
98                metal_matvec: metal_on,
99                cuda_gqa: false,
100            }
101        }
102        #[cfg(not(feature = "metal"))]
103        {
104            FusedOpCaps::default()
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::config::test_dense_fixture;
113
114    #[test]
115    fn plan_from_tiny_config_defaults() {
116        let cfg = test_dense_fixture();
117        let plan = ExecutionPlan::from_config(
118            &cfg,
119            DecoderFamily::StandardGqa,
120            MemoryKind::KvGqa,
121            FusedOpCaps::default(),
122        );
123        assert_eq!(plan.family, DecoderFamily::StandardGqa);
124        assert_eq!(plan.qk_norm, QkNormStyle::WholeVector);
125        assert!(!plan.fused.cuda_gqa);
126    }
127}
ferrox_models/execution_plan.rs

ferrox_models/
execution_plan.rs
