fix(batch): deepseek-v4 serves concurrent requests (one-row prompt batches, CacheList promotion)

CHANGELOG.md

@@ -6,6 +6,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Fixed
+
+- DeepSeek V4 serves concurrent requests: multi-row prompt batches on
+  pooling-cache models failed before prefill, and admission re-merged
+  already-batched caches, killing every request in flight above c=1.
+
 ### Changed
 
 - mlx-kquant floor raised to 0.3.11: MoE prefill gather runs 12-28%

gmlx/apc_pooling.py

@@ -347,6 +347,140 @@ def safe_vlm_maybe_quantize(prompt_cache, quantized_kv_start,
     vlm_common.maybe_quantize_kv_cache = safe_vlm_maybe_quantize
 
 
+def model_has_pools(model) -> bool:
+    """True if the model's cache stack contains a PoolingCache.
+
+    Walks ``make_cache()`` once and memoizes on the model, so the pooled
+    seams can gate without paying a cache build per call.
+    """
+    from .deepseek_v4_cache import PoolingCache
+
+    cached = getattr(model, "_kq_has_pools", None)
+    if cached is not None:
+        return cached
+    make = getattr(model, "make_cache", None)
+    found = False
+    if callable(make):
+        try:
+            stack = list(make() or [])
+        except Exception:
+            stack = []
+        while stack:
+            c = stack.pop()
+            subs = getattr(c, "caches", None)
+            if subs is not None:
+                stack.extend(subs)
+            elif isinstance(c, PoolingCache):
+                found = True
+                break
+    try:
+        model._kq_has_pools = found
+    except Exception:
+        pass
+    return found
+
+
+def install_pooled_prefill_batch_gate() -> None:
+    """Form prompt batches one row at a time on pooling-cache models.
+
+    ``to_batch_cache`` has no PoolingCache arm, so a multi-row prompt batch
+    raises before prefill and every request in it fails. Upstream's
+    single-row fast path sidesteps the conversion entirely (it builds the
+    model's own scalar caches), and the decode side already knows how to
+    join those: ``_extend_cache`` promotes a cache that has ``merge`` and no
+    ``left_padding``, which is exactly PoolingCache, and BatchPoolingCache
+    carries extend/filter/extract. So capping prompt batches at one row
+    yields serial prefill with batched decode, which is where batching pays
+    anyway -- B>1 prefill is not a throughput win at depth (gemma-31b 2x27k
+    measured 130s batched vs 120s serialized).
+
+    Same lever the ckpt tier already pulls in spec_engine, applied for a
+    different reason. Idempotent. Kill switch: GMLX_POOLED_PREFILL_B1=0.
+    """
+    import os
+
+    from mlx_vlm.generate.ar import BatchGenerator
+
+    if getattr(BatchGenerator.__init__, "_kq_pooled_prefill_b1", False):
+        return
+    if os.environ.get("GMLX_POOLED_PREFILL_B1", "1") == "0":
+        return
+
+    _orig_init = BatchGenerator.__init__
+
+    def _gated_init(self, model, processor, **kwargs):
+        _orig_init(self, model, processor, **kwargs)
+        try:
+            if model_has_pools(model) and self.prefill_batch_size != 1:
+                self.prefill_batch_size = 1
+        except Exception:
+            _log.warning("pooled prefill gate failed; continuing",
+                         exc_info=True)
+
+    _gated_init._kq_pooled_prefill_b1 = True
+    BatchGenerator.__init__ = _gated_init
+
+
+def _is_batched_cache(c) -> bool:
+    """True if ``c`` already holds batch rows.
+
+    Batch caches carry ``left_padding``; their scalar counterparts do not.
+    A CacheList carries nothing of its own, so ask its members.
+    """
+    subs = getattr(c, "caches", None)
+    if subs is not None:
+        return any(_is_batched_cache(s) for s in subs)
+    return hasattr(c, "left_padding")
+
+
+def install_batched_cachelist_admission() -> None:
+    """Stop re-merging an already-batched CacheList on admission.
+
+    ``_extend_cache`` lifts a scalar cache into a batch one when it has a
+    ``merge`` and no ``left_padding``. A CacheList defines neither, so a
+    decode batch whose entries are CacheLists (deepseek-v4: rotating window
+    plus two pools per layer) looks scalar on every admission and gets
+    merged a second time. The second merge reaches
+    ``BatchRotatingKVCache.merge``, which calls ``c._temporal_order(c.keys)``
+    -- the scalar signature -- against rows that are already batch caches,
+    and every request in flight dies with a TypeError.
+
+    Replace the promotion test with one that looks through CacheLists.
+    Faithful to the original otherwise: same merge, same in-place extend,
+    same short-circuits. Idempotent. Kill switch:
+    GMLX_BATCHED_CACHELIST_ADMISSION=0.
+    """
+    import os
+
+    from mlx_vlm.generate import ar as _ar
+
+    if getattr(_ar._extend_cache, "_kq_batched_cachelist", False):
+        return
+    if os.environ.get("GMLX_BATCHED_CACHELIST_ADMISSION", "1") == "0":
+        return
+
+    def _promote(c):
+        if _is_batched_cache(c):
+            return c
+        merge = getattr(type(c), "merge", None)
+        return merge([c]) if merge is not None else c
+
+    def _extend_cache(cache_a, cache_b):
+        if not cache_a:
+            return cache_b
+        if not cache_b:
+            return cache_a
+        extended = []
+        for ca, cb in zip(cache_a, cache_b):
+            ca = _promote(ca)
+            ca.extend(_promote(cb))
+            extended.append(ca)
+        return extended
+
+    _extend_cache._kq_batched_cachelist = True
+    _ar._extend_cache = _extend_cache
+
+
 def install_pooled_prompt_kv_quant() -> None:
     """Honor serve kv_bits on pooling-cache models at prompt-batch build.
 
@@ -370,32 +504,7 @@ def install_pooled_prompt_kv_quant() -> None:
     if os.environ.get("GMLX_POOLED_KV_QUANT", "1") == "0":
         return
 
-    from .deepseek_v4_cache import PoolingCache
-
-    def _has_pools(model) -> bool:
-        cached = getattr(model, "_kq_has_pools", None)
-        if cached is not None:
-            return cached
-        make = getattr(model, "make_cache", None)
-        found = False
-        if callable(make):
-            try:
-                stack = list(make() or [])
-            except Exception:
-                stack = []
-            while stack:
-                c = stack.pop()
-                subs = getattr(c, "caches", None)
-                if subs is not None:
-                    stack.extend(subs)
-                elif isinstance(c, PoolingCache):
-                    found = True
-                    break
-        try:
-            model._kq_has_pools = found
-        except Exception:
-            pass
-        return found
+    _has_pools = model_has_pools
 
     _orig_init = ppb.__init__
     _noted = [False]

gmlx/deepseek_v4_cache.py

@@ -520,6 +520,15 @@ def __init__(self, ratio: int, left_padding: List[int]):
         self._prev_kv = None
         self._prev_gate = None
 
+    @property
+    def left_padding(self):
+        # Always zero (the constructor rejects anything else), but present:
+        # batch caches are told apart from their scalar counterparts by
+        # carrying this attribute, and the admission path lifts a cache into
+        # a batch one only when it is missing. Derived rather than stored so
+        # it cannot go stale as extend/filter/extract change the row count.
+        return [0] * len(self._pool_lengths)
+
     @property
     def offset(self):
         return mx.array(self._pool_lengths, dtype=mx.int32)

gmlx/server_patches/__init__.py

@@ -192,13 +192,19 @@ def install_server_patches(cfg, *, reload_fn=None) -> None:
     from ..batch_sched import install_decode_priority_sched
     install_decode_priority_sched()
     from ..apc_pooling import (
+        install_batched_cachelist_admission,
+        install_pooled_prefill_batch_gate,
         install_pooled_prompt_kv_quant,
         install_pooling_apc_support,
         install_safe_kv_quantization,
     )
     install_pooling_apc_support()
     install_safe_kv_quantization()
     install_pooled_prompt_kv_quant()
+    install_pooled_prefill_batch_gate()
+    # Before the model loads, so the cascade stamp wrapper (installed at load
+    # time) wraps this and both survive.
+    install_batched_cachelist_admission()
     install_chat_template_kwargs()
     install_thinking_budget_fix()
     install_openai_stop_sequences()

gmlx/upstream_seams.py

@@ -275,6 +275,14 @@ class Seam:
          "apc_pooling (disk-tier zero-width spill)", critical=True),
     Seam("mlx_vlm.apc", "_clone_cache_entry_for_apc",
          "apc_pooling", critical=True),
+    Seam("mlx_vlm.generate.ar", "BatchGenerator.__init__",
+         "apc_pooling.install_pooled_prefill_batch_gate (prompt batches "
+         "stay B=1 on pooling-cache models; to_batch_cache has no pooled "
+         "arm)", critical=True),
+    Seam("mlx_vlm.generate.ar", "_extend_cache",
+         "apc_pooling.install_batched_cachelist_admission (promotion test "
+         "looks through CacheList so an already-batched one is not merged "
+         "twice)", critical=True),
     Seam("mlx_vlm.apc", "_safetensors_dtype_info",
          "apc_pooling", critical=True),
     # <= 0.6.3 the wrapper delegates to this mlx-lm alias; 0.6.4 inlined it.

tests/test_apc_pooling.py

@@ -452,3 +452,129 @@ def test_pooled_prompt_kv_quant_idempotent(monkeypatch):
 
     install_pooled_prompt_kv_quant()
     assert ppb.__init__ is wrapped
+
+
+# install_pooled_prefill_batch_gate wraps BatchGenerator.__init__; the tests
+# swap in a stand-in class so the wrap is exercised without a real engine.
+DEFAULT_PREFILL_B = 8
+
+
+def _install_gate_on_fake(monkeypatch):
+    from mlx_vlm.generate import ar
+
+    class _FakeBG:
+        def __init__(self, model, processor, **kw):
+            self.model = model
+            self.prefill_batch_size = kw.get(
+                "prefill_batch_size", DEFAULT_PREFILL_B)
+
+    monkeypatch.setattr(ar, "BatchGenerator", _FakeBG)
+    from gmlx.apc_pooling import install_pooled_prefill_batch_gate
+
+    install_pooled_prefill_batch_gate()
+    return _FakeBG
+
+
+def test_prefill_gate_forces_b1_on_pooled(monkeypatch):
+    bg = _install_gate_on_fake(monkeypatch)
+    g = bg(_V4ish(), None)
+    # to_batch_cache has no pooled arm, so multi-row prompt batches would
+    # raise before prefill; one row per batch takes the scalar-cache path.
+    assert g.prefill_batch_size == 1
+
+
+def test_prefill_gate_leaves_non_pooled_alone(monkeypatch):
+    bg = _install_gate_on_fake(monkeypatch)
+    m = SimpleNamespace(make_cache=lambda: [SimpleNamespace(offset=0)])
+    g = bg(m, None)
+    assert g.prefill_batch_size == DEFAULT_PREFILL_B
+
+
+def test_prefill_gate_kill_switch(monkeypatch):
+    monkeypatch.setenv("GMLX_POOLED_PREFILL_B1", "0")
+    bg = _install_gate_on_fake(monkeypatch)
+    assert not getattr(bg.__init__, "_kq_pooled_prefill_b1", False)
+    assert bg(_V4ish(), None).prefill_batch_size == DEFAULT_PREFILL_B
+
+
+def test_prefill_gate_idempotent(monkeypatch):
+    bg = _install_gate_on_fake(monkeypatch)
+    wrapped = bg.__init__
+    from gmlx.apc_pooling import install_pooled_prefill_batch_gate
+
+    install_pooled_prefill_batch_gate()
+    assert bg.__init__ is wrapped
+
+
+def test_model_has_pools_memoizes(monkeypatch):
+    from gmlx.apc_pooling import model_has_pools
+
+    m = _V4ish()
+    calls = []
+    orig = m.make_cache
+    m.make_cache = lambda: (calls.append(1), orig())[1]
+    assert model_has_pools(m) is True
+    assert model_has_pools(m) is True
+    assert len(calls) == 1  # walked once, memoized on the model
+
+
+def test_is_batched_cache_sees_through_cachelist():
+    from gmlx.apc_pooling import _is_batched_cache
+    from gmlx.deepseek_v4_cache import BatchPoolingCache
+
+    scalar = SimpleNamespace(caches=[_pool(rows=0, remainder=0)])
+    batched = SimpleNamespace(caches=[BatchPoolingCache(4, [0, 0])])
+    assert not _is_batched_cache(scalar)
+    assert _is_batched_cache(batched)
+
+
+def test_batched_pool_carries_left_padding():
+    from gmlx.deepseek_v4_cache import BatchPoolingCache
+
+    b = BatchPoolingCache(4, [0, 0, 0])
+    # The admission path lifts only caches missing this attribute.
+    assert b.left_padding == [0, 0, 0]
+    assert not hasattr(_pool(rows=0, remainder=0), "left_padding")
+
+
+def test_admission_does_not_remerge_batched_cachelist(monkeypatch):
+    from mlx_vlm.generate import ar
+
+    calls = []
+
+    class _Sub:
+        def __init__(self, batched):
+            if batched:
+                self.left_padding = [0]
+
+        @classmethod
+        def merge(cls, caches):
+            calls.append("merge")
+            return cls(batched=True)
+
+        def extend(self, other):
+            calls.append("extend")
+
+    class _List:
+        def __init__(self, batched):
+            self.caches = [_Sub(batched)]
+
+        @classmethod
+        def merge(cls, caches):
+            calls.append("list-merge")
+            return cls(batched=True)
+
+        def extend(self, other):
+            calls.append("list-extend")
+
+    monkeypatch.setattr(ar, "_extend_cache", lambda a, b: None, raising=False)
+    from gmlx.apc_pooling import install_batched_cachelist_admission
+
+    install_batched_cachelist_admission()
+    # Already-batched left side: extend only, no second merge.
+    ar._extend_cache([_List(batched=True)], [_List(batched=True)])
+    assert calls == ["list-extend"]
+    # Scalar side still gets lifted.
+    calls.clear()
+    ar._extend_cache([_List(batched=False)], [_List(batched=True)])
+    assert calls == ["list-merge", "list-extend"]