fix(#4621): QK-norm models take the FP16 batched prefill when it fits (D3-SPLIT B, stacked on #4620)

crates/aprender-serve/src/cuda/gpu_profile.rs

@@ -1124,38 +1124,59 @@ impl GpuProfile {
     ///                     argmax). FP16 → 4 of 4 pass.
     ///   qwen2.5-1.5b      (QKV bias, no QK-norm) → 4 of 4 pass under FP8.
     /// The defect is the FP8 GEMM feeding the per-head norm, not the norm; its
-    /// root cause is #3483. Until then a QK-norm model takes the SERIAL prefill
-    /// (per-token `forward_gpu_resident`, byte-correct in the same matrix, no
-    /// prefill weight cache — the FP16 HGEMM cache is 2× the FP8 one and put
-    /// Qwen3-8B + `apr qa`'s CPU twin over 24 GB: `CUDA_ERROR_OUT_OF_MEMORY` in
-    /// `prefill_all_layers_gpu`) with FP8 off, which is what GB10 (cc ≥ 120)
-    /// already does by default. `FP8_PREFILL=1` forces FP8 back for A/B testing,
-    /// and an explicit `BATCHED_PREFILL` keeps its own answer. Returns `true`
-    /// when it changed anything, so the caller can print why.
-    pub fn disable_fp8_for_qk_norm(
+    /// root cause is #3483. FP8 goes off for a QK-norm model; the prefill stays
+    /// on the path `select_prefill_path` chose, since FP16 passes (above).
+    ///
+    /// #3715: this used to move a QK-norm model to the SERIAL prefill as well,
+    /// because the unchunked FP16 batched prefill ran Qwen3-8B beside `apr qa`'s
+    /// CPU twin out of 24 GB (`CUDA_ERROR_OUT_OF_MEMORY` in
+    /// `prefill_all_layers_gpu`). Serial is one `forward_gpu_resident` per prompt
+    /// token: Qwen3-1.7B took 109 s for 7955 tokens and did not finish 26131 in
+    /// 600 s on the 4090 (GPU 99-100 %), so every long-context Qwen3 cell timed
+    /// out. The batched prefill is now chunked to a fixed score budget, which is
+    /// what bounded the OOM. `FP8_PREFILL=1` forces FP8 back for A/B testing.
+    /// Returns `true` when it changed anything, so the caller can print why.
+    pub fn disable_fp8_for_qk_norm(&mut self, has_qk_norm: bool, forced_fp8: Option<&str>) -> bool {
+        if !has_qk_norm || forced_fp8 == Some("1") || !(self.fp8_prefill || self.fp8_decode) {
+            return false;
+        }
+        self.fp8_prefill = false;
+        self.fp8_decode = false;
+        true
+    }
+
+    /// #3715: a QK-norm model whose FP16 batched prefill does not fit on the device
+    /// keeps the serial prefill, as it did before.
+    ///
+    /// With FP8 off, the FP16 weight cache is warmed on either path, but only the
+    /// batched prefill also needs its workspace and score buffers. Qwen3-8B on the
+    /// 24 GB 4090 takes 4.7 GB of weights plus a 15.1 GB FP16 cache, which leaves
+    /// 0.1 GB. The batched workspace then fails to allocate and the dense session
+    /// falls back to the CPU (a 2k chat turn: 382 s). Serial prefills on the device
+    /// in the room that is left. `need_bytes` is weights + cache + KV + the prefill
+    /// reserve. An explicit `BATCHED_PREFILL` keeps its own answer. Returns `true`
+    /// when it moved the path.
+    pub fn serial_when_fp16_prefill_does_not_fit(
         &mut self,
         has_qk_norm: bool,
-        forced_fp8: Option<&str>,
         forced_batched: Option<&str>,
+        need_bytes: usize,
+        free_bytes: usize,
     ) -> bool {
-        if !has_qk_norm || forced_fp8 == Some("1") {
+        if !has_qk_norm
+            || self.fp8_prefill
+            || forced_batched.is_some()
+            || self.prefill_path.path != PrefillPath::Batched
+            || need_bytes <= free_bytes
+        {
             return false;
         }
-        let mut changed = false;
-        if self.fp8_prefill || self.fp8_decode {
-            self.fp8_prefill = false;
-            self.fp8_decode = false;
-            changed = true;
-        }
-        if forced_batched.is_none() && self.prefill_path.path == PrefillPath::Batched {
-            self.prefill_path = PrefillPathChoice {
-                path: PrefillPath::Serial,
-                reason: "qk-norm model (#3413/#3483)",
-                cc: self.cc,
-            };
-            changed = true;
-        }
-        changed
+        self.prefill_path = PrefillPathChoice {
+            path: PrefillPath::Serial,
+            reason: "qk-norm model, FP16 prefill does not fit (#3715)",
+            cc: self.cc,
+        };
+        true
     }
 }
 
@@ -1180,23 +1201,23 @@ mod pmat3477_fp8_qk_norm_tests {
     }
 
     // The measured case: a QK-norm model on sm_89 loses FP8 (Qwen3-8B under FP8
-    // batched prefill scored cosine −0.0972 vs CPU) AND takes the serial prefill
-    // (the FP16 batched cache OOMed 8B beside apr qa's CPU twin on 24 GB).
+    // batched prefill scored cosine −0.0972 vs CPU) and KEEPS the batched prefill:
+    // the serial one could not prefill 26131 tokens in 600 s (#3715).
     #[test]
-    fn qk_norm_model_gets_serial_prefill_without_fp8() {
+    fn qk_norm_model_loses_fp8_keeps_batched_prefill() {
         let mut p = profile(true);
-        assert!(p.disable_fp8_for_qk_norm(true, None, None));
+        assert!(p.disable_fp8_for_qk_norm(true, None));
         assert!(!p.fp8_prefill && !p.fp8_decode);
-        assert_eq!(p.prefill_path.path, PrefillPath::Serial);
-        assert_eq!(p.prefill_path.reason, "qk-norm model (#3413/#3483)");
+        assert_eq!(p.prefill_path.path, PrefillPath::Batched);
+        assert_eq!(p.prefill_path.reason, "default");
     }
 
     // A model without QK-norm (the qwen2.5 control: 4/4 prompts pass under
     // FP8) keeps FP8 and the batched prefill — a model-class gate, not a fleet default.
     #[test]
     fn no_qk_norm_keeps_fp8_and_batched() {
         let mut p = profile(true);
-        assert!(!p.disable_fp8_for_qk_norm(false, None, None));
+        assert!(!p.disable_fp8_for_qk_norm(false, None));
         assert!(p.fp8_prefill && p.fp8_decode);
         assert_eq!(p.prefill_path.path, PrefillPath::Batched);
     }
@@ -1205,26 +1226,65 @@ mod pmat3477_fp8_qk_norm_tests {
     #[test]
     fn forced_fp8_is_honoured() {
         let mut p = profile(true);
-        assert!(!p.disable_fp8_for_qk_norm(true, Some("1"), None));
+        assert!(!p.disable_fp8_for_qk_norm(true, Some("1")));
         assert!(p.fp8_prefill);
         assert_eq!(p.prefill_path.path, PrefillPath::Batched);
     }
 
-    // An explicit BATCHED_PREFILL keeps its own answer; FP8 still goes off.
+    // An explicit BATCHED_PREFILL=0 keeps its own answer; FP8 still goes off.
     #[test]
-    fn explicit_batched_prefill_is_kept() {
+    fn explicit_serial_prefill_is_kept() {
         let mut p = profile(true);
-        p.prefill_path = select_prefill_path(89, Some("1"));
-        assert!(p.disable_fp8_for_qk_norm(true, None, Some("1")));
+        p.prefill_path = select_prefill_path(89, Some("0"));
+        assert!(p.disable_fp8_for_qk_norm(true, None));
         assert!(!p.fp8_prefill);
-        assert_eq!(p.prefill_path.path, PrefillPath::Batched);
+        assert_eq!(p.prefill_path.path, PrefillPath::Serial);
     }
 
     // Already serial and FP8 already off (GB10 defaults, or FP8_PREFILL=0): no change.
     #[test]
     fn already_serial_and_off_is_not_a_change() {
         let mut p = profile(false);
         p.prefill_path = select_prefill_path(121, None);
-        assert!(!p.disable_fp8_for_qk_norm(true, None, None));
+        assert!(!p.disable_fp8_for_qk_norm(true, None));
+        assert_eq!(p.prefill_path.path, PrefillPath::Serial);
+    }
+
+    // #3715: Qwen3-8B on the 4090 (4.7 GB weights + 15.1 GB FP16 cache + 4.6 GB
+    // reserve > 23.4 GB free) keeps serial; Qwen3-1.7B (1.1 + 3.4 + 4.6) keeps batched.
+    #[test]
+    fn fp16_prefill_that_does_not_fit_keeps_serial() {
+        const GB: usize = 1_000_000_000;
+        let mut p = profile(false);
+        assert!(p.serial_when_fp16_prefill_does_not_fit(
+            true,
+            None,
+            24_400_000_000,
+            23_400_000_000
+        ));
+        assert_eq!(p.prefill_path().path, PrefillPath::Serial);
+
+        let mut p = profile(false);
+        assert!(!p.serial_when_fp16_prefill_does_not_fit(true, None, 9_100_000_000, 23 * GB));
+        assert_eq!(p.prefill_path().path, PrefillPath::Batched);
+    }
+
+    #[test]
+    fn fit_check_leaves_non_qk_norm_fp8_and_forced_paths_alone() {
+        let (need, free) = (30_000_000_000, 1_000_000_000);
+        // (has_qk_norm, fp8, forced): no QK norm, FP8 on, BATCHED_PREFILL=1.
+        for (qk, fp8, forced) in [
+            (false, false, None),
+            (true, true, None),
+            (true, false, Some("1")),
+        ] {
+            let mut p = profile(fp8);
+            assert!(!p.serial_when_fp16_prefill_does_not_fit(qk, forced, need, free));
+            assert_eq!(
+                p.prefill_path().path,
+                PrefillPath::Batched,
+                "{qk} {fp8} {forced:?}"
+            );
+        }
     }
 }

crates/aprender-serve/src/gguf/cuda/mod.rs

@@ -40,6 +40,7 @@ mod forward_qwen35_cuda;
 /// #3714: the Qwen3-MoE decoder resident on the GPU.
 mod forward_qwen3_moe_resident;
 mod generation;
+mod resident;
 mod speculative;
 mod weights;
 
@@ -67,6 +68,18 @@ use super::utils::verbose;
 // IMP-800: CUDA-Accelerated Model Wrapper
 // =============================================================================
 
+/// #3715: device bytes load makes resident, from [`OwnedQuantizedModelCuda::resident_estimate`].
+struct ResidentEstimate {
+    /// f32 K+V bytes of one position over every layer.
+    kv_per_pos: usize,
+    /// Quantized weights.
+    weights: usize,
+    /// FP16/FP8 prefill weight cache (0 when this profile warms none).
+    cache: usize,
+    /// GH-178 workspace reserve plus the chunked-prefill score budget.
+    reserve: usize,
+}
+
 /// Error from CUDA model initialization that preserves the unconsumed model.
 ///
 /// When `OwnedQuantizedModelCuda::new()` fails, the model is returned inside this error
@@ -732,10 +745,9 @@ impl OwnedQuantizedModelCuda {
         if executor.gpu_profile.disable_fp8_for_qk_norm(
             model.config.constraints.has_qk_norm,
             std::env::var("FP8_PREFILL").ok().as_deref(),
-            std::env::var("BATCHED_PREFILL").ok().as_deref(),
         ) {
             eprintln!(
-                "[#3413] architecture '{}' uses per-head QK-norm: FP8 prefill off and serial prefill in use — its FP8 batched prefill fails CPU parity (#3483; FP8_PREFILL=1 / BATCHED_PREFILL=1 override)",
+                "[#3413] architecture '{}' uses per-head QK-norm: FP8 prefill off, FP16 batched prefill — its FP8 batched prefill fails CPU parity (#3483; FP8_PREFILL=1 overrides)",
                 model.config.architecture
             );
         }
@@ -765,6 +777,24 @@ impl OwnedQuantizedModelCuda {
         let num_layers = model.layers.len();
         let num_kv_heads = model.config.num_kv_heads; // PAR-021 GQA support
         let head_dim = model.config.hidden_dim / model.config.num_heads;
+        // #3715: the FP16 batched prefill a QK-norm model now takes must fit beside
+        // its cache; when it does not (Qwen3-8B on 24 GB) the model keeps serial.
+        {
+            let r = Self::resident_estimate(&executor, &model);
+            let need = r.weights + r.cache + r.reserve + r.kv_per_pos * max_seq_len;
+            if executor.gpu_profile.serial_when_fp16_prefill_does_not_fit(
+                model.config.constraints.has_qk_norm,
+                std::env::var("BATCHED_PREFILL").ok().as_deref(),
+                need,
+                memory_info.0,
+            ) {
+                eprintln!(
+                    "[#3715] serial prefill: the FP16 batched prefill needs {:.1} GB and {:.1} GB is free (BATCHED_PREFILL=1 overrides)",
+                    need as f64 / 1e9,
+                    memory_info.0 as f64 / 1e9,
+                );
+            }
+        }
 
         if let Err(error) = Self::configure_executor(&mut executor, &model, max_seq_len) {
             return Err(CudaInitError {
@@ -801,6 +831,46 @@ impl OwnedQuantizedModelCuda {
         Ok(cuda_model)
     }
 
+    /// #3715: what load makes resident on the device: the quantized weights, the
+    /// FP16 (or FP8) prefill weight cache when this profile warms one, and the
+    /// GH-178 reserve plus the chunked-prefill score budget. Also the f32 K+V bytes
+    /// of one position over every layer.
+    fn resident_estimate(
+        executor: &crate::cuda::CudaExecutor,
+        model: &OwnedQuantizedModel,
+    ) -> ResidentEstimate {
+        let c = &model.config;
+        let head_dim = c.hidden_dim / c.num_heads.max(1);
+        let kv_per_pos =
+            2 * c.num_kv_heads * head_dim * std::mem::size_of::<f32>() * model.layers.len();
+        let (weights, elems) = resident::projection_sizes(
+            model
+                .layers
+                .iter()
+                .flat_map(resident::layer_projections)
+                .chain(std::iter::once(&model.lm_head_weight)),
+        );
+        let profile = &executor.gpu_profile;
+        let cache_bytes_per_elem = if profile.fp8_prefill {
+            1
+        } else if std::env::var("HGEMM_PREFILL").as_deref() == Ok("0")
+            || (profile.cc >= 120 && std::env::var("FORCE_FP16_CACHE").as_deref() != Ok("1"))
+        {
+            0
+        } else {
+            2
+        };
+        let cache = cache_bytes_per_elem * elems;
+        let reserve =
+            3_500_000_000 + crate::gguf::cuda::forward_qwen35_cuda::PREFILL_SCORES_BUDGET_BYTES;
+        ResidentEstimate {
+            kv_per_pos,
+            weights,
+            cache,
+            reserve,
+        }
+    }
+
     /// GH-129: Free CPU projection weight copies after GPU preload.
     ///
     /// On Jetson (unified memory), the CPU `Vec<u8>` copies are redundant after