From 27b615dc4b8246b2c1a7ed085fecd0bdfe783efb Mon Sep 17 00:00:00 2001
From: dusterbloom <32869278+dusterbloom@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:20:16 +0200
Subject: [PATCH] perf(ds4): batch exact prefill through narrow MMVQ

---
 .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu  |  4 +++-
 server/src/deepseek4/deepseek4_backend.cpp     |  5 +++--
 server/src/deepseek4/deepseek4_graph.cpp       | 18 +++++++++++++++---
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
index 82052f611..663c70b5e 100644
--- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -2787,7 +2787,9 @@ static bool ggml_cuda_try_fuse_mul_mat_glu(
         const int64_t ncols = ids ? src1->ne[2] : src1->ne[1];
         if (ggml_cuda_should_use_mmq(
                 src0->type, cc, ncols,
-                ids ? src0->ne[2] : /*n_experts=*/0)) {
+                ids ? src0->ne[2] : /*n_experts=*/0) &&
+            !(ggml_cuda_mmvq_max_ncols_override > 0 &&
+              ncols <= ggml_cuda_mmvq_max_ncols_override)) {
             ggml_cuda_mul_mat_q_pair(
                 ctx, up->src[0], gate->src[0], src1, ids, up, gate);
             ggml_cuda_op_swiglu_ds4(ctx, glu);
diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp
index c8750bb28..4e9a1695f 100644
--- a/server/src/deepseek4/deepseek4_backend.cpp
+++ b/server/src/deepseek4/deepseek4_backend.cpp
@@ -1896,8 +1896,9 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & tokens,
     const bool hybrid_batch_supported =
         !moe_hybrid_ || cfg_.prefill_mode == PrefillAttentionMode::Sparse;
     const int base_chunk =
-        !prefill_attention_mode_is_approximate(cfg_.prefill_mode) ||
-        !hybrid_batch_supported
+        !hybrid_batch_supported ||
+        (cfg_.prefill_mode == PrefillAttentionMode::Exact &&
+         spec_drafter_ != nullptr)
         ? 1
         : std::max(1, std::min(requested_chunk,
                                layer_major_cap));
diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp
index 8986328b1..4578b9bee 100644
--- a/server/src/deepseek4/deepseek4_graph.cpp
+++ b/server/src/deepseek4/deepseek4_graph.cpp
@@ -6965,7 +6965,17 @@ bool deepseek4_step_layer_range(
     // each sub-forward then writes at most one window and, if present, its
     // boundary is the final token. This preserves the same pool/rotate order
     // as sequential execution while retaining safe batched prefixes.
-    const int first_chunk = deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens);
+    const bool exact_prefill_band =
+        cache.prefill_mode == PrefillAttentionMode::Exact &&
+        allow_decode_graph_reuse && !fused_verify_candidate;
+    const int first_chunk = std::min(
+        deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens),
+        exact_prefill_band ? 4 : n_tokens);
+    const bool exact_multi_token_band =
+        exact_prefill_band && n_tokens > 1 && n_tokens <= 4;
+    ScopedCudaGraphOverrides exact_mmvq_scope(
+        /*disable_graphs=*/false,
+        /*mmvq_max_ncols=*/exact_multi_token_band ? 4 : 0);
     if (first_chunk > 0 && first_chunk < n_tokens &&
         !fused_verify_candidate && !heterogeneous_sparse_prefill &&
         !standard_layer_major_prefill) {
@@ -6989,8 +6999,10 @@ bool deepseek4_step_layer_range(
         }
 
         for (int off = 0; off < n_tokens;) {
-            const int chunk = deepseek4_safe_compressor_batch_tokens(
-                w, kv_start + off, n_tokens - off);
+            const int remaining = n_tokens - off;
+            const int chunk = std::min(
+                deepseek4_safe_compressor_batch_tokens(w, kv_start + off, remaining),
+                exact_prefill_band ? 4 : remaining);
             std::vector<float> chunk_hc;
             std::vector<float> chunk_out;
             std::vector<float> chunk_capture;