fix(serve): coarse decode caps, 2048 MoE prefill chunk and fewer per-step host reads on pre-M5 GPUs

CHANGELOG.md

@@ -6,6 +6,21 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Fixed
+
+- A model in memory decodes on the 0.5.1 command-buffer caps again; the
+  0.5.2 caps decoded 4 to 8 percent slower on an M1 Ultra.
+  `GMLX_CB_FINE_DECODE=1` keeps the 0.5.2 caps.
+- Short prompts to a MoE model in memory prefill faster: the prefill chunk
+  pick visited every module of every layer on each request.
+- A client that disconnects before the server has read its request body
+  no longer leaves an error traceback in the server log.
+- Decode spends less host time per token: the memory governor and the
+  command-buffer phase flip no longer repeat their per-step reads while
+  nothing changed.
+- A sparse MoE on M1 through M4 prefills in the 0.5.1 chunk again, a few
+  percent faster than 0.5.2. `GMLX_MOE_WIDE_STEP=4096` keeps the wider chunk.
+
 ## [0.5.2] - 2026-10-25
 
 ### Added

docs/cli.md

@@ -141,7 +141,7 @@ config:
 | `--pin ID_OR_PATH` | None | Load this model at start and never evict it. Repeat the flag for more models. |
 | `--max-tokens N` | Until the model stops | Cap a completion at this many tokens when the request sets no cap. |
 | `--no-family-defaults` | Off | Do not seed each family's model-card sampling under profiles and requests. |
-| `--prefill-step-size N` | `2048`, `4096` for a sparse MoE when memory allows, `8192` when streaming, `4096` for HY4 when streaming | Prefill in chunks of this many tokens. A lower value caps peak memory. |
+| `--prefill-step-size N` | `2048`, `4096` for a sparse MoE from the M5 on when memory allows, `8192` when streaming, `4096` for HY4 when streaming | Prefill in chunks of this many tokens. A lower value caps peak memory. |
 | `--dtype {auto,bfloat16,float16}` | `auto` | Set the activation width. `auto` picks float16 on M1 and M2. |
 | `--decode-prefill-ratio R` | `auto` | Make each prefill chunk wait until decoding streams have had this multiple of its GPU time. `0` restores stock scheduling. |
 | `--prefill-tick-ms MS` | `500` | Give each prefill chunk this wall-clock budget while streams decode, except on a streamed model. `0` never halves a chunk. |

docs/config.md

@@ -643,7 +643,7 @@ also has a [`gmlx serve` flag](cli.md#gmlx-serve).
 
 | Key | Default | Meaning |
 |-----|---------|---------|
-| <a id="serverprefill_step_size"></a>`prefill_step_size` | `2048`, `4096` for a sparse MoE when memory allows, `8192` when streaming, `4096` for HY4 when streaming | Prompt chunk size in tokens. A lower value lowers a long prompt's memory peak and slows prefill. |
+| <a id="serverprefill_step_size"></a>`prefill_step_size` | `2048`, `4096` for a sparse MoE from the M5 on when memory allows, `8192` when streaming, `4096` for HY4 when streaming | Prompt chunk size in tokens. A lower value lowers a long prompt's memory peak and slows prefill. |
 | <a id="serverdtype"></a>`dtype` | `auto` | Type of activations, unquantized weights and KV cache: `bfloat16`, `float16`, or `auto` (`float16` on M1 and M2) |
 | <a id="serverdecode_prefill_ratio"></a>`decode_prefill_ratio` | `auto` | How a new prompt's prefill shares the GPU with generating requests. `auto` keeps generation above half its speed. |
 | <a id="serverprefill_tick_ms"></a>`prefill_tick_ms` | `500` | Longest prefill chunk, in ms, while other requests generate. A longer chunk is halved, except on a streamed model. `0` turns off the halving. |

docs/internals/debug-switches.md

@@ -36,7 +36,7 @@ is already running does not see a variable you export later.
 | `GMLX_DECODE_PREFILL_FLOOR` | Decode-rate floor that `auto` pacing protects, as a share of a stream's batched rate. The default is `0.5`. |
 | `GMLX_PREFILL_TICK_MS` | Same as `server.prefill_tick_ms`. |
 | `GMLX_PREFILL_MIN_STEP` | Smallest chunk, in tokens, that the tick budget may halve a chunk down to. The default is `256`. |
-| `GMLX_MOE_WIDE_STEP` | Prompt chunk of a sparse MoE with 20 GiB free: 128 tokens or fewer per expert in a 2048-token chunk. The default is `4096`, and `2048` keeps the stock chunk. |
+| `GMLX_MOE_WIDE_STEP` | Chunk of a sparse MoE with 20 GiB free and at most 128 tokens per expert per 2048 rows: `4096` from the M5 on, `2048` before. Set it to force a chunk. |
 | `GMLX_DECODE_BATCH` | Requests that decode together in a step. The default is `8`, or the capacity table's widest fit when smaller. `0` restores the upstream 32. |
 | `GMLX_SSE_KEEPALIVE_S` | Interval in seconds of the SSE keepalive comment while a stream is silent. The default is `15`, and `0` disables it. |
 | `GMLX_FAITHFUL_HISTORY=0` | Restore mlx-vlm's stock chat-history rebuild, which drops `reasoning_content` from plain assistant turns, and its stock `/v1/messages` conversion. |
@@ -163,7 +163,7 @@ fusions.
 | `GMLX_DSA_INDEXER=0` | Compute DeepSeek-V4 and V4.1 indexer scores and top-k on the inline fp32 op chain instead of the mlx-kquant kernels. |
 | `GMLX_DSA_INDEXER_Q=0` | Keep the indexer GEMM on fp16 operands where tensor-op hardware would run the int8 kernel. Output is bit-identical. |
 | `GMLX_CB_PHASE=0` | Disable the per-phase command-buffer caps (same output, slower decode). `MLX_MAX_OPS_PER_BUFFER` or `MLX_MAX_MB_PER_BUFFER` set by hand holds in every phase. |
-| `GMLX_CB_COARSE_DECODE=1` | Decode a model in memory on 400 ops per command buffer instead of 50, both with no size cap. Prefill and streamed models keep their caps. |
+| `GMLX_CB_FINE_DECODE=1` | Decode a model in memory on 50 ops per command buffer instead of 400, both with no size cap. Prefill and streamed models keep their caps. |
 | `GMLX_NATIVE_FP` | MXFP4 and NVFP4 layout. `wire` keeps the file bytes, `packed` repacks them, and `auto` picks wire when streaming or past 90% of the working set. |
 | `GMLX_BATCH_INVARIANT=1` | Run small float `nn.Linear` layers on a row-count-invariant kernel. Raw-array routers, router calls under 64 rows and training are not covered. |
 | `GMLX_BATCH_INVARIANT_MAX_OUT` | Most outputs of a float projection that the batch-invariant kernel takes. The default is `512`. |

gmlx/gen/prefill_decay.py

@@ -47,7 +47,8 @@
 
 A MoE whose routed experts average at most 128 rows of a stock chunk, over
 all of a batch's prompts, reads each expert's weights for few rows, so with
-no arch base step its default base doubles to GMLX_MOE_WIDE_STEP (4096).
+no arch base step its default base doubles to 4096 from the M5 on
+(GMLX_MOE_WIDE_STEP names a chunk on any chip).
 The wider chunk's MoE rows are not in the score fit, so it holds only while
 the free working set keeps 20 GiB, read per chunk, and never on a model
 that streams its experts.
@@ -74,8 +75,8 @@
     GMLX_PREFILL_SCORE_PROFILE=0  ignore arch score profiles (dense model
                                     everywhere, pre-profile behavior)
     GMLX_MOE_WIDE_STEP          base step of a MoE with few rows per expert
-                                    (default 4096; 2048 or less keeps the
-                                    stock step)
+                                    (4096 from the M5 on, 2048 before; 2048
+                                    or less keeps the stock step)
     GMLX_PREFILL_FUSED_256=0    price head dim 256 as unfused, as on a GPU
                                     without NAX
     GMLX_PREFILL_HEADROOM_CAP=1 raise the body cap to min(half the live
@@ -926,6 +927,9 @@ def decayed_seed_step(base: int, depth: int, heads: int,
 # takes the wider base step while the free working set keeps the floor, which
 # keeps a wider chunk off a model near the wired ceiling.
 _MOE_WIDE_MAX_ROWS = 128
+# The wide step pays off from this GPU generation on; older chips keep the
+# stock chunk unless GMLX_MOE_WIDE_STEP names one.
+_MOE_WIDE_MIN_GEN = 17
 _MOE_WIDE_FLOOR_BYTES = 20 << 30
 _EXPERTS = ("num_experts", "num_local_experts", "n_routed_experts",
             "moe_num_experts")
@@ -952,12 +956,28 @@ def prompt_rows(batch) -> int:
     return len(getattr(batch, "uids", None) or ()) or 1
 
 
+def _wide_step() -> int:
+    """GMLX_MOE_WIDE_STEP when it names a number; else 4096 from
+    _MOE_WIDE_MIN_GEN on (an unknown generation counts as new) and the
+    stock chunk before."""
+    raw = os.environ.get("GMLX_MOE_WIDE_STEP", "").strip()
+    if raw:
+        try:
+            return int(raw)
+        except ValueError:
+            pass
+    from gmlx.load.dtypes import gpu_arch_gen
+
+    gen = gpu_arch_gen()
+    return 4096 if gen == 0 or gen >= _MOE_WIDE_MIN_GEN else _STOCK_BASE
+
+
 def moe_base_step(model, rows: int = 1) -> int | None:
-    """GMLX_MOE_WIDE_STEP for a MoE whose experts average at most
+    """The wide step (_wide_step) for a MoE whose experts average at most
     _MOE_WIDE_MAX_ROWS rows of a stock chunk of ``rows`` prompts, keep
     their weights resident, and leave _MOE_WIDE_FLOOR_BYTES of free
     working set; else None, and None with the decay off."""
-    wide = env_int("GMLX_MOE_WIDE_STEP", 4096)
+    wide = _wide_step()
     if wide <= _STOCK_BASE or not _enabled():
         return None
     e, k = _moe_widths(model)

gmlx/serve/cb_phase.py

@@ -1,13 +1,13 @@
 """Per-phase MLX command-buffer split caps.
 
-Decode on a model in RAM runs on RAM_DECODE, 50 ops per buffer with no
+Decode on a model in RAM runs on RAM_DECODE, 400 ops per buffer with no
 size cap: at the device's 50 MB cap a step's encode splits into enough
 command buffers that submission blocks on in-flight drain, so the engine's
-async lookahead hides nothing. GMLX_CB_COARSE_DECODE=1 decodes on COARSE
-instead. Deep
-prefill in RAM needs fine buffers: a giant buffer keeps every layer's
-transients live at once and can exhaust GPU memory on multi-thousand-token
-chunks. A streamed model decodes the other way round. Its decode syncs on
+async lookahead hides nothing. GMLX_CB_FINE_DECODE=1 decodes on
+FINE_DECODE, 50 ops with no size cap, instead. Deep prefill in RAM needs
+fine buffers: a giant buffer keeps every layer's transients live at once
+and can exhaust GPU memory on multi-thousand-token chunks. A streamed
+model decodes the other way round. Its decode syncs on
 the host at every layer to read the routing, so a layer's buffers all
 commit before the sync whatever the caps, and finer buffers let the GPU
 start each layer while the host still encodes it. It prefills on the fine
@@ -32,13 +32,15 @@
 
 import os
 import threading
+import weakref
 
 COARSE = (400, 100000)
 FINE = (50, 50)  # the MLX device defaults on a Max or Ultra chip
-# Decode caps of a model in RAM. Across seven models at 512 to 67k keys,
-# (50, 100000) decoded up to 4.6% faster than COARSE, and slower on none
-# outside the run's spread.
-RAM_DECODE = (50, 100000)
+# Decode caps of a model in RAM. FINE_DECODE, 50 ops with no size cap,
+# decodes within spread of COARSE on M3 Max and M5 Max and 4 to 8 pct
+# slower on M1 Ultra, so every chip decodes on COARSE.
+RAM_DECODE = COARSE
+FINE_DECODE = (50, 100000)
 # Decode caps of a model whose experts stream from the drive. On DeepSeek
 # V4.1 Flash served streamed, (50, 50) decoded 0.63 ms per token faster
 # than (400, 100000), output identical.
@@ -52,10 +54,15 @@
 _STREAMED_ATTR = "_gmlx_cb_streamed"
 _ENV_CAPS = ("MLX_MAX_OPS_PER_BUFFER", "MLX_MAX_MB_PER_BUFFER")
 # pre_lift: the caps end_lift sets back, while a lift holds.
+_UNSET = object()
+# model: a weak reference to the model the caps were last set for, None
+# for a flip without a model, _UNSET when the last set had no model.
 _state: dict = {"phase": None, "caps": None, "kq": None, "lifted": False,
-                "pre_lift": None}
+                "pre_lift": None, "model": _UNSET}
 # Engines on separate threads share the device fields: the compare, the
 # write and the record go together, so the record matches the device.
+# flip's repeat check reads phase and model without the lock: a torn
+# pair at worst skips one flip the other engine is overriding anyway.
 _LOCK = threading.Lock()
 
 
@@ -100,11 +107,11 @@ def _streamed(model) -> bool:
 
 def decode_caps(model) -> tuple[int, int]:
     """The decode caps for ``model``: STREAMED_DECODE when its experts
-    stream from the drive, RAM_DECODE otherwise (COARSE under
-    GMLX_CB_COARSE_DECODE=1)."""
+    stream from the drive, RAM_DECODE otherwise (FINE_DECODE under
+    GMLX_CB_FINE_DECODE=1)."""
     if _streamed(model):
         return STREAMED_DECODE
-    return COARSE if os.environ.get("GMLX_CB_COARSE_DECODE") == "1" else RAM_DECODE
+    return FINE_DECODE if os.environ.get("GMLX_CB_FINE_DECODE") == "1" else RAM_DECODE
 
 
 def prefill_caps(model) -> tuple[int, int]:
@@ -113,11 +120,34 @@ def prefill_caps(model) -> tuple[int, int]:
     return STREAMED_PREFILL if _streamed(model) else FINE
 
 
-def _set(caps: tuple[int, int], ph: str | None) -> tuple[int, int] | None:
+def _ref(model):
+    """A weak reference to ``model``, None for no model, _UNSET for an
+    object that takes no weak reference (no repeat shortcut for it)."""
+    if model is None:
+        return None
+    try:
+        return weakref.ref(model)
+    except TypeError:
+        return _UNSET
+
+
+def _same_model(model) -> bool:
+    """``model`` is the model the caps were last set for."""
+    ref = _state["model"]
+    if ref is _UNSET:
+        return False
+    if ref is None:
+        return model is None
+    return model is not None and ref() is model
+
+
+def _set(caps: tuple[int, int], ph: str | None,
+         model=_UNSET) -> tuple[int, int] | None:
     """Set ``caps`` on the device. Dedupes on the caps set, so a phase
     change that keeps the caps costs nothing and a model change that
-    moves them under one phase still applies. Returns the caps the device
-    held before, None when they are not known."""
+    moves them under one phase still applies. Records the phase and the
+    model it was set for. Returns the caps the device held before, None
+    when they are not known."""
     with _LOCK:
         prev = _state["caps"]
         if prev != caps:
@@ -129,6 +159,7 @@ def _set(caps: tuple[int, int], ph: str | None) -> tuple[int, int] | None:
                 prev = (int(held[0]), int(held[1]))
             _state["caps"] = caps
         _state["phase"] = ph
+        _state["model"] = _UNSET if model is _UNSET else _ref(model)
         return prev
 
 
@@ -150,11 +181,14 @@ def _live() -> bool:
 
 def flip(phase: str, model=None) -> None:
     """Set the caps for ``phase``: prefill_caps(model) for prefill,
-    decode_caps(model) for decode."""
+    decode_caps(model) for decode. A repeat of the last phase for the
+    last model returns at once: the engine flips at every step."""
+    if _state["phase"] == phase and _same_model(model):
+        return
     if not _live():
         return
     caps = decode_caps(model) if phase == "decode" else prefill_caps(model)
-    _set(caps, phase)
+    _set(caps, phase, model)
 
 
 def lift() -> None:

gmlx/serve/governor.py

@@ -63,7 +63,9 @@
 it regrows only into headroom the governor is not using), so orange
 and red reclaim decode speed before they fail a row. Registered
 bytes are sampled on a slow cadence (every 16 ticks and at band
-transitions); the per-tick trajectory rides the free counters and
+transitions), and so is the kernel's unwired count while it sits above
+twice the floor with no prompt step near; the per-tick trajectory rides
+the free counters and
 the admission-path rate stash, never a vars() walk.
 
 Threading: everything here runs on the engine thread inside the
@@ -268,6 +270,10 @@ def __init__(self):
         # wired, and the last wall time a stable (non-collapsing) breach
         # was logged
         self.floor_prev = None
+        self.kernel_sampled_tick = None
+        self.kernel_last = None
+        self.kernel_prompt_prev = False
+        self.prompt_time_seen = None
         self.floor_warned_at = 0.0
         # U5 fault injection: (kind, tick) directives already fired
         self.injected: set = set()
@@ -675,6 +681,51 @@ def _measure_rung(gen, st: _GovState) -> None:
     st.rung_rate_before = rate
 
 
+def _prompt_may_step(gen) -> bool:
+    """A prompt chunk can run on this tick: one is in flight, or the queue
+    holds a prompt the gate is not holding and the batch has room to admit
+    it (ar.py admits when completion_batch_size minus the active rows is
+    at least prefill_batch_size)."""
+    if getattr(gen, "_prompt_batch", None) is not None:
+        return True
+    if not getattr(gen, "_unprocessed_sequences", None):
+        return False
+    if admission_hold_reason(gen) is not None:
+        return False
+    room = getattr(gen, "completion_batch_size", None)
+    if room is None:
+        return True
+    return room - batch_rows(gen) >= getattr(gen, "prefill_batch_size", 1)
+
+
+def _kernel_sample(gen, st: "_GovState", floor: float):
+    """The kernel's unwired bytes for this tick. Read live on the first
+    tick, every 16 ticks after, on a tick where a prompt chunk can run
+    and on the tick after any step that ran one (the tick runs before the
+    step, a chunk can wire GBs, and a request inserted mid-step is
+    admitted at the end of the same step), and on every tick while the
+    last read sat under twice the floor; a plain decode tick between
+    reuses the last read, since a decode step moves memory by a few MB
+    and the read is three sysctl calls. A read that failed is retried
+    next tick. A floor of 0 reads nothing."""
+    if floor <= 0:
+        return None
+    last = st.kernel_last
+    prompt_now = _prompt_may_step(gen)
+    st.prompt_time_seen = getattr(gen, "_prompt_time_counter", None)
+    due = (st.kernel_sampled_tick is None
+           or last is None
+           or st.tick_no - st.kernel_sampled_tick >= 16
+           or last < 2 * floor
+           or prompt_now
+           or st.kernel_prompt_prev)
+    st.kernel_prompt_prev = prompt_now
+    if due:
+        st.kernel_last = _kernel_unwired()
+        st.kernel_sampled_tick = st.tick_no
+    return st.kernel_last
+
+
 def _governor_tick(gen) -> None:
     st = _state(gen)
     st.tick_no += 1
@@ -753,7 +804,7 @@ def _governor_tick(gen) -> None:
     # hundred MB from a livelock (2026-08-24: active 87 GB + cache 27 GB
     # on a 128 GB box, governor green throughout).
     floor = _kernel_floor_bytes()
-    unw = _kernel_unwired() if floor > 0 else None
+    unw = _kernel_sample(gen, st, floor)
     _STATS["kernel_unwired_bytes"] = None if unw is None else int(unw)
     if unw is not None and unw < floor:
         _STATS["kernel_floor_reds"] += 1
@@ -769,6 +820,7 @@ def _governor_tick(gen) -> None:
             freed = _evict_registered(frac)
             mx.clear_cache()
             unw2 = _kernel_unwired()
+            st.kernel_last = unw2
         else:
             freed = 0.0
             unw2 = unw
@@ -904,7 +956,11 @@ def _governor_tick(gen) -> None:
 
 def _harvest_tick(gen, st: _GovState, out) -> None:
     """Post-tick: fold emitted tokens per row into the tokens/tick EMA
-    the demand model uses."""
+    the demand model uses, and mark a step that ran a prompt chunk so the
+    next tick reads the kernel live."""
+    seen = st.prompt_time_seen
+    if seen is not None and getattr(gen, "_prompt_time_counter", seen) != seen:
+        st.kernel_prompt_prev = True
     try:
         responses = out[1]
     except Exception: