(drop arena-miss experts, keep "
"mass P) |")
+ lines.append(" # moe_prestage: (skip routed path with "
"probability P)")
if not models:
diff --git a/gmlx/generation.py b/gmlx/generation.py
index 4b876db..3767e01 100644
--- a/gmlx/generation.py
+++ b/gmlx/generation.py
@@ -432,7 +432,7 @@ def generate(
window=over_generation, inject_critique=inject_critique,
template_kwargs=critique_tk, log_path=over_generation_log,
orig_prompt=_over_prompt, label=over_label,
- params=params, verbose=verbose,
+ params=params, verbose=verbose, reasoning=reasoning,
)
close_progress = None
@@ -553,13 +553,16 @@ def _generate_over(
over_logits_processors=None,
orig_prompt=None,
label=None,
+ reasoning=None,
):
"""Over-generation probe (experimental). Phase 1 generates to the natural
stop (the seam); phase 2 continues from phase 1's KV cache, either forcing
``window`` free tokens past the seam or injecting a follow-up critique turn
and answering it. ``inject_critique`` selects the mode; when both it and
``window`` are set, ``window`` caps the injected reply. Returns the full
- text and appends a JSONL record to ``log_path`` when set. See overgen.py."""
+ text and appends a JSONL record to ``log_path`` when set. See overgen.py.
+ ``reasoning`` styles phase 1's verbose stream like :func:`generate`'s;
+ phase 2 prints raw (the forced continuation is the thing being probed)."""
from mlx_lm.generate import stream_generate
from mlx_lm.models.cache import make_prompt_cache
@@ -580,6 +583,12 @@ def _generate_over(
# the seam is detected here (and its token held out of the kept text)
# instead of stream_generate ending the call.
pre, seam, last = [], None, None
+ emit = close_emit = None
+ if verbose:
+ # Phase 1 is a normal reply: style it like the main verbose path
+ # (thinking rendering included). Phase 2 stays raw: the forced
+ # continuation is the thing being probed.
+ emit, close_emit = _verbose_emitter(prompt, tokenizer, reasoning)
with suppressed_eos(tokenizer):
for r in stream_generate(
model, tokenizer, prompt, max_tokens=max_tokens,
@@ -595,8 +604,10 @@ def _generate_over(
}
break
pre.append(r.text)
- if verbose:
- print(r.text, end="", flush=True)
+ if emit is not None:
+ emit(r.text)
+ if close_emit is not None:
+ close_emit()
pre_text = "".join(pre)
phase1_last = last # the seam response; phase 2 reassigns `last` below
diff --git a/gmlx/loader.py b/gmlx/loader.py
index 6d99af3..8f7ae3b 100644
--- a/gmlx/loader.py
+++ b/gmlx/loader.py
@@ -1609,7 +1609,8 @@ def __call__(self, x, indices, *args, **kwargs):
# adaptive hot-set refresh live here.
gt.on_layer_entry(
self._kq_li,
- getattr(self, "_kq_miss_shed", None))
+ None if getattr(self, "_kq_in_split", False)
+ else getattr(self, "_kq_miss_shed", None))
if (
gt_live
and scores_arg is not None
@@ -1693,7 +1694,13 @@ def __call__(self, x, indices, *args, **kwargs):
t0 = time.perf_counter() if ph is not None else 0.0
if n_tokens == 1:
dfr.ensure_wired()
- ms = getattr(self, "_kq_miss_shed", None)
+ # Miss-shed is decode-only: a single-token leaf of an
+ # arena token split is prefill work, and a shedding
+ # leaf would return a mixed rank-3 output next to a
+ # clean leaf's per-expert rank-4 - the reassembly
+ # concatenate cannot take both.
+ ms = (None if getattr(self, "_kq_in_split", False)
+ else getattr(self, "_kq_miss_shed", None))
sc_f32 = None
if (ms is not None and scores_arg is not None
and n_tokens == 1):
@@ -1747,8 +1754,17 @@ def __call__(self, x, indices, *args, **kwargs):
# layer's gather and the next layers' every-token
# work compute - speculation never competes with
# demand traffic for the SSD.
- for _dst, _ids in la_pred.items():
- dfr.prestage(_dst, _ids)
+ la_keep = (
+ ms if getattr(
+ self, "_kq_prestage_keepers", False)
+ else None)
+ for _dst, (_ids, _sc) in la_pred.items():
+ if la_keep is not None:
+ dfr.prestage(
+ _dst, _ids, keep_mass=la_keep,
+ pred_scores=_sc)
+ else:
+ dfr.prestage(_dst, _ids)
if ph is not None:
ph["prestage"] += time.perf_counter() - t2
if slots is not None:
@@ -1808,13 +1824,21 @@ def __call__(self, x, indices, *args, **kwargs):
if sliceable:
half = n_tokens // 2
parts = []
- for sl in (slice(0, half),
- slice(half, n_tokens)):
- t = tuple(
- [slice(None)] * ax + [sl])
- parts.append(self.__call__(
- x[t], indices[t],
- *[a[t] for a in orig]))
+ prev_split = getattr(
+ self, "_kq_in_split", False)
+ object.__setattr__(
+ self, "_kq_in_split", True)
+ try:
+ for sl in (slice(0, half),
+ slice(half, n_tokens)):
+ t = tuple(
+ [slice(None)] * ax + [sl])
+ parts.append(self.__call__(
+ x[t], indices[t],
+ *[a[t] for a in orig]))
+ finally:
+ object.__setattr__(
+ self, "_kq_in_split", prev_split)
return mx.concatenate(parts, axis=ax)
wedged = dfr is not None and dfr.wedged_at(self._kq_li)
if wedged and dfr.has_dead(self._kq_li):
diff --git a/gmlx/lookahead.py b/gmlx/lookahead.py
index 666b94c..be904fa 100644
--- a/gmlx/lookahead.py
+++ b/gmlx/lookahead.py
@@ -147,7 +147,8 @@ def _gate_module_select(gate):
def select(x):
inds, weights = call(gate, x)
order = mx.argsort(-weights, axis=-1)
- return mx.take_along_axis(inds, order, axis=-1)
+ return (mx.take_along_axis(inds, order, axis=-1),
+ mx.take_along_axis(weights, order, axis=-1))
return select
@@ -162,12 +163,15 @@ def _sigmoid_bias_select(mod):
k = mod.args.num_experts_per_tok
def select(x):
- choice = mx.sigmoid(mod.gate(x.astype(mx.float32)))
- choice = choice + mod.e_score_correction_bias
+ scores = mx.sigmoid(mod.gate(x.astype(mx.float32)))
+ choice = scores + mod.e_score_correction_bias
inds = mx.argpartition(-choice, kth=k - 1, axis=-1)[..., :k]
ch = mx.take_along_axis(choice, inds, axis=-1)
order = mx.argsort(-ch, axis=-1)
- return mx.take_along_axis(inds, order, axis=-1)
+ inds = mx.take_along_axis(inds, order, axis=-1)
+ # Rank order keeps the selection bias; the returned mass is the
+ # bias-free sigmoid score the block's mixing weights renormalize.
+ return inds, mx.take_along_axis(scores, inds, axis=-1)
return select
@@ -181,7 +185,8 @@ def select(x):
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
sc = mx.take_along_axis(gates, inds, axis=-1)
order = mx.argsort(-sc, axis=-1)
- return mx.take_along_axis(inds, order, axis=-1)
+ return (mx.take_along_axis(inds, order, axis=-1),
+ mx.take_along_axis(sc, order, axis=-1))
return select
@@ -238,7 +243,8 @@ def variants(self, probing: bool) -> tuple[str, ...]:
return (v,)
def predict(self, x, variant: str):
- """Lazy ranked expert ids (shape ``indices``-like) for ``dst_li``."""
+ """Lazy ``(ranked ids, matching mass scores)`` (shapes
+ ``indices``-like) for ``dst_li``."""
if variant == "ratio":
x = x * self._ratio.astype(x.dtype)
return self._router_fn(x)
@@ -360,9 +366,9 @@ def predictor(self):
def on_call(self, x, indices) -> dict:
"""Evaluate ``indices`` (the fence the caller needed anyway) plus
every live lookahead prediction in one sync. Returns a dict of
- predicted ranked ids per destination layer when prefetching, each
- trimmed to its rank gate prefix; empty when there is nothing to
- prestage."""
+ predicted ``(ranked ids, mass scores)`` per destination layer when
+ prefetching, both trimmed to the rank gate prefix; empty when
+ there is nothing to prestage."""
probing = self.probe is not None
ph = _LA_PHASE
t0 = time.perf_counter() if ph is not None else 0.0
@@ -384,7 +390,7 @@ def on_call(self, x, indices) -> dict:
t1 = time.perf_counter()
ph["build"] += t1 - t0
try:
- mx.eval(indices, *[a for _, _, a in lazy])
+ mx.eval(indices, *[t for _, _, pair in lazy for t in pair])
except Exception as exc:
# Joint eval: the failing predictor is unattributable, so
# disable every one this hook owns rather than loop forever.
@@ -400,8 +406,9 @@ def on_call(self, x, indices) -> dict:
if ph is not None:
ph["sync"] += time.perf_counter() - t1
by_pred: dict = {}
- for pred, variant, arr in lazy:
- by_pred.setdefault(pred, {})[variant] = np.array(arr)
+ for pred, variant, (ids_l, sc_l) in lazy:
+ by_pred.setdefault(pred, {})[variant] = (
+ np.array(ids_l), np.array(sc_l, dtype=np.float32))
actual_np = None
if self.probe is not None or self.gate is not None:
actual_np = np.array(indices)
@@ -410,10 +417,10 @@ def on_call(self, x, indices) -> dict:
if self.probe is not None:
for pred, pv in by_pred.items():
if pred.depth == 1:
- labeled = pv
+ labeled = {v: a for v, (a, _) in pv.items()}
else:
labeled = {
- f"{v}@d{pred.depth}": a for v, a in pv.items()
+ f"{v}@d{pred.depth}": a for v, (a, _) in pv.items()
}
self.probe.note(pred.dst_li, labeled)
self.probe.actual(self.li, actual_np)
@@ -423,9 +430,10 @@ def on_call(self, x, indices) -> dict:
# With the probe co-installed both variants exist;
# prefetch uses the configured one.
chosen = pred.variants(False)[0]
- ids = pv.get(chosen)
- if ids is None:
- ids = next(iter(pv.values()))
+ pair = pv.get(chosen)
+ if pair is None:
+ pair = next(iter(pv.values()))
+ ids, sc = pair
if self.gate is not None:
# Full width is noted (gated-out ranks keep being
# scored and can re-qualify); only the submission is
@@ -436,7 +444,8 @@ def on_call(self, x, indices) -> dict:
continue
if k < ids.shape[-1]:
ids = ids[..., :k]
- out[pred.dst_li] = ids
+ sc = sc[..., :k]
+ out[pred.dst_li] = (ids, sc)
return out
diff --git a/gmlx/moe_experts.py b/gmlx/moe_experts.py
index bfe9627..d5228e9 100644
--- a/gmlx/moe_experts.py
+++ b/gmlx/moe_experts.py
@@ -450,6 +450,31 @@ def install_moe_miss_shed(model, p: float) -> int:
return n
+def install_moe_prestage_keepers(model) -> int:
+ """Route lookahead prestage through the active miss-shed policy:
+ predictions the shed would drop on arrival are never read, and
+ predicted keepers - the reads the demand path would otherwise do
+ synchronously - stage demand-grade (full width, least-popular
+ eviction). Needs --moe-miss-shed on the same layers to define the
+ policy. Returns the number of layers hooked."""
+ n = 0
+ for m in _streamed_modules(model):
+ if getattr(m, "_kq_decode_feeder", None) is not None:
+ object.__setattr__(m, "_kq_prestage_keepers", True)
+ n += 1
+ if n:
+ loadlog.info(
+ f"[stream] MoE keeper prestage: lookahead reads on {n} offloaded "
+ "MoE layers target only experts the miss-shed policy would keep"
+ )
+ else:
+ print(
+ "[stream] MoE keeper prestage found no decode-feeder MoE layer "
+ "- no effect"
+ )
+ return n
+
+
def install_moe_layer_shed(model, p: float) -> int:
"""Skip each streamed MoE layer's routed path with probability ``p``
per decode token; the layer's shared expert still runs. Lossy.
diff --git a/gmlx/server.py b/gmlx/server.py
index eb55000..3415255 100644
--- a/gmlx/server.py
+++ b/gmlx/server.py
@@ -741,6 +741,13 @@ def _add_serve_args(ap: argparse.ArgumentParser) -> None:
"experts with probability P (0 < P < 1) per token; the "
"shared expert still runs on shed layers (config mode: "
"set `moe_layer_shed: P` per model).")
+ ap.add_argument("--moe-prestage", choices=("ranked", "keepers"),
+ default=None,
+ help="Lookahead prestage targeting for a streamed model; "
+ "default ranked. keepers filters predictions through "
+ "the active miss-shed policy and stages predicted "
+ "keepers demand-grade; needs --moe-miss-shed (config "
+ "mode: set `moe_prestage: keepers` per model).")
ap.add_argument("--prefill-feeder", action=argparse.BooleanOptionalAction,
default=None,
help="Faster prompt processing for streaming models "
@@ -907,6 +914,8 @@ def _abs(p):
out += ["--moe-miss-shed", str(a.moe_miss_shed)]
if getattr(a, "moe_layer_shed", None) is not None:
out += ["--moe-layer-shed", str(a.moe_layer_shed)]
+ if getattr(a, "moe_prestage", None) is not None:
+ out += ["--moe-prestage", str(a.moe_prestage)]
if a.budget_gb is not None:
out += ["--budget-gb", str(a.budget_gb)]
if a.max_models is not None:
@@ -1377,6 +1386,7 @@ def _single_model_cfg(a) -> ServerCfg:
moe_expert_mass=getattr(a, "moe_expert_mass", None),
moe_miss_shed=getattr(a, "moe_miss_shed", None),
moe_layer_shed=getattr(a, "moe_layer_shed", None),
+ moe_prestage=getattr(a, "moe_prestage", None),
prefill_feeder=getattr(a, "prefill_feeder", None),
decode_feeder=getattr(a, "decode_feeder", None),
pin=True, # the single model is always pinned
diff --git a/gmlx/server_bridge_vlm.py b/gmlx/server_bridge_vlm.py
index e97bf74..2b2c27d 100644
--- a/gmlx/server_bridge_vlm.py
+++ b/gmlx/server_bridge_vlm.py
@@ -494,6 +494,7 @@ def load_serveable_model(
moe_expert_mass: float | None = None,
moe_miss_shed: float | None = None,
moe_layer_shed: float | None = None,
+ moe_prestage: str | None = None,
feeder_prefill: bool | None = None,
feeder_decode: bool | None = None,
) -> tuple[object, object, object]:
@@ -533,7 +534,9 @@ def load_serveable_model(
``moe_miss_shed: P`` / ``moe_layer_shed: P``) install their filters/hooks
over the streamed layers after the placement. They ride on ``stream`` -
without a placement each is announced as ignored (there are no streamed
- experts to filter).
+ experts to filter). ``moe_prestage: keepers`` retargets the lookahead
+ prestage through the miss-shed policy and additionally needs
+ ``moe_miss_shed`` (announced as ignored without it).
"""
def _reject_unwired(base_kind: str) -> None:
# Raising beats silently dropping the option on bases that don't
@@ -551,14 +554,15 @@ def _reject_unwired(base_kind: str) -> None:
for key, val in (("moe_experts", moe_experts),
("moe_expert_mass", moe_expert_mass),
("moe_miss_shed", moe_miss_shed),
- ("moe_layer_shed", moe_layer_shed)):
+ ("moe_layer_shed", moe_layer_shed),
+ ("moe_prestage", moe_prestage)):
if val is not None:
print(
f"[stream] {key} ignored: needs stream: experts|cpu "
"(it only applies to streamed MoE layers)"
)
moe_experts = moe_expert_mass = None
- moe_miss_shed = moe_layer_shed = None
+ moe_miss_shed = moe_layer_shed = moe_prestage = None
if mmproj_path is not None and speculative:
# VLM x MTP: text-only requests speculate; image/audio requests prefill media
# into the KV and decode normally (verify is token-only over that cache).
@@ -616,6 +620,13 @@ def _reject_unwired(base_kind: str) -> None:
if moe_miss_shed is not None:
from .moe_experts import install_moe_miss_shed
install_moe_miss_shed(raw_model, moe_miss_shed)
+ if moe_prestage == "keepers":
+ if moe_miss_shed is None:
+ print("[stream] moe_prestage: keepers ignored: it needs "
+ "moe_miss_shed")
+ else:
+ from .moe_experts import install_moe_prestage_keepers
+ install_moe_prestage_keepers(raw_model)
if moe_layer_shed is not None:
from .moe_experts import install_moe_layer_shed
install_moe_layer_shed(raw_model, moe_layer_shed)
@@ -767,14 +778,15 @@ def load_model_resources(model_path, adapter_path=None):
# The feeder overrides (config `prefill_feeder:`/`decode_feeder:` /
# the paired serve flags) and the lossy MoE levers (config
# `moe_experts:`/`moe_expert_mass:`/`moe_miss_shed:`/
- # `moe_layer_shed:` / the paired serve flags) ride along; None
- # keeps the loader default / trained fan-out.
+ # `moe_layer_shed:`/`moe_prestage:` / the paired serve flags) ride
+ # along; None keeps the loader default / trained fan-out.
stream = getattr(spec, "stream", None)
feeders = dict(
moe_experts=getattr(spec, "moe_experts", None),
moe_expert_mass=getattr(spec, "moe_expert_mass", None),
moe_miss_shed=getattr(spec, "moe_miss_shed", None),
moe_layer_shed=getattr(spec, "moe_layer_shed", None),
+ moe_prestage=getattr(spec, "moe_prestage", None),
feeder_prefill=getattr(spec, "prefill_feeder", None),
feeder_decode=getattr(spec, "decode_feeder", None),
)
diff --git a/tests/test_config.py b/tests/test_config.py
index 310a4a8..5ee876c 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -236,31 +236,34 @@ def test_moe_lossy_levers_parsed_and_resolved():
doc = _doc()
doc["models"]["m-moe"] = {"path": "/abs/big.gguf", "stream": "experts",
"moe_experts": 6, "moe_miss_shed": 0.95,
- "moe_layer_shed": "0.1"}
+ "moe_layer_shed": "0.1",
+ "moe_prestage": "Keepers"}
cfg = build_config(doc)
rm = resolve_model("m-moe", cfg)
assert rm.moe_experts == 6
assert rm.moe_miss_shed == 0.95
assert rm.moe_layer_shed == 0.1
+ assert rm.moe_prestage == "keepers" # case-normalized
bare = resolve_model("m-named", cfg)
assert (bare.moe_experts is None and bare.moe_miss_shed is None
- and bare.moe_layer_shed is None)
+ and bare.moe_layer_shed is None and bare.moe_prestage is None)
# Each lever is load-affecting: same GGUF, different value => distinct
# resident entries.
common = dict(path="/p", sampling={}, load={}, cache={}, system=None,
speculative=False, mmproj=None, draft_gguf=None, pin=False,
ttl_s=None, stream="experts")
base_sig = cfgmod.ResolvedModel(id="x", **common).load_signature()
- for key in ("moe_experts", "moe_expert_mass", "moe_miss_shed",
- "moe_layer_shed"):
+ for key, val in (("moe_experts", 2), ("moe_expert_mass", 0.5),
+ ("moe_miss_shed", 0.5), ("moe_layer_shed", 0.5),
+ ("moe_prestage", "keepers")):
sig = cfgmod.ResolvedModel(
- id="x", **{key: 2 if key == "moe_experts" else 0.5},
- **common).load_signature()
+ id="x", **{key: val}, **common).load_signature()
assert sig != base_sig, key
for key, bad in (("moe_experts", 0), ("moe_experts", "many"),
("moe_miss_shed", 1.5), ("moe_miss_shed", 0),
- ("moe_layer_shed", 1.0), ("moe_layer_shed", 0)):
+ ("moe_layer_shed", 1.0), ("moe_layer_shed", 0),
+ ("moe_prestage", "eager")):
doc = _doc()
doc["models"]["m-moe"] = {"path": "/abs/big.gguf",
"stream": "experts", key: bad}
diff --git a/tests/test_decode_feeder.py b/tests/test_decode_feeder.py
index 7326f31..6775242 100644
--- a/tests/test_decode_feeder.py
+++ b/tests/test_decode_feeder.py
@@ -1091,6 +1091,42 @@ def bare(verbose):
assert "wedged" in out
+# Keeper-mode prestage (--moe-prestage keepers): shed-aware targeting
+def test_prestage_keeper_mode_filters_and_evicts(monkeypatch, tmp_path):
+ """Keeper mode never reads a prediction the shed policy would drop on
+ arrival, stages the keepers demand-grade (a colder resident is
+ evicted although the prediction has no popularity), and books no
+ shed stats for the filtering."""
+ feeder, _ = _make_feeder(monkeypatch, tmp_path)
+ feeder.stage(0, np.array([[0, 1]])) # arena full: e0, e1
+ feeder.stage(0, np.array([[0]])) # counts: e0=2 > e1=1
+ # e0 resident; e2 (0.3) survives the 0.3 budget, e3 (0.1) sheds.
+ ids = np.array([[0, 2, 3]])
+ sc = np.array([0.6, 0.3, 0.1], dtype=np.float32)
+ feeder.prestage(0, ids, keep_mass=0.7, pred_scores=sc)
+ assert feeder._la_submitted == 1
+ assert 2 in feeder._pending[0] and 3 not in feeder._pending[0]
+ assert feeder._slot_of[0][1] == -1 # cold resident evicted outright
+ assert feeder._slot_of[0][0] >= 0 # predicted resident shielded
+ assert feeder._shed_n == 0 and feeder._shed_tokens == 0
+ _wait_published(feeder, 0)
+ assert feeder._slot_of[0][2] >= 0
+
+
+def test_prestage_keeper_mode_full_width_when_nothing_sheds(
+ monkeypatch, tmp_path):
+ """A budget that sheds nothing makes every predicted miss a keeper:
+ keeper mode skips the speculative rank cap, ranked mode keeps it."""
+ monkeypatch.setenv("GMLX_DECODE_LOOKAHEAD_K", "1")
+ feeder, _ = _make_feeder(monkeypatch, tmp_path)
+ ids = np.array([[1, 3]])
+ sc = np.array([0.5, 0.5], dtype=np.float32)
+ feeder.prestage(0, ids, keep_mass=1.0, pred_scores=sc)
+ assert feeder._la_submitted == 2 # both, despite the rank cap of 1
+ feeder.prestage(1, ids)
+ assert feeder._la_submitted == 3 # ranked mode: capped to rank 0
+
+
# Lossy miss-shed (--moe-miss-shed): residency-aware expert drop
def test_shed_misses_drops_cold_lowest_first(monkeypatch, tmp_path):
"""Only demand-miss experts shed, lowest score first, capped at
diff --git a/tests/test_generate_over.py b/tests/test_generate_over.py
index 5063aae..5729745 100644
--- a/tests/test_generate_over.py
+++ b/tests/test_generate_over.py
@@ -64,6 +64,34 @@ def _call(**kw):
return _generate_over(object(), _Tok(), "PROMPT", **base)
+def test_verbose_phase1_uses_styled_emitter(monkeypatch, capsys):
+ """Verbose phase 1 streams through the main path's emitter (thinking
+ rendering included) and closes it before the seam marker; phase 2
+ still prints raw."""
+ import gmlx.generation as g
+
+ got = {"chunks": [], "closed": 0, "reasoning": None}
+
+ def fake_emitter(prompt, tokenizer, reasoning):
+ got["reasoning"] = reasoning
+
+ def close():
+ got["closed"] += 1
+
+ return got["chunks"].append, close
+
+ monkeypatch.setattr(g, "_verbose_emitter", fake_emitter)
+ _patch_stream(
+ monkeypatch,
+ phase1=[(11, "think"), (12, "answer"), (2, "<|user|>")],
+ phase2=[(20, "over")],
+ )
+ _call(window=1, verbose=True, reasoning="hide")
+ assert got["chunks"] == ["think", "answer"] # phase 2 never routed here
+ assert got["closed"] == 1 and got["reasoning"] == "hide"
+ assert "over" in capsys.readouterr().out # phase 2 raw print intact
+
+
def test_free_mode_splits_at_seam_and_forces_window(monkeypatch, tmp_path):
_patch_stream(
monkeypatch,
diff --git a/tests/test_lookahead.py b/tests/test_lookahead.py
index c6bbe71..82bd9b5 100644
--- a/tests/test_lookahead.py
+++ b/tests/test_lookahead.py
@@ -115,8 +115,9 @@ def __call__(self, x):
mx.array([[0.1, 0.7, 0.2]]),
)
- ids = np.array(_gate_module_select(_Stub())(mx.zeros((1, DIM))))
- assert ids.tolist() == [[9, 2, 4]]
+ ids, sc = _gate_module_select(_Stub())(mx.zeros((1, DIM)))
+ assert np.array(ids).tolist() == [[9, 2, 4]]
+ assert np.allclose(np.array(sc), [[0.7, 0.2, 0.1]])
def test_sigmoid_bias_select_matches_stock_selection():
@@ -130,12 +131,16 @@ def __init__(self):
mx.random.seed(3)
blk = _Block()
x = mx.random.normal((1, 1, DIM))
- ids = np.array(_sigmoid_bias_select(blk)(x)).reshape(-1)
+ ids_mx, sc_mx = _sigmoid_bias_select(blk)(x)
+ ids = np.array(ids_mx).reshape(-1)
# Reference: the stock forward's selection seam.
choice = mx.sigmoid(blk.gate(x.astype(mx.float32)))
- choice = np.array(choice + blk.e_score_correction_bias).reshape(-1)
- ref = np.argsort(-choice)[:K]
+ biased = np.array(choice + blk.e_score_correction_bias).reshape(-1)
+ ref = np.argsort(-biased)[:K]
assert ids.tolist() == ref.tolist()
+ # Mass scores are the bias-free sigmoid the block's mix renormalizes.
+ raw = np.array(choice).reshape(-1)
+ assert np.allclose(np.array(sc_mx).reshape(-1), raw[ids], atol=1e-6)
def test_router_fn_for_unknown_block_is_none():
@@ -464,6 +469,56 @@ def swapped(self, li):
assert sorted(li for li, _ in df.calls) == [1, 2]
+def test_wrapper_keeper_prestage_passes_policy(monkeypatch):
+ """--moe-prestage keepers: the wrapper hands prestage the active shed
+ share and the prediction's mass scores; with the switch off, or with
+ no miss-shed installed, the plain ranked call is unchanged."""
+ from contextlib import contextmanager
+
+ model = _streaming_model(monkeypatch)
+ install_lookahead(model, model.layers, probe=False, prefetch=True)
+ glu0 = model.layers[0].mlp.switch_mlp
+
+ class _FakeDF:
+ def ensure_wired(self):
+ pass
+
+ calls = []
+
+ def covers(self, li):
+ return True
+
+ def stage(self, li, ids):
+ return ids.astype(np.uint32)
+
+ def prestage(self, li, pred, **kw):
+ self.calls.append((li, pred.copy(), kw))
+
+ def wedged_at(self, li):
+ return False
+
+ @contextmanager
+ def swapped(self, li):
+ yield
+
+ df = _FakeDF()
+ object.__setattr__(glu0, "_kq_decode_feeder", df)
+ object.__setattr__(glu0, "_kq_li", 0)
+ x = mx.random.normal((1, 1, DIM))
+ mx.eval(model.layers[0].mlp(x))
+ assert df.calls and df.calls[-1][2] == {} # ranked: no policy kwargs
+
+ object.__setattr__(glu0, "_kq_prestage_keepers", True)
+ mx.eval(model.layers[0].mlp(x))
+ assert df.calls[-1][2] == {} # keepers without miss-shed: unchanged
+
+ object.__setattr__(glu0, "_kq_miss_shed", 0.8)
+ mx.eval(model.layers[0].mlp(x))
+ li, pred, kw = df.calls[-1]
+ assert li == 1 and kw["keep_mass"] == 0.8
+ assert kw["pred_scores"].shape == pred.shape
+
+
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-q"]))
@@ -486,7 +541,8 @@ def __init__(self):
from gmlx.lookahead import _SIGMOID_BIAS_BLOCKS
assert "KimiK3MoE" in _SIGMOID_BIAS_BLOCKS
x = mx.random.normal((1, 1, DIM))
- ids = np.array(_sigmoid_bias_select(blk)(x)).reshape(-1)
+ ids_mx, _ = _sigmoid_bias_select(blk)(x)
+ ids = np.array(ids_mx).reshape(-1)
choice = mx.sigmoid(blk.gate(x.astype(mx.float32)))
choice = np.array(choice + blk.e_score_correction_bias).reshape(-1)
ref = np.argsort(-choice)[:K]
diff --git a/tests/test_offload.py b/tests/test_offload.py
index f7e95b0..644a4d1 100644
--- a/tests/test_offload.py
+++ b/tests/test_offload.py
@@ -518,3 +518,78 @@ def swapped(self, li):
# out at single-token stage calls that fit under the slot cap.
assert stub.stage_sizes[0] == T * 2
assert min(stub.stage_sizes) == 2
+
+
+def test_arena_split_leaves_never_miss_shed():
+ """Miss-shed is decode-only: single-token leaves inside an arena token
+ split must not shed. A shedding leaf returns a mixed (rank-3) output
+ while a clean leaf returns per-expert rank-4, and the reassembly
+ concatenate crashed on the rank mismatch (kimi-k3 chat prefill with
+ --moe-miss-shed). The split output must match the unsplit reference."""
+ import contextlib
+
+ import numpy as np
+
+ mx.random.seed(5)
+ glu = SwitchGLU(16, 32, 8)
+ mx.eval(glu.parameters())
+ T = 2
+ x = mx.random.normal((1, T, 16))
+ inds = mx.array(
+ (np.arange(T * 2).reshape(1, T, 2) % 8).astype(np.uint32))
+ weights = mx.softmax(mx.random.normal((1, T, 2)), axis=-1)
+ ref = (glu(x, inds) * weights[..., None]).sum(axis=-2)
+ mx.eval(ref)
+
+ class _ShedStub:
+ def __init__(self, cap):
+ self.cap = cap
+ self.shed_calls = 0
+
+ def covers(self, li):
+ return True
+
+ def can_stage_smaller(self, li):
+ return True
+
+ def wedged_at(self, li):
+ return False
+
+ def ensure_wired(self):
+ pass
+
+ def shed_misses(self, li, ids, sc, p):
+ # First leaf: everything resident, nothing shed. Second leaf:
+ # shed all but the first routed expert. The asymmetry is what
+ # produced the mixed-rank concatenate.
+ self.shed_calls += 1
+ if self.shed_calls == 1:
+ return None
+ return np.array([0])
+
+ def stage(self, li, ids):
+ if len(np.unique(ids.reshape(-1))) > self.cap:
+ return None
+ return ids # identity slots: arena views == real weights
+
+ @contextlib.contextmanager
+ def swapped(self, li):
+ yield
+
+ model = _holder_model(glu)
+ n, _ = install_expert_streaming(model)
+ assert n == 1
+ stub = _ShedStub(cap=2)
+ object.__setattr__(glu, "_kq_decode_feeder", stub)
+ object.__setattr__(glu, "_kq_cpu_only", True)
+ object.__setattr__(glu, "_kq_li", 0)
+ object.__setattr__(glu, "_kq_miss_shed", 0.8)
+ # Scores-sink call shape: the block hands routing weights as the third
+ # argument (kimi-k3 latent MoE does this whenever the sink is advertised).
+ out = glu(x, inds, weights)
+ mx.eval(out)
+ # Leaves stay on the per-expert contract; mix like the block would.
+ assert out.ndim == x.ndim + 1
+ mixed = (out * weights[..., None]).sum(axis=-2)
+ assert mx.allclose(ref, mixed, atol=1e-6, rtol=1e-6)
+ assert stub.shed_calls == 0 # decode-only: split leaves never shed
diff --git a/tests/test_serving_mtp.py b/tests/test_serving_mtp.py
index d00924c..a41ec2a 100644
--- a/tests/test_serving_mtp.py
+++ b/tests/test_serving_mtp.py
@@ -356,7 +356,7 @@ def spy(model_path, *, mmproj_path=None, hf_source=None,
speculative=False, draft_gguf_path=None, chat_template=None,
adapter_gguf=None, stream=None, moe_experts=None,
moe_expert_mass=None, moe_miss_shed=None, moe_layer_shed=None,
- feeder_prefill=None, feeder_decode=None):
+ moe_prestage=None, feeder_prefill=None, feeder_decode=None):
serveable_calls.append(
(model_path, mmproj_path, hf_source, speculative, draft_gguf_path,
chat_template, adapter_gguf)
diff --git a/tests/test_serving_vlm.py b/tests/test_serving_vlm.py
index b4ebd6a..e3a137b 100644
--- a/tests/test_serving_vlm.py
+++ b/tests/test_serving_vlm.py
@@ -148,7 +148,7 @@ def spy(model_path, *, mmproj_path=None, hf_source=None,
speculative=False, draft_gguf_path=None, chat_template=None,
adapter_gguf=None, stream=None, moe_experts=None,
moe_expert_mass=None, moe_miss_shed=None, moe_layer_shed=None,
- feeder_prefill=None, feeder_decode=None):
+ moe_prestage=None, feeder_prefill=None, feeder_decode=None):
serveable_calls.append((model_path, mmproj_path, hf_source, chat_template))
return ("SERVEABLE", model_path, mmproj_path)