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

memra_engine/
sigrouter_contract.rs
