Linux CUDA GPU opportunities for R-MSRM, GP-MAP and GP posterior¶

Measurement date: September 21, 2026. Status: research note on the development source at Git HEAD 5a49a54; no package release, default change, or production optimization is made by this note.

This note builds on the prior Apple MPS findings and the post-release Linux evaluation plan. It answers three questions with measurements on the user's Linux workstation: where GPU acceleration could help each model, how much float64 throughput the installed NVIDIA cards actually provide, and which bottlenecks are algorithmic rather than hardware limits.

Summary¶

1. R-MSRM does not need a GPU. A bounded fit at the recorded one-eighth empirical scale spends its time in many small per-block NumPy/SciPy operations and Python-level loops. The largest single operation, the dense reconstruction residual, is about 30 ms. Restarts already run as parallel processes across the 64 CPU cores. Keep R on CPU.
2. GP-MSRM already runs on CUDA with no code change. The grouped and dense likelihoods are JAX; the guarded symmetric factorization selects the CPU LAPACK callback only on the CPU platform and native Cholesky elsewhere. Installing jax[cuda13] and starting Python with the CUDA platform is sufficient to run the released code on a GPU in float64. GPU support in the package is therefore an installation, configuration, evidence-recording and testing task, not a port.
3. Float64 throughput on these cards is limited by design, and the PRO 6000 is the useful one. Both the RTX 3090 (Ampere, compute capability 8.6) and the RTX PRO 6000 Blackwell (compute capability 12.0) execute float64 at 1/64 of their float32 rate. Measured float64 matrix multiplication peaks near 0.55 TFLOPS on the 3090 and 2.1 TFLOPS on the PRO 6000, against roughly 1 TFLOPS from the 64-core CPU under the day's background load. For the model's actual objective-and-gradient at the one-eighth scale the PRO 6000 took 66 ms, the 3090 172 ms and the best CPU row 644 ms. Float32 is 30 to 80 times faster on the GPUs (and JAX's default float32 matmul on CUDA silently uses TF32) but is outside the float64 numerical contract, and the MPS note already recorded a float32 factorization failing gradient qualification.
4. The current CPU GP path had a large avoidable cost in this environment. The CPU LAPACK callback took 0.4 to 2.9 s for a 2,727-dimension factorization depending on OpenBLAS thread count and background load, while direct POTRF+POTRI with 16 threads took 60 ms and the 3090 took 135 ms. Thread control is the first fix for the CPU path; it needs no GPU.
5. For full recordings the dense grouped algebra is the wall, not the device. Cost grows with the cube of unique observation nodes times factors. A single run eight times longer than the one-eighth case at K=5 implies a 36,000-dimension dense factorization per gradient, which is tens of seconds on any device here and tens of gigabytes of memory. If the full data are eight separate runs, cost grows only linearly and runs are independent. The state-space formulation, with grouped observation updates and a parallel-in-time scan, is the candidate that changes the scaling; the plan already names it.
6. The posterior is where a GPU pays, and it is still expensive. One NUTS leapfrog step for the K = 3 one-eighth case cost 2.0 s on the CPU, 0.37 s on the 3090 and 0.16 s on the PRO 6000, in each case two to three times the bare gradient. Vectorizing two chains on one card gained about 1.35x throughput, not 2x. At realistic budgets (a thousand warmup and twelve hundred draws at a few hundred steps each) one chain at this scale is on the order of a day on the 3090 and half a day on the PRO 6000, versus about a week on the CPU. The ArviZ diagnostics phase took 70 to 100 s for only 15 draws over 57,757 parameters on every device and grows with draw count; it needs its own profile.
7. The state-space backend cannot be measured at empirical scale as written. Two Gaussian responses give 50 states per factor, so the state dimension is 150 at K = 3. The implementation stages one transition and one process-noise matrix per scalar observation, 17.6 GB for a 48,000-observation run, and reverse-mode differentiation through the scan stores per-step covariances on top: the 96 GB card requested a single 214 GB allocation for the gradient on a 36-feature reduction, and the 3090 failed earlier. At 3.6 million observations the same arrays would be terabytes. Grouping observations by unique node (909 here) reduces the transition arrays to 327 MB and the scan length by a factor of 50. Grouped observation updates, with gradient checkpointing across the scan, are therefore a prerequisite, not an optimization, for state-space at empirical scale on any device.

The rest of this note records the evidence and its limits.

Hardware and runtime inventory¶

Both GPUs were in use by another lab member's jobs during the first part of the session (about 10 GB resident on the 3090, 5 to 8 GB on the PRO 6000, intermittent utilization), and the CPU load average was between 25 and 65 from other users' work throughout. All timings below are therefore upper bounds on contended hardware, recorded with XLA_PYTHON_CLIENT_PREALLOCATE=false. At the user's request, measurement first moved to the 3090 only; the full PRO 6000 series ran later once that card showed zero utilization and under 1 GB resident, with the user's authorization. The CPU rows remain the least reliable and need an idle-machine repeat.

The MPS note's lessons were followed: one frozen source tree, separate preparation, compilation and warm timings, device synchronization before stopping timers, callback and fallback behavior observed, and float64 throughout.

Raw float64 linear algebra¶

bench_linalg.py times a square matrix product, Cholesky, LU and Cholesky-inverse (cho_solve against the identity, the pattern used by the native symmetric solve) with JAX on each device. Median of three warm repeats after one compile.

Readings:

• In float64 the 3090 is at rough parity with the loaded CPU for the factorization sizes the grouped model produces today, and slower for matrix products. The PRO 6000 is about 3 to 4 times faster than the CPU in float64 at these sizes, sustains its rate to 16,384 dimensions, and is 3.6x the 3090 at 8,192. Neither approaches the float32 numbers, which are excluded by the numerical contract.
• The PRO 6000 float32 matrix product exceeds the card's nominal float32 rate, which shows that JAX's default float32 matmul precision on CUDA uses TF32 tensor cores. Any float32 experiment must set jax_default_matmul_precision="highest" or it measures a reduced-precision path.
• The CPU rows at n = 8,192 are anomalous (matmul throughput halves) and reflect the background load, not a hardware limit.
• Memory, not time, bounds single-run scaling on the 3090: a 36,000-dimension float64 matrix is 10.4 GB and the inverse path holds several. The PRO 6000's 96 GB removes that bound for a single run, at about 70 s per factorization-plus-inverse by cubic extrapolation from 16,384.

R-MSRM: CPU profile at the recorded scale¶

Synthetic workload shaped like the recorded one-eighth empirical case: four participants with a 3,600-feature brain stream sampled every 1.8 s (252 samples), five participants with rating and physiology streams on 1.2 s and 1.5 s clocks, one run of about 454 s, K = 3, latent_dt = 0.9 (505 grid points), Identity brain response, Gaussian rating lag and Gaussian physiology width and lag estimated with shared pooling. profile_r.py ran one restart for three alternating iterations with 8 OpenBLAS threads under cProfile.

A repeat at twice the recording length (1,009 grid points, 7.3 million observed values, load average 47 to 58) took 6.6 s for the same three iterations: 1.7 s of one-time block preparation, then about 1.6 s per iteration with the reconstruction residual at 40 ms per call and the Gaussian operator construction at 5 ms per call. Growth is close to linear in recording length once the background load is discounted.

Extrapolation: a 100-iteration restart costs on the order of one to two minutes at this scale, and the per-iteration cost scales roughly linearly with recording length because every term is per-observation-row or banded. Eight times the recording with 30 participants implies roughly 10 to 30 minutes per restart, and restarts are independent processes. A GPU would accelerate only the 30 ms dense residual (a memory-bound weighted GEMM) and the loading solves; the operator construction, sparse assembly, L-BFGS-B kernel loop and objective bookkeeping are Python and SciPy sparse code with no device path. The prior MPS note measured the same loading update as 3 ms on CPU. Recommendation: no GPU work for R. The worthwhile CPU items, if profiling of the real data confirms them, are (a) reusing the reconstruction residual between complete_objective calls inside one alternating iteration, since the objective is evaluated up to a dozen times per iteration on unchanged loadings, and (b) a banded Cholesky (solveh_banded) in place of SuperLU for the latent normal equations. Neither changes the objective.

GP-MSRM grouped likelihood: where the time goes¶

Same synthetic data, linear_algebra="grouped", K = 3, 909 unique (modality, time) nodes, factorization dimension 2,727, 57,757 parameters, 3.64 million observed scalars. bench_gp_parts.py times each stage under jax.jit with warm medians of three.

Two findings matter more than the device comparison:

• The CPU callback path was pathological here. symmetric_solve on the CPU platform calls SciPy's POTRF and POTRI through jax.pure_callback. Direct POTRF+POTRI at this size takes 420 ms with one OpenBLAS thread and 60 ms with 16 threads (potri_threads.py). Through the callback in a JAX process the same work took 655 ms with one thread, 386 ms with 16 threads at load 42, 2.5 s with 16 threads at load 60, and 2.9 s with the default 64 threads at load 36. OpenBLAS threads, XLA's 128-thread pool and the other users' processes oversubscribe the cores, and the callback's serialization with XLA's thread pool amplifies contention. The CI convention of one BLAS thread is safe but slow; 8 to 16 threads is the right region on this machine when it is otherwise idle, and it must be set before Python starts. An uncontended CPU measurement is still needed. The Mac note's 116 ms whole-gradient timing was obtained with one thread on an idle machine, which is why the Linux CPU number looks worse than the Mac number.
• The observation-count part is cheap on the GPU and not on XLA CPU. Going from 360 to 3,600 brain features (0.37 to 3.64 million observations) added 4.8 s to the CPU gradient and 30 ms to the 3090 gradient. XLA's CPU scatter-add and its transpose are effectively single-threaded; on the GPU they are memory-bound and negligible. This is the one place where the GPU wins by a large factor even in float64, and it grows with feature count.

Gradient agreement: the objective and gradient fingerprints agreed between CPU and GPU at the float64 level in these runs (differences reported in bench_gp.jsonl). This is algebraic agreement at one parameter vector, not convergence qualification.

GP objective-and-gradient, NUTS and state-space timings¶

bench_gp.py builds the synthetic problem, evaluates the jitted objective-and-gradient at one seeded parameter vector (warm median of five), and optionally runs a short NUTS (15 warmup, 15 draws, diagonal mass matrix, tree depth at most 5) through the package's own sample entry point so that initialization, adaptation, draw conversion, log-likelihood and diagnostics phases are all exercised.

Objective and gradient¶

The objectives agreed between CPU and 3090 to 16 significant digits in all three cases (for example 5323210.5539216 at K = 3) and the sampled gradient entries agreed to about 13 digits. The two K = 3 CPU rows differ by 5x with identical settings and only the background load changed (43 versus 27), which is the clearest evidence that CPU numbers on this shared machine need an idle-machine repeat. Against the better CPU row the 3090 was 3.3x faster at K = 3 and the PRO 6000 9.8x; against the contended rows the 3090 was 16x at K = 3, 15x at K = 5 for the one-eighth case, and 7x at the doubled length. The PRO 6000 was 2.6 to 3.8x faster than the 3090 across the matched cases. On the PRO 6000 each doubling of recording length as one run raised the K = 5 gradient cost 6.0x and then 7.6x, so a full recording as one run extrapolates to roughly 70 s per gradient there, and the 96 GB card held the 18,530-dimension case that the 3090 could not. Doubling the recording length as one run raised the K = 5 gradient cost 8.4x on the GPU and 3.9x on the CPU, consistent with cubic scaling on the GPU and with the CPU still being dominated by the contended callback at the smaller size. Against an idle, well-threaded CPU the expected ratio is 2 to 4x at these sizes, from the raw factorization sweep. Peak GPU memory was not captured by the script (the allocator statistic returned nothing with preallocation disabled); capture it with nvidia-smi sampling in the next campaign.

Short NUTS, K = 3, one-eighth scale¶

The same seed produced the same 465-step trajectory length on all three devices. Per leapfrog step the PRO 6000 was 2.3x the 3090 and 13x the CPU, and vectorizing two chains gave 1.37x throughput on either card. At 15 draws the diagnostics phase already exceeds warmup plus sampling on the PRO 6000. The per-leapfrog figures include NUTS kernel compilation amortized over a short run, so they overstate the steady-state cost; the ratio of leapfrog cost to bare gradient (2.2x on the GPU, 3.2x on the CPU) also includes the diagonal mass-matrix adaptation and NumPyro's tree building, which are serial small operations that suit the CPU less than the dense algebra does. The 3090 was 5.5x faster than the CPU per leapfrog step in this pairing. Sampling ran with XLA_PYTHON_CLIENT_PREALLOCATE=false on a card shared with other jobs. Diagnostics run on the CPU through ArviZ over every parameter, cost about 70 s on either device at 15 draws, and are the second-largest phase.

State-space backend¶

ResponseStateSpace.prepare gives 50 states per factor for this response configuration (two Gaussian responses on rational or Laguerre banks plus the Matérn pair). StateSpaceSystem.prepare materializes (observations, 150, 150) transition and process-covariance arrays, 17.6 GB at K = 3 and 48.8 GB at K = 5 for this reduced run, and reverse-mode differentiation through lax.scan additionally stores the per-step predicted and filtered covariances, which is what the 214 GB request on the PRO 6000 reflects. Checkpointed differentiation (jax.checkpoint over scan chunks) bounds that residual memory, as the plan's section 7 already requires. At the empirical 3.6 million observations the same arrays would be 1.3 TB. With observations grouped to the 909 unique nodes the arrays are 327 MB and 909 MB. This is the plan's section 7 "exact grouped state-space observation updates" candidate, and it is required before any state-space timing at empirical scale is meaningful, on GPU or CPU. A parallel associative scan would additionally replace 909 sequential 150-dimension updates with a logarithmic-depth combination, which is where a GPU would help this backend.

Scaling to the full recordings¶

Let N be the number of unique observation nodes per run and K the factor count. The grouped algebra forms and factorizes an NK-square dense matrix per run per gradient, and the guarded path also materializes its inverse for the implicit derivative. Costs scale as (NK)^3 in time and (NK)^2 in memory per matrix.

These extrapolations use the measured n = 2,700 to 16,384 cubic trend (the PRO 6000 column interpolates its measured 4,096, 8,192 and 16,384 rows) and should be replaced by measurements. Their point is structural: with one long run, no device in this machine makes a NUTS chain feasible at full length (a thousand warmup and twelve hundred draws at up to 1,023 leapfrog steps each would need on the order of a year of gradients). With runs at their natural boundaries, the full data costs about eight times the one-eighth case, runs are independent and can be spread over devices, and the GPU speedup of the per-run gradient carries over.

Whether the user's recordings are one run or several is the single most important fact to verify in the plan's Stage A inventory before choosing implementation work. The plan's Stage D duration ladder should be built at run boundaries, as it already says.

What "GPU support" would mean in the package¶

For GP-MSRM the code already dispatches by platform. The remaining work is contractual:

1. Installation extra and documentation. Add an optional bayesian-cuda extra (or documented jax[cuda13] instructions) with the same pinned JAX version, plus a getting-started section on JAX_PLATFORMS, XLA_PYTHON_CLIENT_PREALLOCATE, thread variables and CUDA_VISIBLE_DEVICES. Keep CPU-only installation the default.
2. Evidence recording. execution_info already records backend and devices for sampling; MAP diagnostics and fitted archives should record the same, plus dtype, so a fit made on a GPU is identifiable. Warmup checkpoints already bind to device identity.
3. Numerical qualification on CUDA. Run the existing factorization, analytic-score, multifactor and posterior smoke tests on the CUDA platform and compare objectives, physical gradients and MAP endpoints against CPU at identical parameters, as the plan's Stage C requires. Pin jax_default_matmul_precision="highest" in the test harness; it only affects float32 but prevents a silent TF32 path if any float32 array ever appears.
4. Thread and device configuration guard. Emit a warning when the CPU platform runs with more BLAS threads than physical cores or when pure_callback is on the hot path with an oversubscribed pool; document 8 to 16 threads as the tested CPU setting on this hardware.
5. Chain scheduling. Keep chain_method="sequential" as default; qualify vectorized on one GPU (padded trajectories waste work when tree depths differ across chains) and document running separate processes per device for independent chains, each with its own seed and archive.
6. Memory guard for the inverse path. symmetric_solve materializes the full inverse. On a 24 GB card the practical dimension ceiling is about 15,000 to 18,000 in float64; document it and keep the LU reference path as the fallback. A reverse-mode rule that uses solves against the cotangent instead of the explicit inverse would remove that ceiling and reduce work; it is an algebraic change to the derivative rule, not to the density, and would need the existing finite-difference and higher-order derivative tests.

For R-MSRM, no GPU backend. A NumPy array-API port for the two dense operations would touch a third of the fitting code for a few percent of runtime.

Candidate ranking for the plan's Stage C¶

Evidence retained locally¶

Scripts and raw outputs are retained under ignored local_data/linux-performance-v1/2026-09-21-gpu-opportunities/: bench_linalg.py with bench_linalg.log (CPU and partial PRO 6000 rows), bench_linalg_3090.json and bench_linalg_pro6000.json; synth.py; bench_gp.py with bench_gp.jsonl (every objective, gradient, NUTS and error row); bench_gp_parts.py; potri_threads.py; callback_probe.py; profile_r.py with profile_r_F3600.log; run_gp.sh, sequence.sh with sequence.log (3090 and CPU series, including the scale-2 R profile) and sequence_blackwell.sh with sequence_blackwell.log. No empirical data were used; all inputs are synthetic and seeded.

Limits¶

• Synthetic data of the recorded shape, not the empirical arrays; feature count, node count and parameter count were matched from the MPS note's description, not from the actual inventory.
• Contended hardware throughout; the CPU rows especially vary by a factor of two to five with background load. The GPU rows were taken with other users' contexts resident on the cards and are consistent across repeats, but not idle-machine measurements.
• Single parameter vector timings; no optimizer trajectories, no convergence qualification, no posterior diagnostics beyond the short NUTS timing run.
• One run per dataset; the multi-run parallel case is extrapolated, not measured.