Studio: account for DeepSeek-V4 compute buffer in context auto-fit (#6940)

* Studio: account for DeepSeek-V4 compute buffer in context auto-fit

DeepSeek-V4-Flash's lightning indexer plus compressed sparse attention reserve a
large context-scaling compute buffer that _compute_buffer_ctx_bytes did not model
(the KQ-mask and dequant-scratch rates both miss it, even with an f16 cache).
Measured on UD-Q4_K_XL at ub 512 it is about 65.5 GiB at 1M context, which the
mask estimate puts near 1.5 GiB, so the auto-fit kept the full 1M train context
and llama-server OOM'd allocating the ~70 GB buffer, then spilled to CPU (~4
tok/s). Add a deepseek4-gated flat plus per-token term so the fit caps the context
(about 256k on a B200) and the model stays fully on GPU.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

studio/backend/core/inference/llama_cpp.py

@@ -3143,6 +3143,12 @@ class LlamaCppBackend:
     _CTX_COMPUTE_BYTES_PER_EMBD = 2.25  # quantized KV, regular attention (dequant scratch)
     _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25  # quantized KV, MLA (compressed attn: measured 0.94x)
     _CTX_COMPUTE_F16_MASK_SAFETY = 1.5  # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok)
+    # DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large
+    # context-scaling compute buffer the rates above miss (present even with an f16
+    # cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without
+    # it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU.
+    _DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3  # ctx-independent indexer scratch
+    _DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000  # per token at ub=512 (~72 GiB at 1M)
 
     def _estimate_compute_buffer_bytes(
         self,
@@ -3192,6 +3198,14 @@ class LlamaCppBackend:
         if n_embd <= 0 or n_ctx <= 0:
             return 0
         ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
+        if getattr(self, "_architecture", None) == "deepseek4":
+            # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires
+            # for any KV type -- the indexer scratch is present even with an f16 cache.
+            ub_scale = ub / self._DEFAULT_N_UBATCH
+            return int(
+                self._DSV4_CTX_COMPUTE_FLAT_BYTES
+                + self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale
+            )
         if _kv_bytes_per_elem(cache_type_kv) < 2.0:
             # Quantized cache: the dequant scratch dominates and scales with n_embd.
             # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on

studio/backend/tests/test_compute_buffer.py

@@ -65,12 +65,14 @@ def _backend(
     vocab = 248320,
     embd = 5120,
     mla = None,
+    arch = None,
 ):
     """Backend with just the dims the compute-buffer estimate reads."""
     b = LlamaCppBackend.__new__(LlamaCppBackend)
     b._vocab_size = vocab
     b._embedding_length = embd
     b._key_length_mla = mla  # non-None -> MLA (compressed attention)
+    b._architecture = arch  # GGUF general.architecture (e.g. 'deepseek4')
     return b
 
 
@@ -290,3 +292,62 @@ class TestContextBufferMLA:
         b = _backend(embd = 6144, mla = 256)
         est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB
         assert est <= 4141 * 1.7
+
+
+class TestContextBufferDSV4:
+    """DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention
+    compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache).
+    Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit
+    must see this so it does not commit the full 1M train context and OOM (spilling
+    to CPU at ~4 tok/s)."""
+
+    _MEASURED_1M_GIB = 65.5  # 70353790464 B compute-graph reserve that OOM'd at 1M ctx
+    GIB = 1024**3
+
+    def test_covers_measured_1m_buffer(self):
+        b = _backend(embd = 4096, arch = "deepseek4")
+        gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
+        assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}"
+
+    def test_not_wildly_over_at_1m(self):
+        # Within ~1.3x of measured so the fit still grants a large (~256k) context.
+        b = _backend(embd = 4096, arch = "deepseek4")
+        gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
+        assert gib <= self._MEASURED_1M_GIB * 1.3
+
+    def test_fires_for_f16_cache(self):
+        # The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must
+        # reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx.
+        dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes(
+            262144, cache_type_kv = "f16"
+        )
+        other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes(
+            262144, cache_type_kv = "f16"
+        )
+        assert dsv4 > 40 * other
+
+    def test_cache_type_independent(self):
+        # Indexer scratch is present for an f16 and a quantized cache alike.
+        b = _backend(embd = 4096, arch = "deepseek4")
+        assert b._compute_buffer_ctx_bytes(
+            262144, cache_type_kv = "f16"
+        ) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
+
+    def test_flat_floor_at_small_ctx(self):
+        # ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB).
+        b = _backend(embd = 4096, arch = "deepseek4")
+        assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0
+
+    def test_scales_with_context_and_ubatch(self):
+        b = _backend(embd = 4096, arch = "deepseek4")
+        assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536)
+        assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes(
+            131072, n_ubatch = 256
+        )
+
+    def test_non_dsv4_unchanged(self):
+        # Regression guard: a non-deepseek4 model keeps the mask-only f16 rate.
+        b = _backend(embd = 4096, arch = "llama")
+        per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000
+        expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY
+        assert per_tok == pytest.approx(expected, rel = 1e-6)