From 513645e39403bf558f48c753b9a5b809aa18c42f Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 09:56:34 -0700
Subject: [PATCH 1/8] fix(serve): models in memory decode on the coarse caps
again, the prefill chunk pick keeps its streaming verdict, a client gone
before its body is no error
---
CHANGELOG.md | 11 +++++
docs/internals/debug-switches.md | 2 +-
gmlx/serve/cb_phase.py | 21 ++++-----
gmlx/serve/patches/media_gate.py | 17 +++++++-
gmlx/stream/expert_streaming.py | 19 ++++++++-
tests/serve/test_cb_phase.py | 22 +++++-----
tests/serve/test_media_gate.py | 37 ++++++++++++++++
tests/stream/test_streamed_verdict.py | 61 +++++++++++++++++++++++++++
8 files changed, 165 insertions(+), 25 deletions(-)
create mode 100644 tests/stream/test_streamed_verdict.py
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc84cdc8c..94d5e436f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,17 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased]
+### Fixed
+
+- A model in memory decodes on 400 ops per command buffer again, the 0.5.1
+ caps. The 50-op caps of 0.5.2 decoded 4 to 8 percent slower on an M1
+ Ultra and no faster on an M3 Max or M5 Max. `GMLX_CB_FINE_DECODE=1`
+ keeps the 50-op caps.
+- Short prompts to a MoE model in memory prefill about 10 percent faster:
+ the prefill chunk pick walked every expert module 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.
+
## [0.5.2] - 2026-10-25
### Added
diff --git a/docs/internals/debug-switches.md b/docs/internals/debug-switches.md
index 1c9ff16fd..e9f1a29f5 100644
--- a/docs/internals/debug-switches.md
+++ b/docs/internals/debug-switches.md
@@ -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`. |
diff --git a/gmlx/serve/cb_phase.py b/gmlx/serve/cb_phase.py
index 1b3064825..9b68897ba 100644
--- a/gmlx/serve/cb_phase.py
+++ b/gmlx/serve/cb_phase.py
@@ -1,10 +1,10 @@
"""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
+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
@@ -35,10 +35,11 @@
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.
@@ -100,11 +101,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]:
diff --git a/gmlx/serve/patches/media_gate.py b/gmlx/serve/patches/media_gate.py
index 6a5e987d1..e441816b8 100644
--- a/gmlx/serve/patches/media_gate.py
+++ b/gmlx/serve/patches/media_gate.py
@@ -890,6 +890,9 @@ def install_media_gate(allow_urls: bool = False, app=None) -> None:
a URL unless ``allow_urls``, or that sends a form to a route that takes
JSON. Install it before the API-key middleware, so an unauthenticated
request gets 401 and learns nothing from the gate. Idempotent."""
+ from fastapi.responses import Response
+ from starlette.requests import ClientDisconnect
+
global _allow_urls
_allow_urls = bool(allow_urls)
stock = app is None
@@ -906,6 +909,16 @@ def install_media_gate(allow_urls: bool = False, app=None) -> None:
app.add_middleware(_RefuseWebSockets)
async def _media_gate(request, call_next):
+ try:
+ verdict = await _gate(request)
+ except ClientDisconnect:
+ # The client left before its body arrived. Nothing reads the
+ # answer, and a traceback would say nothing.
+ return Response(status_code=400)
+ return verdict if verdict is not None else await call_next(request)
+
+ async def _gate(request):
+ """The refusal for ``request``, or None when the route may run."""
# Every method: a route reads a body whatever the method is.
upload = _is_upload(request)
session = request.scope.get(SESSION_SCOPE_KEY) is not None
@@ -921,7 +934,7 @@ async def _media_gate(request, call_next):
return _refusal(request.url.path, 413,
f"the tensor body is larger than the "
f"{limit >> 20} MiB limit")
- return await call_next(request)
+ return None
raw = await _read_body(request, body_limit(upload, session))
if raw is None:
return _refusal(request.url.path, 413,
@@ -933,7 +946,7 @@ async def _media_gate(request, call_next):
del parsed
except MediaRefused as e:
return _refusal(request.url.path, 400, str(e))
- return await call_next(request)
+ return None
app.middleware_stack = None # allow install after a stack build
app.middleware("http")(_media_gate)
diff --git a/gmlx/stream/expert_streaming.py b/gmlx/stream/expert_streaming.py
index b8e927806..270ca6c72 100644
--- a/gmlx/stream/expert_streaming.py
+++ b/gmlx/stream/expert_streaming.py
@@ -1850,10 +1850,16 @@ def __call__(self, x, indices, *args, **kwargs):
return n_wrapped, offloaded
+_WALK_ATTR = "_kq_streamed_walk"
+
+
def experts_streamed(model) -> bool:
"""True when the install streams this model's experts from the drive.
Reads the install's verdict on the model or a module it wraps, since a
- CPU-only expert codec marks its modules the way streaming does."""
+ CPU-only expert codec marks its modules the way streaming does. A
+ model with no verdict is walked once and the answer kept on it: the
+ walk visits every module of every layer, and the serve path asks at
+ each request. A verdict the install sets later still wins."""
from .installs import wrapper_chain
chain = list(wrapper_chain(model))
@@ -1861,7 +1867,16 @@ def experts_streamed(model) -> bool:
flag = getattr(m, "_kq_streaming", None)
if flag is not None:
return bool(flag)
- return any(moe_streaming_active(m) for m in chain)
+ if not chain:
+ return False
+ walked = getattr(chain[0], _WALK_ATTR, None)
+ if walked is None:
+ walked = any(moe_streaming_active(m) for m in chain)
+ try:
+ object.__setattr__(chain[0], _WALK_ATTR, walked)
+ except (AttributeError, TypeError):
+ pass
+ return walked
def _resolve_prefill_step(model, requested: int | None) -> tuple[int | None, bool]:
diff --git a/tests/serve/test_cb_phase.py b/tests/serve/test_cb_phase.py
index 2818e4fba..5a104b331 100644
--- a/tests/serve/test_cb_phase.py
+++ b/tests/serve/test_cb_phase.py
@@ -26,7 +26,7 @@ def set_cb_caps(ops, mb):
# setenv first, so teardown also drops a key the test sets.
monkeypatch.setenv(k, "")
monkeypatch.delenv(k)
- monkeypatch.delenv("GMLX_CB_COARSE_DECODE", raising=False)
+ monkeypatch.delenv("GMLX_CB_FINE_DECODE", raising=False)
monkeypatch.setitem(cb_phase._state, "kq",
SimpleNamespace(set_cb_caps=set_cb_caps))
monkeypatch.setitem(cb_phase._state, "phase", None)
@@ -305,26 +305,28 @@ def test_chat_installs_the_phased_steps_as_run_does(monkeypatch, tmp_path,
def test_decode_caps_follow_the_streaming_verdict(monkeypatch):
"""A streamed decode syncs on the host at every layer, so it decodes on
- the streamed caps; a model in RAM decodes on RAM_DECODE."""
+ the streamed caps; a model in RAM decodes on RAM_DECODE, the coarse
+ caps on every chip."""
import gmlx.serve.cb_phase as cb_phase
- monkeypatch.delenv("GMLX_CB_COARSE_DECODE", raising=False)
+ monkeypatch.delenv("GMLX_CB_FINE_DECODE", raising=False)
assert cb_phase.STREAMED_DECODE != cb_phase.RAM_DECODE
+ assert cb_phase.RAM_DECODE == cb_phase.COARSE
assert cb_phase.decode_caps(_model(True)) == cb_phase.STREAMED_DECODE
assert cb_phase.decode_caps(_model(False)) == cb_phase.RAM_DECODE
assert cb_phase.decode_caps(None) == cb_phase.RAM_DECODE
@pytest.mark.parametrize("value", ["1", "0"])
-def test_the_coarse_decode_switch_moves_decode_only(monkeypatch, fake_kq, value):
- """GMLX_CB_COARSE_DECODE=1 decodes a model in RAM on COARSE. A streamed
- decode and every prefill keep their caps, and any other value keeps
- RAM_DECODE."""
+def test_the_fine_decode_switch_moves_decode_only(monkeypatch, fake_kq, value):
+ """GMLX_CB_FINE_DECODE=1 decodes a model in RAM on FINE_DECODE. A
+ streamed decode and every prefill keep their caps, and any other value
+ keeps RAM_DECODE."""
import gmlx.serve.cb_phase as cb_phase
- monkeypatch.setenv("GMLX_CB_COARSE_DECODE", value)
- want = cb_phase.COARSE if value == "1" else cb_phase.RAM_DECODE
- assert cb_phase.RAM_DECODE != cb_phase.COARSE
+ monkeypatch.setenv("GMLX_CB_FINE_DECODE", value)
+ want = cb_phase.FINE_DECODE if value == "1" else cb_phase.RAM_DECODE
+ assert cb_phase.FINE_DECODE != cb_phase.RAM_DECODE
assert cb_phase.decode_caps(_model(False)) == want
assert cb_phase.decode_caps(_model(True)) == cb_phase.STREAMED_DECODE
assert cb_phase.prefill_caps(_model(False)) == cb_phase.FINE
diff --git a/tests/serve/test_media_gate.py b/tests/serve/test_media_gate.py
index 6ab326f07..1ee54a7db 100644
--- a/tests/serve/test_media_gate.py
+++ b/tests/serve/test_media_gate.py
@@ -639,3 +639,40 @@ def test_a_body_that_is_not_json_is_refused():
asyncio.run(mg.check_request("POST", "/v1/chat/completions",
{"content-type": "application/x-www-form-urlencoded"},
b'{"model": "m"}'))
+
+
+def test_a_client_that_leaves_before_its_body_is_no_error(caplog):
+ """A client that disconnects before the gate has read its body gets no
+ traceback in the log: the gate answers the dead connection and stops."""
+ import asyncio
+ import logging
+
+ app = FastAPI()
+ mg.install_media_gate(app=app)
+ ran = []
+
+ @app.post("/v1/chat/completions")
+ async def chat(request: Request):
+ ran.append(await request.body())
+ return {}
+
+ scope = {"type": "http", "http_version": "1.1", "method": "POST",
+ "scheme": "http", "path": "/v1/chat/completions",
+ "raw_path": b"/v1/chat/completions", "query_string": b"",
+ "root_path": "", "headers": [(b"content-type", b"application/json"),
+ (b"content-length", b"2")],
+ "client": ("127.0.0.1", 1), "server": ("127.0.0.1", 80)}
+
+ async def receive():
+ return {"type": "http.disconnect"}
+
+ sent = []
+
+ async def send(message):
+ sent.append(message)
+
+ with caplog.at_level(logging.ERROR):
+ asyncio.run(app(scope, receive, send))
+ assert not ran
+ assert not [r for r in caplog.records if r.levelno >= logging.ERROR]
+ assert sent and sent[0]["status"] == 400
diff --git a/tests/stream/test_streamed_verdict.py b/tests/stream/test_streamed_verdict.py
new file mode 100644
index 000000000..68a8c2ce2
--- /dev/null
+++ b/tests/stream/test_streamed_verdict.py
@@ -0,0 +1,61 @@
+"""experts_streamed walks a model without an install verdict once, and the
+install's verdict wins over the walk."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+from gmlx.stream import expert_streaming as es
+
+
+class _Layer:
+ def __init__(self, walks, cpu_only=False):
+ self._walks = walks
+ self._cpu_only = cpu_only
+
+ def modules(self):
+ self._walks.append(self)
+ return [SimpleNamespace(_kq_cpu_only=self._cpu_only)]
+
+
+def _model(walks, cpu_only=False, layers=3):
+ return SimpleNamespace(layers=[_Layer(walks, cpu_only) for _ in range(layers)])
+
+
+def test_a_model_without_a_verdict_is_walked_once():
+ walks = []
+ model = _model(walks)
+ assert es.experts_streamed(model) is False
+ assert len(walks) == 3
+ assert es.experts_streamed(model) is False
+ assert es.experts_streamed(model) is False
+ assert len(walks) == 3
+
+
+def test_the_walk_reads_a_cpu_only_codec_as_streamed():
+ walks = []
+ assert es.experts_streamed(_model(walks, cpu_only=True)) is True
+
+
+def test_the_installs_verdict_wins_over_the_kept_walk():
+ walks = []
+ model = _model(walks)
+ assert es.experts_streamed(model) is False
+ model._kq_streaming = True
+ assert es.experts_streamed(model) is True
+ assert len(walks) == 3
+
+
+def test_the_walk_is_kept_on_the_outer_model_of_a_wrapper_chain():
+ walks = []
+ inner = _model(walks)
+ outer = SimpleNamespace(language_model=inner)
+ assert es.experts_streamed(outer) is False
+ assert es.experts_streamed(outer) is False
+ assert len(walks) == 3
+ assert not hasattr(outer, "_kq_streaming")
+ assert not hasattr(inner, "_kq_streaming")
+
+
+def test_no_model_is_not_streamed():
+ assert es.experts_streamed(None) is False
From bb200948127d334f9043a028b1a9c1c8f4515332 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 10:19:38 -0700
Subject: [PATCH 2/8] perf(serve): the governor reads the kernel's unwired
count every 16 steps above the floor, and a repeat phase flip for the same
model returns at once; review nits of the caps and verdict fixes
---
CHANGELOG.md | 15 ++++----
gmlx/serve/cb_phase.py | 51 +++++++++++++++++++++------
gmlx/serve/governor.py | 25 +++++++++++--
tests/serve/test_cb_phase.py | 45 +++++++++++++++++++++--
tests/serve/test_governor.py | 43 ++++++++++++++++++++++
tests/serve/test_media_gate.py | 33 +++++++++++++++++
tests/stream/test_streamed_verdict.py | 2 +-
7 files changed, 193 insertions(+), 21 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 94d5e436f..8c2028fa8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,14 +8,17 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
### Fixed
-- A model in memory decodes on 400 ops per command buffer again, the 0.5.1
- caps. The 50-op caps of 0.5.2 decoded 4 to 8 percent slower on an M1
- Ultra and no faster on an M3 Max or M5 Max. `GMLX_CB_FINE_DECODE=1`
- keeps the 50-op caps.
-- Short prompts to a MoE model in memory prefill about 10 percent faster:
- the prefill chunk pick walked every expert module on each request.
+- 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.
+- Less host work per decode token: the memory governor reads the kernel's
+ unwired count every 16 steps while it sits above the floor, instead of
+ three sysctl calls on every step, and the command-buffer phase flip
+ returns at once on a repeat of the last phase for the last model.
## [0.5.2] - 2026-10-25
diff --git a/gmlx/serve/cb_phase.py b/gmlx/serve/cb_phase.py
index 9b68897ba..3ac3a2e78 100644
--- a/gmlx/serve/cb_phase.py
+++ b/gmlx/serve/cb_phase.py
@@ -4,10 +4,10 @@
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_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
+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,6 +32,7 @@
import os
import threading
+import weakref
COARSE = (400, 100000)
FINE = (50, 50) # the MLX device defaults on a Max or Ultra chip
@@ -53,8 +54,11 @@
_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.
_LOCK = threading.Lock()
@@ -114,12 +118,36 @@ 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:
+ _state["model"] = _UNSET if model is _UNSET else _ref(model)
prev = _state["caps"]
if prev != caps:
kq = _kq()
@@ -151,11 +179,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:
diff --git a/gmlx/serve/governor.py b/gmlx/serve/governor.py
index d93b925e5..8e68a848c 100644
--- a/gmlx/serve/governor.py
+++ b/gmlx/serve/governor.py
@@ -63,7 +63,8 @@
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
+the floor; 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 +269,8 @@ 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.floor_warned_at = 0.0
# U5 fault injection: (kind, tick) directives already fired
self.injected: set = set()
@@ -675,6 +678,23 @@ def _measure_rung(gen, st: _GovState) -> None:
st.rung_rate_before = rate
+def _kernel_sample(st, floor: float):
+ """The kernel's unwired bytes for this tick. Read live on the first
+ tick, every 16 ticks after, and on every tick while the last read sat
+ under the floor; the ticks between reuse the last read. The read is
+ three sysctl calls, and a decode step moves memory by a few MB. A
+ floor of 0 reads nothing."""
+ if floor <= 0:
+ return None
+ due = (st.kernel_sampled_tick is None
+ or st.tick_no - st.kernel_sampled_tick >= 16
+ or (st.kernel_last is not None and st.kernel_last < floor))
+ 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 +773,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(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 +789,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
diff --git a/tests/serve/test_cb_phase.py b/tests/serve/test_cb_phase.py
index 5a104b331..61acab753 100644
--- a/tests/serve/test_cb_phase.py
+++ b/tests/serve/test_cb_phase.py
@@ -33,6 +33,7 @@ def set_cb_caps(ops, mb):
monkeypatch.setitem(cb_phase._state, "caps", None)
monkeypatch.setitem(cb_phase._state, "lifted", False)
monkeypatch.setitem(cb_phase._state, "pre_lift", None)
+ monkeypatch.setitem(cb_phase._state, "model", cb_phase._UNSET)
return calls
@@ -45,9 +46,16 @@ def held(request, monkeypatch, fake_kq):
return key, fake_kq
+class _Marked:
+ """A model the streaming install marked, in RAM or streamed. A class,
+ not a SimpleNamespace: the flip keeps a weak reference to the model."""
+
+ def __init__(self, streamed):
+ self._kq_streaming = streamed
+
+
def _model(streamed):
- """A model the streaming install marked, in RAM or streamed."""
- return SimpleNamespace(_kq_streaming=streamed)
+ return _Marked(streamed)
@pytest.fixture
@@ -719,3 +727,36 @@ def set_cb_caps(ops, mb):
streamed.join(5)
assert device == [cb_phase.RAM_DECODE, cb_phase.STREAMED_DECODE]
assert cb_phase._state["caps"] == device[-1]
+
+
+def test_a_repeat_flip_for_the_same_model_reads_nothing(monkeypatch, fake_kq):
+ """The engine flips at every step. A repeat of the last phase for the
+ last model returns before the caps are recomputed; a phase change or
+ another model recomputes them."""
+ import gmlx.serve.cb_phase as cb_phase
+
+ picks = []
+ monkeypatch.setattr(cb_phase, "decode_caps",
+ lambda m: (picks.append(("decode", id(m))), cb_phase.RAM_DECODE)[1])
+ monkeypatch.setattr(cb_phase, "prefill_caps",
+ lambda m: (picks.append(("prefill", id(m))), cb_phase.FINE)[1])
+ a, b = _model(False), _model(False)
+ cb_phase.flip("decode", a)
+ cb_phase.flip("decode", a)
+ cb_phase.flip("decode", a)
+ assert picks == [("decode", id(a))]
+ cb_phase.flip("decode", b)
+ assert picks == [("decode", id(a)), ("decode", id(b))]
+ cb_phase.flip("prefill", b)
+ cb_phase.flip("prefill", b)
+ assert picks[-1] == ("prefill", id(b)) and len(picks) == 3
+ cb_phase.lift()
+ cb_phase.flip("prefill", b)
+ assert len(picks) == 4
+ # A model that died and a new one at its address are not the same.
+ cb_phase.flip("decode", _model(False))
+ cb_phase.flip("decode", _model(False))
+ assert len(picks) == 6
+ cb_phase.flip("decode", None)
+ cb_phase.flip("decode", None)
+ assert len(picks) == 7
diff --git a/tests/serve/test_governor.py b/tests/serve/test_governor.py
index d83b3e07f..4d1e5d735 100644
--- a/tests/serve/test_governor.py
+++ b/tests/serve/test_governor.py
@@ -956,3 +956,46 @@ def test_dead_arena_entry_is_replaced(rig, monkeypatch):
assert gov._ARENA_REFS["arena:fixed"]() is arena
gov.unregister_arena(arena)
assert "arena:fixed" not in gov._REG and "arena:fixed" not in gov._ARENA_REFS
+
+
+def test_kernel_floor_reads_every_16_ticks_above_the_floor(rig, monkeypatch):
+ # The read is three sysctl calls on the decode step. Above the floor
+ # it runs on the first tick and every 16th after; a tick in between
+ # reuses the last read.
+ reads = []
+ monkeypatch.setattr(gov, "_kernel_unwired",
+ lambda: (reads.append(1), rig["kernel"])[1])
+ gen = FakeGen(rows=2)
+ for _ in range(33):
+ gov._governor_tick(gen)
+ assert len(reads) == 3
+ assert gov._STATS["kernel_unwired_bytes"] == int(rig["kernel"])
+
+
+def test_kernel_floor_reads_every_tick_under_the_floor(rig, monkeypatch):
+ # A breach is watched tick by tick until a read comes back above the
+ # floor; then the slow cadence resumes.
+ reads = []
+ monkeypatch.setattr(gov, "_kernel_unwired",
+ lambda: (reads.append(1), rig["kernel"])[1])
+ gen = FakeGen(rows=2)
+ gov._governor_tick(gen)
+ assert len(reads) == 1
+ st = gov._state(gen)
+ st.kernel_sampled_tick = st.tick_no - 16 # the next tick is due
+ rig["kernel"] = 2e9
+ gov._governor_tick(gen) # breach: sample, reclaim re-read
+ n = len(reads)
+ assert n >= 2
+ gov._governor_tick(gen) # still under: read again
+ assert len(reads) > n
+ n = len(reads)
+ rig["kernel"] = 40e9
+ gov._governor_tick(gen) # the read that clears it
+ assert len(reads) > n
+ n = len(reads)
+ for _ in range(15):
+ gov._governor_tick(gen)
+ assert len(reads) == n
+ gov._governor_tick(gen)
+ assert len(reads) == n + 1
diff --git a/tests/serve/test_media_gate.py b/tests/serve/test_media_gate.py
index 1ee54a7db..8ad495f1b 100644
--- a/tests/serve/test_media_gate.py
+++ b/tests/serve/test_media_gate.py
@@ -676,3 +676,36 @@ async def send(message):
assert not ran
assert not [r for r in caplog.records if r.levelno >= logging.ERROR]
assert sent and sent[0]["status"] == 400
+
+
+def test_a_disconnect_inside_a_route_still_raises():
+ """The gate answers only a client that left while it read the body. A
+ ClientDisconnect raised by route code is not its to catch."""
+ import asyncio
+
+ from starlette.requests import ClientDisconnect
+
+ app = FastAPI()
+ mg.install_media_gate(app=app)
+
+ @app.post("/v1/chat/completions")
+ async def chat(request: Request):
+ raise ClientDisconnect()
+
+ body = b"{}"
+ scope = {"type": "http", "http_version": "1.1", "method": "POST",
+ "scheme": "http", "path": "/v1/chat/completions",
+ "raw_path": b"/v1/chat/completions", "query_string": b"",
+ "root_path": "", "headers": [(b"content-type", b"application/json"),
+ (b"content-length", str(len(body)).encode())],
+ "client": ("127.0.0.1", 1), "server": ("127.0.0.1", 80)}
+ messages = iter([{"type": "http.request", "body": body, "more_body": False}])
+
+ async def receive():
+ return next(messages)
+
+ async def send(message):
+ pass
+
+ with pytest.raises(ClientDisconnect):
+ asyncio.run(app(scope, receive, send))
diff --git a/tests/stream/test_streamed_verdict.py b/tests/stream/test_streamed_verdict.py
index 68a8c2ce2..7b2389608 100644
--- a/tests/stream/test_streamed_verdict.py
+++ b/tests/stream/test_streamed_verdict.py
@@ -46,7 +46,7 @@ def test_the_installs_verdict_wins_over_the_kept_walk():
assert len(walks) == 3
-def test_the_walk_is_kept_on_the_outer_model_of_a_wrapper_chain():
+def test_a_wrapper_chain_is_walked_once_and_gets_no_install_verdict():
walks = []
inner = _model(walks)
outer = SimpleNamespace(language_model=inner)
From 69bac15f642d2c668fa7f0d53f414e0aaa4d571e Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 10:27:21 -0700
Subject: [PATCH 3/8] fix(serve): the governor reads the kernel live under a
prompt in flight and within twice the floor; the cap record lands with the
phase; trend test reads live
---
CHANGELOG.md | 7 +++----
gmlx/serve/cb_phase.py | 2 +-
gmlx/serve/governor.py | 20 +++++++++++++-------
tests/serve/test_governor.py | 31 ++++++++++++++++++++++++++++++-
4 files changed, 47 insertions(+), 13 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8c2028fa8..e538dfa81 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,10 +15,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
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.
-- Less host work per decode token: the memory governor reads the kernel's
- unwired count every 16 steps while it sits above the floor, instead of
- three sysctl calls on every step, and the command-buffer phase flip
- returns at once on a repeat of the last phase for the last model.
+- 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.
## [0.5.2] - 2026-10-25
diff --git a/gmlx/serve/cb_phase.py b/gmlx/serve/cb_phase.py
index 3ac3a2e78..3fae21e99 100644
--- a/gmlx/serve/cb_phase.py
+++ b/gmlx/serve/cb_phase.py
@@ -147,7 +147,6 @@ def _set(caps: tuple[int, int], ph: str | None,
model it was set for. Returns the caps the device held before, None
when they are not known."""
with _LOCK:
- _state["model"] = _UNSET if model is _UNSET else _ref(model)
prev = _state["caps"]
if prev != caps:
kq = _kq()
@@ -158,6 +157,7 @@ def _set(caps: tuple[int, int], ph: str | 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
diff --git a/gmlx/serve/governor.py b/gmlx/serve/governor.py
index 8e68a848c..a72230cd6 100644
--- a/gmlx/serve/governor.py
+++ b/gmlx/serve/governor.py
@@ -678,17 +678,23 @@ def _measure_rung(gen, st: _GovState) -> None:
st.rung_rate_before = rate
-def _kernel_sample(st, floor: float):
+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, and on every tick while the last read sat
- under the floor; the ticks between reuse the last read. The read is
- three sysctl calls, and a decode step moves memory by a few MB. A
- floor of 0 reads nothing."""
+ tick, every 16 ticks after, on a tick with a prompt in flight (a
+ prompt chunk can wire GBs), 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
due = (st.kernel_sampled_tick is None
+ or last is None
or st.tick_no - st.kernel_sampled_tick >= 16
- or (st.kernel_last is not None and st.kernel_last < floor))
+ or last < 2 * floor
+ or getattr(gen, "_prompt_batch", None) is not None
+ or bool(getattr(gen, "_unprocessed_sequences", None)))
if due:
st.kernel_last = _kernel_unwired()
st.kernel_sampled_tick = st.tick_no
@@ -773,7 +779,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_sample(st, floor)
+ 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
diff --git a/tests/serve/test_governor.py b/tests/serve/test_governor.py
index 4d1e5d735..142962b75 100644
--- a/tests/serve/test_governor.py
+++ b/tests/serve/test_governor.py
@@ -678,6 +678,8 @@ def test_kernel_floor_trend_resets_above_floor(rig, monkeypatch):
gov._governor_tick(gen)
rig["kernel"] = 30e9 # recovered
gov._governor_tick(gen)
+ st = gov._state(gen)
+ st.kernel_sampled_tick = st.tick_no - 16 # the next tick reads live
rig["kernel"] = 6.5e9 # new breach, stable level
gov._governor_tick(gen)
assert failed == []
@@ -983,7 +985,7 @@ def test_kernel_floor_reads_every_tick_under_the_floor(rig, monkeypatch):
assert len(reads) == 1
st = gov._state(gen)
st.kernel_sampled_tick = st.tick_no - 16 # the next tick is due
- rig["kernel"] = 2e9
+ rig["kernel"] = 7e9
gov._governor_tick(gen) # breach: sample, reclaim re-read
n = len(reads)
assert n >= 2
@@ -999,3 +1001,30 @@ def test_kernel_floor_reads_every_tick_under_the_floor(rig, monkeypatch):
assert len(reads) == n
gov._governor_tick(gen)
assert len(reads) == n + 1
+
+
+def test_kernel_floor_reads_live_near_the_floor_and_under_a_prompt(rig, monkeypatch):
+ # Within twice the floor every tick reads; so does a tick with a
+ # prompt in flight, since a prompt chunk can wire GBs at once.
+ reads = []
+ monkeypatch.setattr(gov, "_kernel_unwired",
+ lambda: (reads.append(1), rig["kernel"])[1])
+ gen = FakeGen(rows=2)
+ rig["kernel"] = 15e9 # under 2 x the 8 GB floor
+ for _ in range(5):
+ gov._governor_tick(gen)
+ assert len(reads) == 5
+ rig["kernel"] = 40e9
+ gov._governor_tick(gen) # the read that lifts it
+ n = len(reads)
+ for _ in range(5):
+ gov._governor_tick(gen)
+ assert len(reads) == n
+ gen._prompt_batch = object()
+ for _ in range(3):
+ gov._governor_tick(gen)
+ assert len(reads) == n + 3
+ gen._prompt_batch = None
+ gen._unprocessed_sequences = [object()]
+ gov._governor_tick(gen)
+ assert len(reads) == n + 4
From 578bfb7b6f0899c536150d0e207c018fcf0cb910 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 10:39:30 -0700
Subject: [PATCH 4/8] fix(serve): the governor reads the kernel on the tick
after a prompt step and only when a queued prompt can be admitted; a failed
read retries; the flip's lock-free repeat check is documented
---
gmlx/serve/cb_phase.py | 2 ++
gmlx/serve/governor.py | 39 +++++++++++++++++++++++++++---------
tests/serve/test_governor.py | 27 +++++++++++++++++++++++++
3 files changed, 59 insertions(+), 9 deletions(-)
diff --git a/gmlx/serve/cb_phase.py b/gmlx/serve/cb_phase.py
index 3fae21e99..1c931e88b 100644
--- a/gmlx/serve/cb_phase.py
+++ b/gmlx/serve/cb_phase.py
@@ -61,6 +61,8 @@
"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()
diff --git a/gmlx/serve/governor.py b/gmlx/serve/governor.py
index a72230cd6..761b78852 100644
--- a/gmlx/serve/governor.py
+++ b/gmlx/serve/governor.py
@@ -64,7 +64,8 @@
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), and so is the kernel's unwired count while it sits above
-the floor; the per-tick trajectory rides the free counters and
+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
@@ -271,6 +272,7 @@ def __init__(self):
self.floor_prev = None
self.kernel_sampled_tick = None
self.kernel_last = None
+ self.kernel_prompt_prev = False
self.floor_warned_at = 0.0
# U5 fault injection: (kind, tick) directives already fired
self.injected: set = set()
@@ -678,23 +680,42 @@ 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 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
+ gb = getattr(gen, "_generation_batch", None)
+ room = getattr(gen, "completion_batch_size", None)
+ if gb is None or room is None:
+ return True
+ return room - len(gb) >= 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 with a prompt in flight (a
- prompt chunk can wire GBs), 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."""
+ tick, every 16 ticks after, on a tick where a prompt chunk can run
+ and on the tick after it (the tick runs before the step, and a chunk
+ can wire GBs), 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)
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 getattr(gen, "_prompt_batch", None) is not None
- or bool(getattr(gen, "_unprocessed_sequences", None)))
+ 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
diff --git a/tests/serve/test_governor.py b/tests/serve/test_governor.py
index 142962b75..dd348b528 100644
--- a/tests/serve/test_governor.py
+++ b/tests/serve/test_governor.py
@@ -37,6 +37,8 @@ def __init__(self, rows=2, rate=1e6, live=8e9):
self._kq_admit_live_bytes = live
self._kq_admit_n_caches = 2
self.prefill_step_size = 2048
+ self.completion_batch_size = 8
+ self.prefill_batch_size = 1
self.removed = []
def remove(self, uid):
@@ -1028,3 +1030,28 @@ def test_kernel_floor_reads_live_near_the_floor_and_under_a_prompt(rig, monkeypa
gen._unprocessed_sequences = [object()]
gov._governor_tick(gen)
assert len(reads) == n + 4
+ # The tick runs before the step: a one-chunk prompt admitted on the
+ # last tick has wired its KV by now, so this tick reads too.
+ gen._unprocessed_sequences = []
+ gov._governor_tick(gen)
+ assert len(reads) == n + 5
+ gov._governor_tick(gen)
+ assert len(reads) == n + 5
+ # A queued prompt the full batch cannot admit runs no chunk: no read.
+ gen._unprocessed_sequences = [object()]
+ gen.completion_batch_size = len(gen._generation_batch)
+ for _ in range(3):
+ gov._governor_tick(gen)
+ assert len(reads) == n + 5
+
+
+def test_kernel_floor_retries_a_failed_read_next_tick(rig, monkeypatch):
+ samples = iter([None, 40e9]) # a third read would raise
+ monkeypatch.setattr(gov, "_kernel_unwired", lambda: next(samples))
+ gen = FakeGen(rows=2)
+ gov._governor_tick(gen) # sysctl failed: None
+ assert gen._kq_governor.kernel_last is None
+ gov._governor_tick(gen) # retried, not reused
+ assert gen._kq_governor.kernel_last == 40e9
+ gov._governor_tick(gen) # above 2 x floor: reused
+ assert gen._kq_governor.kernel_last == 40e9
From e03f08d019f031aa450aafb026bdb48c2f31ebd2 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 10:45:04 -0700
Subject: [PATCH 5/8] fix(serve): the governor reads the kernel on the tick
after any step that ran a prompt chunk, including a request admitted
mid-step; a held or roomless queue is not a prompt step
---
gmlx/serve/governor.py | 34 +++++++++++++++++++++-------------
tests/serve/test_governor.py | 12 ++++++++++++
2 files changed, 33 insertions(+), 13 deletions(-)
diff --git a/gmlx/serve/governor.py b/gmlx/serve/governor.py
index 761b78852..6a91e8f7f 100644
--- a/gmlx/serve/governor.py
+++ b/gmlx/serve/governor.py
@@ -273,6 +273,7 @@ def __init__(self):
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()
@@ -682,33 +683,36 @@ def _measure_rung(gen, st: _GovState) -> None:
def _prompt_may_step(gen) -> bool:
"""A prompt chunk can run on this tick: one is in flight, or the queue
- holds a prompt 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)."""
+ 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
- gb = getattr(gen, "_generation_batch", None)
+ if admission_hold_reason(gen) is not None:
+ return False
room = getattr(gen, "completion_batch_size", None)
- if gb is None or room is None:
+ if room is None:
return True
- return room - len(gb) >= getattr(gen, "prefill_batch_size", 1)
+ 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 it (the tick runs before the step, and a chunk
- can wire GBs), 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."""
+ 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
@@ -952,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:
diff --git a/tests/serve/test_governor.py b/tests/serve/test_governor.py
index dd348b528..0cce2b6b7 100644
--- a/tests/serve/test_governor.py
+++ b/tests/serve/test_governor.py
@@ -39,6 +39,7 @@ def __init__(self, rows=2, rate=1e6, live=8e9):
self.prefill_step_size = 2048
self.completion_batch_size = 8
self.prefill_batch_size = 1
+ self._prompt_time_counter = 0.0
self.removed = []
def remove(self, uid):
@@ -1043,6 +1044,17 @@ def test_kernel_floor_reads_live_near_the_floor_and_under_a_prompt(rig, monkeypa
for _ in range(3):
gov._governor_tick(gen)
assert len(reads) == n + 5
+ gen._unprocessed_sequences = []
+ # A request inserted during the step is admitted and, when it fits
+ # one chunk, fully prefilled before the step returns: the harvest
+ # sees the prompt-time counter move and the next tick reads.
+ gov._governor_tick(gen)
+ gen._prompt_time_counter += 0.05
+ gov._harvest_tick(gen, gen._kq_governor, ([], []))
+ gov._governor_tick(gen)
+ assert len(reads) == n + 6
+ gov._governor_tick(gen)
+ assert len(reads) == n + 6
def test_kernel_floor_retries_a_failed_read_next_tick(rig, monkeypatch):
From 706ab410cb8284338859f8526317541477be89c4 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 14:49:42 -0700
Subject: [PATCH 6/8] perf(prefill): the MoE wide prompt chunk applies from the
M5 on; M1 through M4 keep the 0.5.1 chunk unless GMLX_MOE_WIDE_STEP names one
---
CHANGELOG.md | 3 +++
docs/internals/debug-switches.md | 2 +-
gmlx/gen/prefill_decay.py | 18 ++++++++++++++++--
tests/gen/test_prefill_decay.py | 17 +++++++++++++++++
4 files changed, 37 insertions(+), 3 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e538dfa81..4f1792c31 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
- 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 long prompts in the 0.5.1 chunk
+ again; the wider 0.5.2 chunk, which pays off from the M5 on, cost
+ a few percent of prefill there. `GMLX_MOE_WIDE_STEP=4096` keeps it.
## [0.5.2] - 2026-10-25
diff --git a/docs/internals/debug-switches.md b/docs/internals/debug-switches.md
index e9f1a29f5..31f73ea1d 100644
--- a/docs/internals/debug-switches.md
+++ b/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` | Prompt chunk of a sparse MoE with 20 GiB free and few tokens per expert: `4096` from the M5 on, `2048` before. Set it to use either chunk on any chip. |
| `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. |
diff --git a/gmlx/gen/prefill_decay.py b/gmlx/gen/prefill_decay.py
index dbb95b9ca..5bdc48fa1 100644
--- a/gmlx/gen/prefill_decay.py
+++ b/gmlx/gen/prefill_decay.py
@@ -926,6 +926,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 +955,23 @@ def prompt_rows(batch) -> int:
return len(getattr(batch, "uids", None) or ()) or 1
+def _wide_step() -> int:
+ """GMLX_MOE_WIDE_STEP when set; else 4096 from _MOE_WIDE_MIN_GEN on
+ (an unknown generation counts as new) and the stock chunk before."""
+ if os.environ.get("GMLX_MOE_WIDE_STEP"):
+ return env_int("GMLX_MOE_WIDE_STEP", 4096)
+ 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)
diff --git a/tests/gen/test_prefill_decay.py b/tests/gen/test_prefill_decay.py
index 631bc6d5f..59c14b803 100644
--- a/tests/gen/test_prefill_decay.py
+++ b/tests/gen/test_prefill_decay.py
@@ -1025,14 +1025,31 @@ def _moe(e, k, nested=False):
@pytest.fixture
def resident(monkeypatch):
+ import gmlx.load.dtypes as dtypes
import gmlx.stream.expert_streaming as es
monkeypatch.setattr(es, "experts_streamed", lambda m: False)
+ monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 17)
monkeypatch.delenv("PREFILL_STEP_SIZE", raising=False)
monkeypatch.delenv("GMLX_MOE_WIDE_STEP", raising=False)
_headroom(monkeypatch, 40 * GB)
+def test_moe_wide_step_waits_for_a_new_enough_gpu(monkeypatch, resident):
+ import gmlx.load.dtypes as dtypes
+
+ # one generation under the line keeps the stock chunk by default
+ monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 16)
+ assert pd._run_base(_moe(256, 8), None, None) == 2048
+ # an explicit GMLX_MOE_WIDE_STEP applies on any chip
+ monkeypatch.setenv("GMLX_MOE_WIDE_STEP", "4096")
+ assert pd._run_base(_moe(256, 8), None, None) == 4096
+ monkeypatch.delenv("GMLX_MOE_WIDE_STEP")
+ # an unknown generation counts as new
+ monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 0)
+ assert pd._run_base(_moe(256, 8), None, None) == 4096
+
+
def test_moe_with_few_rows_per_expert_takes_the_wide_step(resident):
# 256 experts at top 6 or 8 and 128 at top 8: 48, 64 and 128 rows per
# expert of a 2048 chunk
From cc54ed6bf92108a4fee7cf1ffd4b1a0ce59dafc5 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 14:55:59 -0700
Subject: [PATCH 7/8] fix(prefill): the resident MoE budget prices the chunk
prefill runs on this GPU; a GMLX_MOE_WIDE_STEP that is not a number counts as
unset; help and docs name the M5 line
---
CHANGELOG.md | 5 ++---
docs/cli.md | 2 +-
docs/config.md | 2 +-
docs/internals/debug-switches.md | 2 +-
gmlx/gen/prefill_decay.py | 20 +++++++++++++-------
gmlx/serve/server.py | 3 ++-
gmlx/stream/budget.py | 12 +++++++-----
tests/gen/test_prefill_decay.py | 4 ++++
tests/stream/test_budget.py | 19 ++++++++++++++++++-
9 files changed, 49 insertions(+), 20 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4f1792c31..ba1e8ded3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,9 +18,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
- 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 long prompts in the 0.5.1 chunk
- again; the wider 0.5.2 chunk, which pays off from the M5 on, cost
- a few percent of prefill there. `GMLX_MOE_WIDE_STEP=4096` keeps it.
+- 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
diff --git a/docs/cli.md b/docs/cli.md
index 4fe187760..7360db1e5 100644
--- a/docs/cli.md
+++ b/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. |
diff --git a/docs/config.md b/docs/config.md
index eb5d2bf70..f688a2977 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -643,7 +643,7 @@ also has a [`gmlx serve` flag](cli.md#gmlx-serve).
| Key | Default | Meaning |
|-----|---------|---------|
-| `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. |
+| `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. |
| `dtype` | `auto` | Type of activations, unquantized weights and KV cache: `bfloat16`, `float16`, or `auto` (`float16` on M1 and M2) |
| `decode_prefill_ratio` | `auto` | How a new prompt's prefill shares the GPU with generating requests. `auto` keeps generation above half its speed. |
| `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. |
diff --git a/docs/internals/debug-switches.md b/docs/internals/debug-switches.md
index 31f73ea1d..a352b538c 100644
--- a/docs/internals/debug-switches.md
+++ b/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 and few tokens per expert: `4096` from the M5 on, `2048` before. Set it to use either chunk on any chip. |
+| `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. |
diff --git a/gmlx/gen/prefill_decay.py b/gmlx/gen/prefill_decay.py
index 5bdc48fa1..6f0d20d14 100644
--- a/gmlx/gen/prefill_decay.py
+++ b/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
@@ -956,10 +957,15 @@ def prompt_rows(batch) -> int:
def _wide_step() -> int:
- """GMLX_MOE_WIDE_STEP when set; else 4096 from _MOE_WIDE_MIN_GEN on
- (an unknown generation counts as new) and the stock chunk before."""
- if os.environ.get("GMLX_MOE_WIDE_STEP"):
- return env_int("GMLX_MOE_WIDE_STEP", 4096)
+ """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()
diff --git a/gmlx/serve/server.py b/gmlx/serve/server.py
index 972a47b60..9cd94b33c 100644
--- a/gmlx/serve/server.py
+++ b/gmlx/serve/server.py
@@ -1985,7 +1985,8 @@ def _add_serve_args(ap: argparse.ArgumentParser) -> None:
"reload restores).")
ap.add_argument("--prefill-step-size", type=int, default=None, metavar="N",
help="Prefill chunk size in tokens (default 2048, 4096 for a sparse MoE "
- "when memory allows, 8192 for a streamed model) - lower "
+ "from the M5 on when memory allows, 8192 for a streamed "
+ "model) - lower "
"it to cap peak memory on long prompts, at some "
"prefill-throughput cost. Also via PREFILL_STEP_SIZE; "
"config mode: server.prefill_step_size.")
diff --git a/gmlx/stream/budget.py b/gmlx/stream/budget.py
index 91255890e..e6c3f40aa 100644
--- a/gmlx/stream/budget.py
+++ b/gmlx/stream/budget.py
@@ -264,19 +264,21 @@ def resident_prefill_rows(step: int | None = None,
"""Rows of the widest prefill chunk serve takes for a model whose
experts stay in memory, as ``prefill_decay._run_base`` takes it: the
requested ``step`` or ``PREFILL_STEP_SIZE`` when set, else the wide
- MoE step (``GMLX_MOE_WIDE_STEP``) for a MoE of ``experts`` routed
- experts, ``used`` a token, that ``prefill_decay.moe_base_step``
- widens, else the stock 2048. An arch with its own step (its step env,
+ MoE step (``prefill_decay._wide_step``: ``GMLX_MOE_WIDE_STEP``, else
+ 4096 from the M5 on) for a MoE of ``experts`` routed experts, ``used``
+ a token, that ``prefill_decay.moe_base_step`` widens, else the stock
+ 2048. An arch with its own step (its step env,
else ``_ARCH_RESIDENT_STEP``) takes the wider of that and the rest,
since a batch its profile is off for takes the rest. A resident
prefill merges no tail."""
- from gmlx.gen.prefill_decay import _MOE_WIDE_MAX_ROWS, _STOCK_BASE
+ from gmlx.gen.prefill_decay import (_MOE_WIDE_MAX_ROWS, _STOCK_BASE,
+ _wide_step)
if not (isinstance(step, int) and step > 0):
step = env_int("PREFILL_STEP_SIZE", 0)
if step > 0:
return step
- wide = env_int("GMLX_MOE_WIDE_STEP", 4096)
+ wide = _wide_step()
rest = (wide if wide > _STOCK_BASE and experts and used
and _STOCK_BASE * used <= _MOE_WIDE_MAX_ROWS * experts
else _STOCK_BASE)
diff --git a/tests/gen/test_prefill_decay.py b/tests/gen/test_prefill_decay.py
index 59c14b803..8eef6b83a 100644
--- a/tests/gen/test_prefill_decay.py
+++ b/tests/gen/test_prefill_decay.py
@@ -1048,6 +1048,10 @@ def test_moe_wide_step_waits_for_a_new_enough_gpu(monkeypatch, resident):
# an unknown generation counts as new
monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 0)
assert pd._run_base(_moe(256, 8), None, None) == 4096
+ # a value that is not a number counts as unset
+ monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 16)
+ monkeypatch.setenv("GMLX_MOE_WIDE_STEP", "yes")
+ assert pd._run_base(_moe(256, 8), None, None) == 2048
def test_moe_with_few_rows_per_expert_takes_the_wide_step(resident):
diff --git a/tests/stream/test_budget.py b/tests/stream/test_budget.py
index a38d0c83e..1663174b6 100644
--- a/tests/stream/test_budget.py
+++ b/tests/stream/test_budget.py
@@ -69,8 +69,25 @@ def test_room_prices_ctx_width_transient_reserve(header, monkeypatch):
@pytest.fixture
def chunk_env(monkeypatch):
- for k in ("PREFILL_STEP_SIZE", "GMLX_STREAM_PREFILL_TAIL_MERGE"):
+ import gmlx.load.dtypes as dtypes
+
+ for k in ("PREFILL_STEP_SIZE", "GMLX_STREAM_PREFILL_TAIL_MERGE",
+ "GMLX_MOE_WIDE_STEP"):
monkeypatch.delenv(k, raising=False)
+ monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 17)
+
+
+def test_a_moe_kept_in_memory_prices_the_chunk_of_its_gpu(chunk_env, monkeypatch):
+ import gmlx.load.dtypes as dtypes
+
+ # one generation under the wide step's line: the stock chunk, as
+ # prefill_decay runs it there
+ monkeypatch.setattr(dtypes, "gpu_arch_gen", lambda: 16)
+ assert budget.resident_prefill_rows(model_type="qwen3_5_moe",
+ experts=256, used=8) == 2048
+ monkeypatch.setenv("GMLX_MOE_WIDE_STEP", "4096")
+ assert budget.resident_prefill_rows(model_type="qwen3_5_moe",
+ experts=256, used=8) == 4096
def test_chunk_rows_follow_the_streamed_chunk_width(chunk_env, monkeypatch):
From 26462a43acdd730cc01bdb4ffe0eccea922968e0 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 26 Oct 2026 06:08:41 -0700
Subject: [PATCH 8/8] test(models): the minicpm-v window merger test allows for
tf32 gemm on m5
---
tests/models/test_minicpmv_vlm.py | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/tests/models/test_minicpmv_vlm.py b/tests/models/test_minicpmv_vlm.py
index 2be5a63de..278723964 100644
--- a/tests/models/test_minicpmv_vlm.py
+++ b/tests/models/test_minicpmv_vlm.py
@@ -500,6 +500,9 @@ def test_gelu_kinds_match_the_reference():
def test_window_merger_folds_each_2x2_window_in_row_major_order():
+ # The merger folds all windows in one f32 GEMM, which M5 runs at TF32
+ # (about 1e-3 here), while each reference window is a 1-row GEMV at
+ # f32. A wrong window or position order moves the output far more.
mx.random.seed(0)
vc = _vision()
vm = VitMerger(vc)
@@ -516,7 +519,7 @@ def test_window_merger_folds_each_2x2_window_in_row_major_order():
want = vm.linear_2(vm.act(vm.linear_1(
vm.pre_norm(win.reshape(1, -1)))))[0] + win.mean(axis=0)
np.testing.assert_allclose(np.asarray(got[r * (w // 2) + c]),
- np.asarray(want), atol=1e-5)
+ np.asarray(want), atol=5e-3)
def test_tower_returns_one_token_per_4x4_patches_in_crop_order():