memra_engine/
sigrouter_contract.rs

1//! Cross-surface contracts for DeepSeek-class sigmoid routing.
2
3use std::sync::OnceLock;
4
5const SERVED_LOGIT_MAGIC: &[u8; 8] = b"MSIGRPL1";
6
7pub struct ServedLogitRecord {
8    pub layer: u32,
9    pub tokens: usize,
10    pub n_expert: usize,
11    pub n_used: usize,
12    pub scaling_factor: f32,
13    pub route_norm: bool,
14    pub active: Vec<u8>,
15    pub bias: Vec<f32>,
16    pub logits: Vec<f32>,
17}
18
19struct ServedLogitWriter {
20    path: std::path::PathBuf,
21    file: std::fs::File,
22    seen_layers: std::collections::HashSet<u32>,
23}
24
25static SERVED_LOGIT_WRITER: OnceLock<std::sync::Mutex<Option<ServedLogitWriter>>> =
26    OnceLock::new();
27
28/// Reject an undersubscribed active set before sorting, slicing, or launching CUDA work.
29pub fn validate_active_count(n_used: usize, active_count: usize) -> Result<(), String> {
30    if active_count < n_used {
31        return Err(format!(
32            "sigmoid router requires active_count >= n_used: active_count={active_count}, n_used={n_used}",
33        ));
34    }
35    Ok(())
36}
37
38/// Probe the runtime scalar expf against the bit patterns that froze the host/device oracle.
39pub fn verify_host_expf() -> Result<(), String> {
40    static RESULT: OnceLock<Result<(), String>> = OnceLock::new();
41    RESULT
42        .get_or_init(|| {
43            const CASES: &[(u32, u32)] = &[
44                (0xc0f4_bca4, 0x39fa_13b9), (0x4152_4c00, 0x48f9_5e97),
45                (0xc100_cdee, 0x39a7_4153), (0xc07a_9e68, 0x3ca3_33fa),
46                (0x412b_32ca, 0x472d_3f68), (0xc153_f532, 0x35ec_e4cd),
47                (0x4031_a0d8, 0x4180_5da2), (0xc0e3_fbc0, 0x3a53_10be),
48                (0x4141_525e, 0x482c_a0b2), (0x40f0_a9a8, 0x44e6_bc14),
49                (0xc158_4cca, 0x35b4_96ea), (0xc123_1dda, 0x381c_b7e3),
50                (0x3f72_ed60, 0x4025_4f28), (0x4179_f55a, 0x4ab9_e5a0),
51                (0x415c_15c0, 0x4965_e07c), (0x3fa0_5800, 0x405f_fb8f),
52                (0x416d_e91a, 0x4a2f_193b), (0x417c_082a, 0x4ad3_9e5b),
53                (0x4133_a8f6, 0x4792_ffbd), (0xc17b_4c4a, 0x3422_1cc9),
54                (0x4049_f238, 0x41bb_b376), (0x4122_c18a, 0x46cc_6dd0),
55                (0x40c3_8904, 0x43e1_46c6), (0x411f_3fd2, 0x46a4_31c2),
56            ];
57            for (case, &(input_bits, expected_bits)) in CASES.iter().enumerate() {
58                let input = std::hint::black_box(f32::from_bits(input_bits));
59                let actual_bits = input.exp().to_bits();
60                if actual_bits != expected_bits {
61                    return Err(format!(
62                        "host expf byte probe mismatch at case {case}: input=0x{input_bits:08x}, expected=0x{expected_bits:08x}, actual=0x{actual_bits:08x}",
63                    ));
64                }
65            }
66            Ok(())
67        })
68        .clone()
69}
70
71pub fn served_logit_trace_enabled() -> bool {
72    std::env::var_os("MEMRA_SIG_ROUTER_LOGIT_TRACE").is_some()
73}
74
75/// Persist the first real decode router row for each layer. All floats are stored as raw f32 bits,
76/// so replay tests the exact served inputs rather than a decimal serialization of them.
77#[allow(clippy::too_many_arguments)]
78pub fn capture_served_logits(
79    layer: u32,
80    tokens: usize,
81    n_expert: usize,
82    n_used: usize,
83    scaling_factor: f32,
84    route_norm: bool,
85    active: &[u8],
86    bias: &[f32],
87    logits: &[f32],
88) -> Result<(), String> {
89    use std::io::Write as _;
90
91    let Some(path) = std::env::var_os("MEMRA_SIG_ROUTER_LOGIT_TRACE") else {
92        return Ok(());
93    };
94    if tokens != 1 {
95        return Ok(());
96    }
97    if active.len() != n_expert || bias.len() != n_expert || logits.len() < n_expert {
98        return Err(format!(
99            "served sigmoid-logit trace shape mismatch at layer {layer}: active={} bias={} logits={} n_expert={n_expert}",
100            active.len(), bias.len(), logits.len(),
101        ));
102    }
103    let path = std::path::PathBuf::from(path);
104    let state = SERVED_LOGIT_WRITER.get_or_init(|| std::sync::Mutex::new(None));
105    let mut state = state
106        .lock()
107        .map_err(|_| "served sigmoid-logit trace writer lock is poisoned".to_string())?;
108    if state.is_none() {
109        let mut file = std::fs::OpenOptions::new()
110            .create(true)
111            .truncate(true)
112            .write(true)
113            .open(&path)
114            .map_err(|error| format!("cannot create served sigmoid-logit trace {}: {error}", path.display()))?;
115        file.write_all(SERVED_LOGIT_MAGIC)
116            .map_err(|error| format!("cannot write served sigmoid-logit trace header: {error}"))?;
117        *state = Some(ServedLogitWriter {
118            path: path.clone(),
119            file,
120            seen_layers: std::collections::HashSet::new(),
121        });
122    }
123    let writer = state.as_mut().unwrap();
124    if writer.path != path {
125        return Err("MEMRA_SIG_ROUTER_LOGIT_TRACE changed after capture started".into());
126    }
127    if !writer.seen_layers.insert(layer) {
128        return Ok(());
129    }
130
131    for value in [
132        layer,
133        tokens as u32,
134        n_expert as u32,
135        n_used as u32,
136        scaling_factor.to_bits(),
137        u32::from(route_norm),
138    ] {
139        writer.file.write_all(&value.to_le_bytes())
140            .map_err(|error| format!("cannot write served sigmoid-logit trace row: {error}"))?;
141    }
142    writer.file.write_all(active)
143        .map_err(|error| format!("cannot write served sigmoid-logit active mask: {error}"))?;
144    for value in bias.iter().chain(logits[..n_expert].iter()) {
145        writer.file.write_all(&value.to_bits().to_le_bytes())
146            .map_err(|error| format!("cannot write served sigmoid-logit f32 row: {error}"))?;
147    }
148    writer.file.flush()
149        .map_err(|error| format!("cannot flush served sigmoid-logit trace: {error}"))?;
150    Ok(())
151}
152
153fn read_u32(bytes: &[u8], cursor: &mut usize) -> Result<u32, String> {
154    let end = cursor.saturating_add(4);
155    let raw: [u8; 4] = bytes
156        .get(*cursor..end)
157        .ok_or_else(|| "truncated served sigmoid-logit trace".to_string())?
158        .try_into()
159        .unwrap();
160    *cursor = end;
161    Ok(u32::from_le_bytes(raw))
162}
163
164pub fn read_served_logits(path: &std::path::Path) -> Result<Vec<ServedLogitRecord>, String> {
165    let bytes = std::fs::read(path)
166        .map_err(|error| format!("cannot read served sigmoid-logit trace {}: {error}", path.display()))?;
167    if bytes.get(..SERVED_LOGIT_MAGIC.len()) != Some(SERVED_LOGIT_MAGIC) {
168        return Err("served sigmoid-logit trace has wrong or missing v1 header".into());
169    }
170    let mut cursor = SERVED_LOGIT_MAGIC.len();
171    let mut records = Vec::new();
172    while cursor < bytes.len() {
173        let layer = read_u32(&bytes, &mut cursor)?;
174        let tokens = read_u32(&bytes, &mut cursor)? as usize;
175        let n_expert = read_u32(&bytes, &mut cursor)? as usize;
176        let n_used = read_u32(&bytes, &mut cursor)? as usize;
177        let scaling_factor = f32::from_bits(read_u32(&bytes, &mut cursor)?);
178        let route_norm = match read_u32(&bytes, &mut cursor)? {
179            0 => false,
180            1 => true,
181            value => return Err(format!("invalid route_norm={value} in served sigmoid-logit trace")),
182        };
183        if tokens != 1 || n_expert == 0 || n_used == 0 || n_expert > 1024 {
184            return Err(format!(
185                "invalid served sigmoid-logit record shape: layer={layer} tokens={tokens} n_expert={n_expert} n_used={n_used}",
186            ));
187        }
188        let active_end = cursor.saturating_add(n_expert);
189        let active = bytes
190            .get(cursor..active_end)
191            .ok_or_else(|| "truncated served sigmoid-logit active mask".to_string())?
192            .to_vec();
193        cursor = active_end;
194        let mut read_f32_row = || -> Result<Vec<f32>, String> {
195            (0..n_expert)
196                .map(|_| read_u32(&bytes, &mut cursor).map(f32::from_bits))
197                .collect()
198        };
199        let bias = read_f32_row()?;
200        let logits = read_f32_row()?;
201        records.push(ServedLogitRecord {
202            layer,
203            tokens,
204            n_expert,
205            n_used,
206            scaling_factor,
207            route_norm,
208            active,
209            bias,
210            logits,
211        });
212    }
213    Ok(records)
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn active_count_error_quotes_both_counts() {
222        assert_eq!(
223            validate_active_count(8, 7).unwrap_err(),
224            "sigmoid router requires active_count >= n_used: active_count=7, n_used=8",
225        );
226        assert!(validate_active_count(8, 8).is_ok());
227    }
228
229    #[test]
230    fn pinned_host_expf_matches() {
231        verify_host_expf().unwrap();
232    }
233}
memra_engine/sigrouter_contract.rs

memra_engine/
sigrouter_contract.rs
