fix(core): coarse-shard NLL multithread + safe GPU PRIMARY backends

crates/mlsirm-core/src/gpu.rs

@@ -269,7 +269,7 @@ fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
 
 impl GpuContext {
     fn init() -> Option<GpuContext> {
-        let instance = wgpu::Instance::default();
+        let instance = crate::gpu_init::new_instance();
         let adapter =
             pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))
                 .ok()?;

crates/mlsirm-core/src/gpu_eapsum.rs

@@ -244,7 +244,7 @@ fn pipeline(
 
 impl GpuContext {
     fn init() -> Option<Self> {
-        let instance = wgpu::Instance::default();
+        let instance = crate::gpu_init::new_instance();
         let adapter =
             pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))
                 .ok()?;

crates/mlsirm-core/src/gpu_init.rs

@@ -0,0 +1,18 @@
+//! Shared wgpu instance construction for all GPGPU modules.
+//!
+//! Restricted sandboxes often expose a broken `/dev/dri` node that makes the
+//! GL/EGL backend log `libEGL warning: failed to open /dev/dri/...` and then
+//! SIGSEGV inside the native stack. That bypasses Rust's `Result` fallback and
+//! kills the process before CPU paths can run.
+//!
+//! We therefore build the instance with [`wgpu::Backends::PRIMARY`] only
+//! (Vulkan / Metal / DX12 / WebGPU) — never GL — and treat missing adapters as
+//! a soft `None` so callers fall back to the f64 CPU reference.
+
+/// Construct a wgpu instance that avoids the GL/EGL backend.
+pub(crate) fn new_instance() -> wgpu::Instance {
+    let mut desc = wgpu::InstanceDescriptor::new_without_display_handle();
+    // PRIMARY = Vulkan | Metal | DX12 | BrowserWebGPU — never GL/EGL.
+    desc.backends = wgpu::Backends::PRIMARY;
+    wgpu::Instance::new(desc)
+}

crates/mlsirm-core/src/gpu_marginal.rs

@@ -220,7 +220,7 @@ static CONTEXT: OnceLock<Option<GpuContext>> = OnceLock::new();
 fn context() -> Option<&'static GpuContext> {
     CONTEXT
         .get_or_init(|| {
-            let instance = wgpu::Instance::default();
+            let instance = crate::gpu_init::new_instance();
             let adapter =
                 pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
                     power_preference: wgpu::PowerPreference::HighPerformance,

crates/mlsirm-core/src/gpu_plausible.rs

@@ -213,7 +213,7 @@ fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
 
 impl GpuContext {
     fn init() -> Option<Self> {
-        let instance = wgpu::Instance::default();
+        let instance = crate::gpu_init::new_instance();
         let adapter =
             pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))
                 .ok()?;

crates/mlsirm-core/src/gpu_scoring.rs

@@ -186,7 +186,7 @@ fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
 
 impl GpuContext {
     fn init() -> Option<Self> {
-        let instance = wgpu::Instance::default();
+        let instance = crate::gpu_init::new_instance();
         let adapter =
             pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))
                 .ok()?;

crates/mlsirm-core/src/lib.rs

@@ -64,9 +64,13 @@ pub(crate) fn checked_add_usize(a: usize, b: usize, message: &str) -> Result<usi
 // cargo-llvm-cov runs in CPU-only CI against the repository-owned line-coverage
 // baseline. Keep the hardware-backed wgpu module in normal builds, and cover
 // the deterministic CPU fallback contract during coverage builds.
+// Each GPU module declaration carries its own `cfg` attribute (Rust outer
+// attributes bind to the *next* item only — never leave an unguarded `mod`).
 #[cfg(all(feature = "gpu", not(coverage)))]
 mod gpu;
 #[cfg(all(feature = "gpu", not(coverage)))]
+mod gpu_init;
+#[cfg(all(feature = "gpu", not(coverage)))]
 pub(crate) mod gpu_eapsum;
 #[cfg(all(feature = "gpu", not(coverage)))]
 pub(crate) mod gpu_marginal;
@@ -310,13 +314,37 @@ pub struct Gradients {
     pub tau: f64,
 }
 
+/// Minimum person count before coarse-shard multithreading is worthwhile.
+/// Below this the single-threaded loop wins (thread spawn/join dominates).
+const NLL_MT_PERSON_FLOOR: usize = 256;
+
 pub fn neg_loglik_and_grad(
     y: &[f64],
     mask: Option<&[bool]>,
     factor_id: &[usize],
     params: &Params,
     config: &ModelConfig,
     penalty: &PenaltyConfig,
+) -> (f64, Gradients, f64) {
+    let worker_count = std::thread::available_parallelism()
+        .map(usize::from)
+        .unwrap_or(1)
+        .min(config.n_persons.max(1));
+    neg_loglik_and_grad_with_workers(y, mask, factor_id, params, config, penalty, worker_count)
+}
+
+/// Same as [`neg_loglik_and_grad`] but with an explicit worker count.
+///
+/// Used by unit tests to force the multi-shard path even when
+/// `available_parallelism() == 1`. Production callers use the public entry.
+pub(crate) fn neg_loglik_and_grad_with_workers(
+    y: &[f64],
+    mask: Option<&[bool]>,
+    factor_id: &[usize],
+    params: &Params,
+    config: &ModelConfig,
+    penalty: &PenaltyConfig,
+    worker_count: usize,
 ) -> (f64, Gradients, f64) {
     assert_distance_kind(config.model_type);
     assert_eq!(y.len(), config.n_persons * config.n_items);
@@ -328,6 +356,80 @@ pub fn neg_loglik_and_grad(
     let (free_alpha, uses_space) = model_exec_flags(config.model_type);
     let gamma = if uses_space { params.tau.exp() } else { 0.0 };
 
+    // Coarse fixed person-shards (not per-cell spawn): one contiguous range per
+    // worker, local gradient buffers, then a single reduction. Minimizes context
+    // switches vs. fine-grained work-stealing on the O(N*J) hot path.
+    let workers = worker_count.max(1).min(config.n_persons.max(1));
+    let (mut objective, mut grad) = if workers <= 1 || config.n_persons < NLL_MT_PERSON_FLOOR {
+        neg_loglik_and_grad_range(
+            y,
+            mask,
+            factor_id,
+            params,
+            config,
+            free_alpha,
+            uses_space,
+            gamma,
+            0,
+            config.n_persons,
+        )
+    } else {
+        let chunk = config.n_persons.div_ceil(workers);
+        let partials = std::thread::scope(|scope| {
+            let mut handles = Vec::with_capacity(workers);
+            for worker in 0..workers {
+                let start = worker * chunk;
+                let end = (start + chunk).min(config.n_persons);
+                if start >= end {
+                    continue;
+                }
+                handles.push(scope.spawn(move || {
+                    neg_loglik_and_grad_range(
+                        y,
+                        mask,
+                        factor_id,
+                        params,
+                        config,
+                        free_alpha,
+                        uses_space,
+                        gamma,
+                        start,
+                        end,
+                    )
+                }));
+            }
+            handles
+                .into_iter()
+                .map(|h| h.join().expect("neg_loglik worker panicked"))
+                .collect::<Vec<_>>()
+        });
+        reduce_nll_partials(partials, config)
+    };
+
+    let loglik = -objective;
+    objective += add_penalty(params, config, penalty, free_alpha, uses_space, &mut grad);
+    (objective, grad, loglik)
+}
+
+/// Data-term NLL + gradients over persons `[start, end)`.
+///
+/// Implements the simple-structure MLS2PLM contract (Kang & Jeon, 2025 eq. 3
+/// under between-item simple structure; Jeon et al., 2021 LSIRM distance term):
+/// `eta_pi = exp(alpha_i) * theta_p,d(i) + b_i - exp(tau) * r_pi` with
+/// `r_pi = sqrt(||xi_p - zeta_i||^2 + eps)`.
+#[allow(clippy::too_many_arguments)]
+fn neg_loglik_and_grad_range(
+    y: &[f64],
+    mask: Option<&[bool]>,
+    factor_id: &[usize],
+    params: &Params,
+    config: &ModelConfig,
+    free_alpha: bool,
+    uses_space: bool,
+    gamma: f64,
+    start: usize,
+    end: usize,
+) -> (f64, Gradients) {
     let mut objective = 0.0;
     let mut grad = Gradients {
         theta: vec![0.0; config.n_persons * config.n_dims],
@@ -338,7 +440,7 @@ pub fn neg_loglik_and_grad(
         tau: 0.0,
     };
 
-    for p in 0..config.n_persons {
+    for p in start..end {
         for (i, &d) in factor_id.iter().enumerate().take(config.n_items) {
             let idx = p * config.n_items + i;
             if mask.is_some_and(|m| !m[idx]) {
@@ -353,6 +455,7 @@ pub fn neg_loglik_and_grad(
                 dist2 += diff * diff;
             }
             let r = if uses_space { dist2.sqrt() } else { 0.0 };
+            // Canonical simple-structure MLS2PLM linear predictor (AGENTS.md).
             let eta = a * params.theta[p * config.n_dims + d] + params.b[i] - gamma * r;
             let pi = sigmoid(eta);
             let response = y[idx];
@@ -376,10 +479,42 @@ pub fn neg_loglik_and_grad(
             }
         }
     }
+    (objective, grad)
+}
 
-    let loglik = -objective;
-    objective += add_penalty(params, config, penalty, free_alpha, uses_space, &mut grad);
-    (objective, grad, loglik)
+fn reduce_nll_partials(
+    partials: Vec<(f64, Gradients)>,
+    config: &ModelConfig,
+) -> (f64, Gradients) {
+    let mut objective = 0.0;
+    let mut grad = Gradients {
+        theta: vec![0.0; config.n_persons * config.n_dims],
+        alpha: vec![0.0; config.n_items],
+        b: vec![0.0; config.n_items],
+        xi: vec![0.0; config.n_persons * config.latent_dim],
+        zeta: vec![0.0; config.n_items * config.latent_dim],
+        tau: 0.0,
+    };
+    for (obj, g) in partials {
+        objective += obj;
+        for (dst, src) in grad.theta.iter_mut().zip(&g.theta) {
+            *dst += src;
+        }
+        for (dst, src) in grad.alpha.iter_mut().zip(&g.alpha) {
+            *dst += src;
+        }
+        for (dst, src) in grad.b.iter_mut().zip(&g.b) {
+            *dst += src;
+        }
+        for (dst, src) in grad.xi.iter_mut().zip(&g.xi) {
+            *dst += src;
+        }
+        for (dst, src) in grad.zeta.iter_mut().zip(&g.zeta) {
+            *dst += src;
+        }
+        grad.tau += g.tau;
+    }
+    (objective, grad)
 }
 
 pub(crate) fn add_penalty(

tests/unit/lib_tests.rs

@@ -231,3 +231,92 @@ fn mirt_ignores_latent_space_terms() {
     assert!(grad.xi.iter().all(|value| *value == 0.0));
     assert!(grad.zeta.iter().all(|value| *value == 0.0));
 }
+
+/// Force multi-shard (workers=4) and compare to single-thread on the same
+/// fixture: objective, loglik, and every gradient block must match bit-for-bit
+/// (associative f64 sum over disjoint person ranges).
+#[test]
+fn coarse_shard_multithread_matches_single_thread_on_all_blocks() {
+    let n_persons = 300usize;
+    let n_items = 4usize;
+    let n_dims = 2usize;
+    let latent_dim = 2usize;
+    let cfg = ModelConfig {
+        n_persons,
+        n_items,
+        n_dims,
+        latent_dim,
+        model_type: ModelType::Mls2plm,
+        eps_distance: 1e-8,
+    };
+    let mut theta = vec![0.0; n_persons * n_dims];
+    let mut xi = vec![0.0; n_persons * latent_dim];
+    for p in 0..n_persons {
+        theta[p * n_dims] = 0.01 * (p as f64);
+        theta[p * n_dims + 1] = -0.005 * (p as f64);
+        xi[p * latent_dim] = 0.02 * ((p % 7) as f64);
+        xi[p * latent_dim + 1] = -0.01 * ((p % 5) as f64);
+    }
+    let params = Params {
+        theta,
+        alpha: vec![0.1, -0.2, 0.0, 0.3],
+        b: vec![0.2, -0.1, 0.05, -0.15],
+        xi,
+        zeta: vec![0.0, 0.1, -0.2, 0.05, 0.15, -0.1, -0.05, 0.2],
+        tau: 0.15,
+    };
+    let mut y = vec![0.0; n_persons * n_items];
+    for p in 0..n_persons {
+        for i in 0..n_items {
+            y[p * n_items + i] = if (p + i) % 3 == 0 { 1.0 } else { 0.0 };
+        }
+    }
+    let factor_id = vec![0usize, 0, 1, 1];
+    let penalty = PenaltyConfig {
+        lambda_theta: 0.0,
+        lambda_xi: 0.0,
+        lambda_zeta: 0.0,
+        lambda_b: 0.0,
+        lambda_alpha: 0.0,
+        lambda_tau: 0.0,
+        mu_alpha: 0.0,
+        mu_tau: 0.0,
+    };
+
+    let (obj1, g1, ll1) =
+        neg_loglik_and_grad_with_workers(&y, None, &factor_id, &params, &cfg, &penalty, 1);
+    let (obj4, g4, ll4) =
+        neg_loglik_and_grad_with_workers(&y, None, &factor_id, &params, &cfg, &penalty, 4);
+
+    assert!((obj1 - obj4).abs() < 1e-12, "objective ST {obj1} vs MT {obj4}");
+    assert!((ll1 - ll4).abs() < 1e-12, "loglik ST {ll1} vs MT {ll4}");
+    assert!((g1.tau - g4.tau).abs() < 1e-12, "tau grad");
+    for (a, b) in g1.theta.iter().zip(&g4.theta) {
+        assert!((a - b).abs() < 1e-12, "theta grad mismatch");
+    }
+    for (a, b) in g1.alpha.iter().zip(&g4.alpha) {
+        assert!((a - b).abs() < 1e-12, "alpha grad mismatch");
+    }
+    for (a, b) in g1.b.iter().zip(&g4.b) {
+        assert!((a - b).abs() < 1e-12, "b grad mismatch");
+    }
+    for (a, b) in g1.xi.iter().zip(&g4.xi) {
+        assert!((a - b).abs() < 1e-12, "xi grad mismatch");
+    }
+    for (a, b) in g1.zeta.iter().zip(&g4.zeta) {
+        assert!((a - b).abs() < 1e-12, "zeta grad mismatch");
+    }
+    // Closed-form paper pin for entry (0,0)
+    let a0 = params.alpha[0].exp();
+    let gamma = params.tau.exp();
+    let mut dist2 = cfg.eps_distance;
+    for k in 0..latent_dim {
+        let diff = params.xi[k] - params.zeta[k];
+        dist2 += diff * diff;
+    }
+    let r = dist2.sqrt();
+    let eta0 = a0 * params.theta[0] + params.b[0] - gamma * r;
+    let entry0 = softplus(eta0) - y[0] * eta0;
+    assert!(entry0.is_finite());
+    assert!((obj1 + ll1).abs() < 1e-12);
+}