Refinement: its GPU modules loaded in the model warm-up, not at Step 11

The intermittent 1.1-2.7 s stall at the start of live refinement (TIMING
refine.engine) is the CUDA driver's lazy module load. Stack samples of the broker in
a stalled run show the engine's first kernel launch (phase_table_kernel in
RefinementGPUCard) inside cuLibraryLoadData for the whole stall, and per-thread
/proc samples show the process using no CPU meanwhile: it is waiting for the
executable's pages. On this machine the binaries are on a busy, nearly full HDD and
the page cache is under pressure, so pages of the executable touched for the first
time late in a run - the refinement's fatbin, and in the worst runs other code and
data too (refine.topology 0.015 -> 2.5 s, refine.placement 0.15 -> 2.3 s) - can take
seconds to come back. Reproduced on purpose: dropping the broker's unmapped pages
from the page cache with a 4 GB write to /data beside it gives refine.engine
0.45 / 0.85 s (0.02 s normally).

WarmUpModelValidation, which the broker runs when a model is uploaded, now also loads
the modules of the refinement engine's and the scale fit's kernels (the two
refinement TUs a validation does not launch) on every card. With that, the same
induced load gives refine.engine 0.020 / 0.020 / 0.021 s. Other cold pages can still
stall under such load; keeping the binaries on a local SSD avoids all of it.

Live thau replays, no load added: refine.engine 0.020 s in 3 of 3, refinement 2.9 s
(rc184-port 3.1-3.2 s, with the two previous levers in). The warm-up only loads
modules, so the results are unchanged (same R-work/R-free and evaluations live;
ModelValidation_WarmUpChangesNothing passes).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQsU5a8AJUdJPgWevAK

image_analysis/structure_refinement/ModelScaleGPU.cu

@@ -402,6 +402,11 @@ size_t ModelScaleGPU::RobustBytes(size_t max_points) {
     return MAX_ROBUST_FITS * max_points * sizeof(float) + MAX_ROBUST_FITS * sizeof(double);
 }
 
+void ModelScaleGPU::LoadModule() {
+    cudaFuncAttributes attributes;
+    cuda_err(cudaFuncGetAttributes(&attributes, solvent_amplitudes));
+}
+
 size_t ModelScaleGPU::DeviceBytes(size_t max_points) {
     return max_points * (3 * sizeof(int) + sizeof(double) + sizeof(float) + sizeof(unsigned char))
            + MAX_FITS * max_points * sizeof(float)

image_analysis/structure_refinement/ModelScaleGPU.h

@@ -43,6 +43,8 @@ struct ModelScaleFitState {
 class ModelScaleGPU {
 public:
     static size_t DeviceBytes(size_t max_points);
+    // The module of its kernels loaded on the current device now, not at their first launch (lazy loading).
+    static void LoadModule();
     ModelScaleGPU(cudaStream_t stream, size_t max_points);
 
     // Per zone: the Scaling points in gemmi prepare_points() order, adp_symmetry_constraints(sg) rows

image_analysis/structure_refinement/ModelValidation.cpp

@@ -51,6 +51,7 @@
 #include "../indexing/CUDAMemHelpers.h"   // cuda_trim_allocation_pool
 #include "RigidBodyGPU.h"
 #include "RigidBodyGPUEngine.h"   // RigidBodyGPUEngine::CurrentDevice
+#include "RefinementGPUEngine.h"  // RefinementGPUEngine::LoadModules
 #include "../../common/CUDAWrapper.h"
 #include "../../common/JFJochException.h"
 #else
@@ -2188,6 +2189,13 @@ void WarmUpModelValidation(std::shared_ptr<const RugnuxModel> model, size_t nthr
         const auto start = std::chrono::steady_clock::now();
         // Held for as long as it runs: work beside a run waits for it where the card is too small for both.
         GPUWorkBeside beside;
+        // The refinement's kernels that a validation does not launch, on every card: lazily their modules would be
+        // loaded at their first launch, in the refinement, from the executable - 1-3 s in a live run where its
+        // pages had left the page cache and the disk was busy.
+        for (int d = 0; d < get_gpu_count(); d++) {
+            const ScopedGpuPin pin(d);
+            RefinementGPUEngine::LoadModules();
+        }
         Logger quiet = Logger::Buffered();
         if (const SearchModel search = PrepareSearchModel(*model); search.failure_reason.empty() && search.predicted) {
             std::string failure;

image_analysis/structure_refinement/RefinementGPU.cu

@@ -12,6 +12,7 @@
 
 #include "../../common/JFJochException.h"
 #include "../indexing/CUDAMemHelpers.h"
+#include "ModelScaleGPU.h"   // ModelScaleGPU::LoadModule
 
 namespace {
 
@@ -1015,6 +1016,12 @@ size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t reco
     return bytes;
 }
 
+void RefinementGPUEngine::LoadModules() {
+    cudaFuncAttributes attributes;
+    cuda_err(cudaFuncGetAttributes(&attributes, phase_table_kernel));
+    ModelScaleGPU::LoadModule();
+}
+
 std::shared_ptr<RefinementGPUCard> RefinementGPUEngine::MakeCard(int device, const RefinementGPUSetup &setup) {
     DeviceGuard guard(device);
     return std::make_shared<RefinementGPUCard>(device, setup);

image_analysis/structure_refinement/RefinementGPUEngine.h

@@ -57,6 +57,9 @@ public:
     static size_t DeviceBytes(const RefinementGPUSetup &setup, size_t records, size_t slots);
     static size_t FFTWorkBytes(int device, const RefinementGPUSetup &setup);
     static std::shared_ptr<RefinementGPUCard> MakeCard(int device, const RefinementGPUSetup &setup);
+    // The modules of the engines' kernels and of the scale fit's (ModelScaleGPU) loaded on the current device now:
+    // lazily the driver loads each at its first launch, reading it from the executable.
+    static void LoadModules();
 
     // `setup` is read for as long as the engine lives; `card` is MakeCard(device, setup)'s. `own_lane`: on a lane of
     // working buffers of its own, else on the card's lane with the fewest engines.