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