Skip to main content

keyhog_scanner/gpu/
self_test.rs

1/// Result from an explicit GPU adapter and dispatch self-test.
2#[derive(Debug, Clone, PartialEq, Eq)]
3pub struct GpuSelfTest {
4    /// Human-readable adapter name reported by wgpu.
5    pub adapter_name: String,
6    /// Approximate storage-buffer capability in MiB when available.
7    pub vram_mb: Option<u64>,
8    /// Number of scores produced by the compute dispatch.
9    pub scores: usize,
10}
11
12/// Result from an explicit VYRE GPU scanner self-test.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct VyreGpuSelfTest {
15    /// Number of direct GPU matches produced by `GpuLiteralSet::scan`.
16    pub direct_matches: usize,
17    /// Number of matches produced by one coalesced scanner GPU dispatch.
18    pub coalesced_matches: usize,
19}
20
21#[cfg(feature = "gpu")]
22static GPU_SELF_TEST_CACHE: std::sync::OnceLock<std::result::Result<GpuSelfTest, String>> =
23    std::sync::OnceLock::new();
24
25/// Force a GPU compute dispatch and validate the returned scores.
26///
27/// This is stricter than [`crate::gpu::gpu_available`]: it proves that a
28/// non-software wgpu adapter initialized and that the MoE compute shader can run
29/// at least one production-sized batch.
30pub fn gpu_self_test() -> Result<GpuSelfTest, String> {
31    #[cfg(not(feature = "gpu"))]
32    {
33        Err(
34            "GPU support not compiled in (lean ci build). Rebuild with `--features gpu` \
35             (or the default profile) to exercise the wgpu/CUDA path."
36                .to_string(),
37        )
38    }
39    #[cfg(feature = "gpu")]
40    {
41        GPU_SELF_TEST_CACHE
42            .get_or_init(|| {
43                let gpu = super::backend::get_gpu()
44                    .map_err(|error| error.to_string())?
45                    .ok_or_else(|| {
46                        "GPU adapter unavailable; install or enable a non-software GPU adapter and driver"
47                            .to_string()
48                    })?;
49
50                // PARITY, not "in range". The prior check scored ALL-ZERO feature
51                // vectors and only asserted each result was finite and within
52                // [0,1], which a GPU that returns 0.0 for EVERY input trivially
53                // passes. That masked a real shipped fault: the MoE shader scored
54                // genuine secrets ~0.0 (CPU scored them ~1.0), so on a GPU host the
55                // ML gate silently dropped findings and `--self-test` still reported
56                // HEALTHY. Assert the actual contract instead: the GPU MoE must
57                // reproduce the CPU MoE (`ml_scorer::score_features`, the reference
58                // every confidence floor is tuned/benched against) within tolerance
59                // on a probe that includes real secrets. This is the SAME verdict
60                // the scan path enforces (`batch_score_features` fails closed to CPU
61                // on the same divergence), so doctor and the scanner never disagree.
62                let max_abs = super::backend::gpu_moe_parity_max_divergence(
63                    std::time::Duration::from_millis(
64                        crate::scanner_config::ScannerTuningConfig::GPU_MOE_TIMEOUT_MS_DEFAULT,
65                    ),
66                )?;
67                if max_abs > super::backend::GPU_MOE_PARITY_TOLERANCE {
68                    return Err(format!(
69                        "GPU MoE compute shader diverges from the CPU MoE reference by {max_abs:.4} \
70                         (tolerance {:.4}); the GPU would score findings differently from the \
71                         CPU/SIMD path. Indicates a shader miscompile, weights-packing mismatch, \
72                         or driver bug. Scans record the GPU MoE degrade and use the CPU MoE \
73                         (correct + deterministic), \
74                         so detection is unaffected, but GPU ML acceleration is OFF on this host.",
75                        super::backend::GPU_MOE_PARITY_TOLERANCE
76                    ));
77                }
78
79                Ok(GpuSelfTest {
80                    adapter_name: gpu.gpu_name().to_string(),
81                    vram_mb: gpu.vram_mb(),
82                    scores: crate::ml_scorer::GPU_BATCH_THRESHOLD,
83                })
84            })
85            .clone()
86    }
87}
88
89/// Force the VYRE GPU scanner and coalesced scanner paths.
90///
91/// Proves the scanner-side GPU dependency is available independently from
92/// KeyHog's MoE GPU scorer. Both counts are populated from real GPU scans.
93pub fn vyre_gpu_self_test() -> Result<VyreGpuSelfTest, String> {
94    #[cfg(not(feature = "gpu"))]
95    {
96        Err(
97            "VYRE GPU self-test not available in the lean CI build (no WGPU driver compiled in). \
98             Rebuild with `--features gpu`."
99                .to_string(),
100        )
101    }
102    #[cfg(feature = "gpu")]
103    {
104        vyre_gpu_self_test_impl()
105    }
106}
107
108#[cfg(feature = "gpu")]
109fn vyre_gpu_self_test_impl() -> Result<VyreGpuSelfTest, String> {
110    use vyre_driver_wgpu::WgpuBackend;
111    use vyre_libs::scan::GpuLiteralSet;
112
113    let patterns: Vec<Vec<u8>> = vec![b"needle".to_vec()];
114    let pattern_refs: Vec<&[u8]> = patterns.iter().map(Vec::as_slice).collect();
115
116    let backend = WgpuBackend::shared().map_err(|e| format!("failed to init wgpu backend: {e}"))?;
117    let scanner = GpuLiteralSet::compile(&pattern_refs);
118
119    let direct = scanner
120        .scan(backend.as_ref(), b"needle", 100)
121        .map_err(|error| format!("vyre direct GPU scan failed: {error}"))?;
122    if direct.len() != 1 || direct[0].pattern_id != 0 || direct[0].start != 0 {
123        return Err(format!(
124            "vyre direct GPU scan returned unexpected matches: {direct:?}"
125        ));
126    }
127
128    const COALESCED_ITEMS: usize = 100;
129    let items: Vec<Vec<u8>> = (0..COALESCED_ITEMS)
130        .map(|index| format!("id-{index:03}-needle").into_bytes())
131        .collect();
132    let mut buffer = Vec::with_capacity(items.iter().map(Vec::len).sum());
133    for item in &items {
134        buffer.extend_from_slice(item);
135    }
136
137    let coalesced = scanner
138        .scan(backend.as_ref(), &buffer, 10_000)
139        .map_err(|error| format!("vyre coalesced GPU scan failed: {error}"))?;
140    if coalesced.len() != COALESCED_ITEMS {
141        return Err(format!(
142            "vyre coalesced GPU scan returned {} matches, expected {COALESCED_ITEMS}",
143            coalesced.len()
144        ));
145    }
146
147    Ok(VyreGpuSelfTest {
148        direct_matches: direct.len(),
149        coalesced_matches: coalesced.len(),
150    })
151}
152
153/// One acquired peer proven by the production GPU region-presence self-test.
154#[derive(Debug, Clone, PartialEq, Eq)]
155pub struct GpuRegionPresencePeerSelfTest {
156    /// Exact scanner route exercised by the test.
157    pub backend: crate::hw_probe::ScanBackend,
158    /// `VyreBackend::id()` of the driver that ran the test.
159    pub backend_id: &'static str,
160    /// Number of findings emitted through the production GPU trigger path.
161    pub matches: usize,
162}
163
164/// Status report from the production GPU region-presence self-test.
165#[derive(Debug, Clone, PartialEq, Eq)]
166pub struct GpuRegionPresenceSelfTest {
167    /// Every acquired CUDA or WGPU peer. All entries passed exact CPU parity.
168    pub peers: Vec<GpuRegionPresencePeerSelfTest>,
169}
170
171/// Honest aggregate failure from the peer self-test.
172#[derive(Debug, Clone, PartialEq, Eq)]
173pub struct GpuRegionPresenceSelfTestFailure {
174    /// Exact peers acquired before parity execution began.
175    pub acquired_backends: Vec<crate::hw_probe::ScanBackend>,
176    /// Peer-specific acquisition, dispatch, or parity diagnostics.
177    pub message: String,
178}
179
180impl std::fmt::Display for GpuRegionPresenceSelfTestFailure {
181    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
182        formatter.write_str(&self.message)
183    }
184}
185
186impl std::error::Error for GpuRegionPresenceSelfTestFailure {}
187
188/// Build a minimal one-detector `CompiledScanner` and dispatch a scan through
189/// the production GPU backend. A PASS proves device acquisition, compilation,
190/// lowering, dispatch, and host readback on this host.
191pub fn gpu_region_presence_self_test(
192) -> Result<GpuRegionPresenceSelfTest, GpuRegionPresenceSelfTestFailure> {
193    #[cfg(not(feature = "gpu"))]
194    {
195        Err(GpuRegionPresenceSelfTestFailure {
196            acquired_backends: Vec::new(),
197            message: "GPU region-presence self-test not available in the lean ci build. Rebuild with `--features gpu` to exercise the production GPU trigger path.".to_string(),
198        })
199    }
200    #[cfg(feature = "gpu")]
201    {
202        gpu_region_presence_self_test_impl()
203    }
204}
205
206#[cfg(feature = "gpu")]
207fn gpu_region_presence_self_test_impl(
208) -> Result<GpuRegionPresenceSelfTest, GpuRegionPresenceSelfTestFailure> {
209    use crate::engine::CompiledScanner;
210    use crate::hw_probe::ScanBackend;
211    use keyhog_core::{Chunk, ChunkMetadata, DetectorFile};
212
213    // The probe MUST be a credential keyhog actually REPORTS, not one it
214    // suppresses. A plain dictionary word (e.g. "needle") triggers phase-1 and
215    // extracts, but is then correctly dropped by low-entropy/placeholder
216    // suppression on EVERY backend - so asserting "GPU found > 0" on such a word
217    // is a false failure that has nothing to do with the GPU kernel. This probe
218    // mirrors the proven `scan_engine_self_test` shape: a distinctive literal
219    // PREFIX ("KHGPUSELFTEST_") that the GPU literal-set anchors on to drive
220    // region-presence -> trigger, followed by a mixed-case high-entropy suffix
221    // that survives suppression so the match is emitted end to end.
222    const PLANTED: &str = "KHGPUSELFTEST_A1b2C3d4E5f6";
223    let detector =
224        toml::from_str::<DetectorFile>(include_str!("../../data/gpu-self-test-detector.toml"))
225            .map(|file| file.detector)
226            .map_err(|error| GpuRegionPresenceSelfTestFailure {
227                acquired_backends: Vec::new(),
228                message: format!("bundled GPU self-test detector TOML is invalid: {error}"),
229            })?;
230
231    let scanner = CompiledScanner::compile(vec![detector]).map_err(|error| {
232        GpuRegionPresenceSelfTestFailure {
233            acquired_backends: Vec::new(),
234            message: format!("CompiledScanner::compile failed during self-test: {error}"),
235        }
236    })?;
237
238    let candidates = scanner.gpu_backend_candidates();
239    let acquired_backends: Vec<_> = candidates
240        .iter()
241        .filter(|candidate| candidate.is_eligible())
242        .map(|candidate| candidate.backend)
243        .collect();
244    if acquired_backends.is_empty() {
245        let diagnostics = candidates
246            .iter()
247            .map(|candidate| {
248                let diagnostic = match candidate.acquisition_error.as_deref() {
249                    Some(reason) => reason,
250                    None => "driver was not acquired and returned no diagnostic",
251                };
252                format!("{}: {diagnostic}", candidate.backend.label())
253            })
254            .collect::<Vec<_>>()
255            .join("; ");
256        return Err(GpuRegionPresenceSelfTestFailure {
257            acquired_backends,
258            message: format!("no GPU region-presence peer was acquired ({diagnostics})"),
259        });
260    }
261
262    let make_chunk = || Chunk {
263        data: format!("gpu_secret = {PLANTED}").into(),
264        metadata: ChunkMetadata::default(),
265    };
266
267    // CPU baseline on the SAME detector+chunk. This is the oracle: it proves the
268    // planted secret is detectable AT ALL on this build, so a low GPU count means
269    // a real GPU phase-1 divergence rather than an invalid/suppressed probe.
270    let cpu_results = scanner
271        .scan_chunks_with_backend(&[make_chunk()], ScanBackend::CpuFallback)
272        .map_err(|error| GpuRegionPresenceSelfTestFailure {
273            acquired_backends: acquired_backends.clone(),
274            message: format!("CPU baseline dispatch failed during GPU self-test: {error}"),
275        })?;
276    let cpu_total: usize = cpu_results.iter().map(Vec::len).sum();
277    if cpu_total == 0 {
278        return Err(GpuRegionPresenceSelfTestFailure {
279            acquired_backends,
280            message: "GPU self-test probe matched on no backend (CPU baseline is zero); fix the self-test probe so it survives suppression.".to_string(),
281        });
282    }
283
284    let mut peers = Vec::with_capacity(acquired_backends.len());
285    let mut failures = Vec::new();
286    for candidate in candidates
287        .into_iter()
288        .filter(|candidate| candidate.is_eligible())
289    {
290        let route = candidate.backend;
291        let Some(backend_id) = candidate.driver_id else {
292            failures.push(format!(
293                "{}: acquired driver returned no identity",
294                route.label()
295            ));
296            continue;
297        };
298        let degrade_before = scanner.runtime_status().gpu_degrade_count;
299        let results = match scanner.scan_coalesced_gpu_region_presence(
300            &[make_chunk()],
301            route,
302            scanner.execution_route_for_backend(route),
303        ) {
304            Ok(results) => results,
305            Err(error) => {
306                failures.push(format!(
307                    "{} ({backend_id}): dispatch failed: {error}",
308                    route.label()
309                ));
310                continue;
311            }
312        };
313        if scanner.runtime_status().gpu_degrade_count > degrade_before {
314            let diagnostic = match scanner.last_gpu_degrade_reason() {
315                Some(reason) => reason,
316                None => "runtime degrade recorded without a diagnostic".to_owned(),
317            };
318            failures.push(format!("{} ({backend_id}): {diagnostic}", route.label()));
319            continue;
320        }
321        let total: usize = results.iter().map(Vec::len).sum();
322        if total != cpu_total {
323            failures.push(format!(
324                "{} ({backend_id}): found {total} match(es), CPU found {cpu_total}",
325                route.label()
326            ));
327            continue;
328        }
329        peers.push(GpuRegionPresencePeerSelfTest {
330            backend: route,
331            backend_id,
332            matches: total,
333        });
334    }
335    if !failures.is_empty() {
336        let passed = peers
337            .iter()
338            .map(|peer| format!("{} ({})", peer.backend.label(), peer.backend_id))
339            .collect::<Vec<_>>()
340            .join(", ");
341        let passed = if passed.is_empty() {
342            "none".to_string()
343        } else {
344            passed
345        };
346        return Err(GpuRegionPresenceSelfTestFailure {
347            acquired_backends,
348            message: format!(
349                "GPU region-presence peer parity failed: {}; passed peers: {passed}",
350                failures.join("; ")
351            ),
352        });
353    }
354    Ok(GpuRegionPresenceSelfTest { peers })
355}