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    /// Probe Metal fused-op availability without changing the Llama
81    /// default path. Returns conservative caps when Metal is off.
82    pub fn probe_metal_caps() -> FusedOpCaps {
83        #[cfg(feature = "metal")]
84        {
85            let metal_on = ferrox_core::metal_dense_enabled();
86            FusedOpCaps {
87                metal_flash_attn: metal_on && ferrox_metal::attn::metal_attn_enabled(),
88                metal_swiglu: metal_on,
89                metal_matvec: metal_on,
90                cuda_gqa: false,
91            }
92        }
93        #[cfg(not(feature = "metal"))]
94        {
95            FusedOpCaps::default()
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::config::test_dense_fixture;
104
105    #[test]
106    fn plan_from_tiny_config_defaults() {
107        let cfg = test_dense_fixture();
108        let plan = ExecutionPlan::from_config(
109            &cfg,
110            DecoderFamily::StandardGqa,
111            MemoryKind::KvGqa,
112            FusedOpCaps::default(),
113        );
114        assert_eq!(plan.family, DecoderFamily::StandardGqa);
115        assert_eq!(plan.qk_norm, QkNormStyle::WholeVector);
116        assert!(!plan.fused.cuda_gqa);
117    }
118}
ferrox_models/execution_plan.rs

ferrox_models/
execution_plan.rs
