From 3df638a4edccbca326d5a00d45ae7645a2c7afa4 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Tue, 15 Sep 2026 20:18:39 -0700
Subject: [PATCH] fix(stream): lookahead prestage casts bf16 router scores
 before the host copy

---
 CHANGELOG.md                   |  3 +++
 gmlx/stream/lookahead.py       |  6 +++++-
 tests/stream/test_lookahead.py | 23 +++++++++++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c026ecf5..13578f6d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - A stalled or dropped read no longer abandons a `pull`: the transfer
   retries with backoff from the bytes already on disk, tunable by
   `GMLX_PULL_RETRIES` and `GMLX_PULL_TIMEOUT`.
+- Streaming decode with the lookahead prestage no longer crashes on a
+  model whose router scores are bfloat16 ("'bfloat16' is not a valid
+  PEP 3118 buffer format string").
 
 ## [0.4.13] - 2026-09-12
 
diff --git a/gmlx/stream/lookahead.py b/gmlx/stream/lookahead.py
index fd904dbc..2139c72a 100644
--- a/gmlx/stream/lookahead.py
+++ b/gmlx/stream/lookahead.py
@@ -259,7 +259,11 @@ def predict(self, x, variant: str):
         ``indices``-like) for ``dst_li``."""
         if variant == "ratio":
             x = x * self._ratio.astype(x.dtype)
-        return self._router_fn(x)
+        ids, scores = self._router_fn(x)
+        # The hook reads the scores back through numpy, which has no
+        # bfloat16 buffer format; cast inside the lazy graph so the joint
+        # eval covers it and a bf16 router never crashes decode.
+        return ids, scores.astype(mx.float32)
 
 
 class LookaheadProbe:
diff --git a/tests/stream/test_lookahead.py b/tests/stream/test_lookahead.py
index e7b9de4e..23b6789a 100644
--- a/tests/stream/test_lookahead.py
+++ b/tests/stream/test_lookahead.py
@@ -239,6 +239,29 @@ def test_install_and_probe_end_to_end(monkeypatch):
     assert cells[K] == [float(K), float(K)]
 
 
+def test_bf16_router_scores_reach_numpy(monkeypatch):
+    """A bf16 router's scores have no numpy buffer format; the predictor
+    hands them over as float32 so on_call's host copy works."""
+    model = _streaming_model(monkeypatch)
+    monkeypatch.setenv("GMLX_DECODE_LOOKAHEAD_PROBE", "1")
+    assert install_lookahead(model, model.layers, probe=True) == 1
+    la0 = model.layers[0].mlp.switch_mlp._kq_lookahead
+    stock = la0.predictor._router_fn
+
+    def bf16_router(x):
+        ids, scores = stock(x)
+        return ids, scores.astype(mx.bfloat16)
+
+    monkeypatch.setattr(la0.predictor, "_router_fn", bf16_router)
+    x = mx.random.normal((1, 1, DIM))
+    mx.eval(model.layers[0].mlp(x))
+    assert not la0.predictor.dead
+    preds = la0.probe._pending[la0.predictor.dst_li]
+    assert preds and all(isinstance(v, np.ndarray) for v in preds.values())
+    mx.eval(model.layers[1].mlp(x))
+    assert la0.probe._recall
+
+
 def test_dead_predictor_disables_cleanly(monkeypatch):
     model = _streaming_model(monkeypatch)
     install_lookahead(model, model.layers, probe=True)