diff --git a/CHANGELOG.md b/CHANGELOG.md index fca533f..04c1131 100644 --- a/CHANGELOG.md +++ b/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% diff --git a/gmlx/apc_pooling.py b/gmlx/apc_pooling.py index 86819b6..b03a438 100644 --- a/gmlx/apc_pooling.py +++ b/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] diff --git a/gmlx/deepseek_v4_cache.py b/gmlx/deepseek_v4_cache.py index 8b109c4..8320e6a 100644 --- a/gmlx/deepseek_v4_cache.py +++ b/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) diff --git a/gmlx/server_patches/__init__.py b/gmlx/server_patches/__init__.py index 68a2a06..beb26ee 100644 --- a/gmlx/server_patches/__init__.py +++ b/gmlx/server_patches/__init__.py @@ -192,6 +192,8 @@ 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, @@ -199,6 +201,10 @@ def install_server_patches(cfg, *, reload_fn=None) -> None: 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() diff --git a/gmlx/upstream_seams.py b/gmlx/upstream_seams.py index 760b69f..353fc9b 100644 --- a/gmlx/upstream_seams.py +++ b/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. diff --git a/tests/test_apc_pooling.py b/tests/test_apc_pooling.py index 67833bc..a5d73c7 100644 --- a/tests/test_apc_pooling.py +++ b/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"]