From a7984da46f00aadfc9410cf971642e29d581f8f4 Mon Sep 17 00:00:00 2001 From: Demetrios Agourakis Date: Fri, 11 Sep 2026 08:31:14 -0300 Subject: [PATCH 1/2] fix(streaming): stop layer discovery on IndexError too install_streaming_experts(num_layers=None) counts layers by probing spec.block_of(model, n) with increasing n until it fails, but it only catches AttributeError. Layers live in a list (MoESpec.block_of indexes "layers are plain lists in most families"), so running off the end raises IndexError and the call crashes instead of returning. The shipped engines always pass num_layers, which is why this never showed. Adds test_install_discovers_layer_count: a three-layer list-based model (one MoE layer, two dense) installed with num_layers=None. It raised IndexError before this change; now it returns [twin, None, None] and the twin is swapped into the MoE block. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0146Hf9MXz9hHfRnq3MNXXWX --- src/edge0/streaming/install.py | 2 +- tests/test_streaming_math.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/edge0/streaming/install.py b/src/edge0/streaming/install.py index 37b0e7e..acc3585 100644 --- a/src/edge0/streaming/install.py +++ b/src/edge0/streaming/install.py @@ -41,7 +41,7 @@ def install_streaming_experts( while True: try: spec.block_of(model, n) - except AttributeError: + except (AttributeError, IndexError): # past the last layer break n += 1 if n == 0: diff --git a/tests/test_streaming_math.py b/tests/test_streaming_math.py index e77154e..3fb0b6f 100644 --- a/tests/test_streaming_math.py +++ b/tests/test_streaming_math.py @@ -299,3 +299,33 @@ def test_double_buffered_swap(layer): ref2 = lay(x, mx.array([second], dtype=mx.int32)) assert mx.allclose(out1, ref1).item() assert mx.allclose(out2, ref2).item() + + +def test_install_discovers_layer_count(tmp_path): + """install_streaming_experts(num_layers=None) probes block_path with + increasing layer indices until it stops resolving. Layers live in a + list, so running off the end raises IndexError, not AttributeError.""" + from types import SimpleNamespace + + from edge0.streaming.install import install_streaming_experts + + path = tmp_path / "w.safetensors" + _write_shard(path, fuse_gu=False) + spec = MoESpec( + num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, + quant=QuantSpec(bits=4, group_size=64), + layout=WeightLayout.SEPARATE, + key_template="layers.{layer}.mlp.switch_mlp", + block_path="layers.{layer}.mlp", + ) + moe = SimpleNamespace(switch_mlp=object()) + model = SimpleNamespace(layers=[SimpleNamespace(mlp=moe), + SimpleNamespace(mlp=SimpleNamespace()), + SimpleNamespace(mlp=SimpleNamespace())]) + twins = install_streaming_experts( + model, [SafetensorsMmap(str(path))], spec, options=_options()) + assert len(twins) == 3 + assert isinstance(twins[0], StreamingSwitchGLU) + assert twins[1] is None and twins[2] is None # dense layers + assert moe.switch_mlp is twins[0] + twins[0].close() From d2178440fd94133bfd0c04321265ee82b9bac900 Mon Sep 17 00:00:00 2001 From: Demetrios Agourakis Date: Tue, 22 Sep 2026 17:11:29 -0300 Subject: [PATCH 2/2] fix(moe): isolate layer-index errors from other numeric segments in discovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two coverage nits from the PR review of the previous commit: - The AttributeError half of the widened except was untested: added test_layer_exists_stops_on_attribute_error, using an attribute-based layer container (no list, no __getitem__) to reach it independently of the IndexError case. - Catching (AttributeError, IndexError) around the whole block_of call could mask an IndexError raised by a different, non-layer-index numeric segment further down block_path (e.g. a fixed expert-slot index) — a real bug, not end-of-list. No shipped spec hits this, but it would have silently under-counted layers instead of surfacing the break. Fixes it by adding MoESpec.layer_exists(model, layer), which resolves block_path segment-by-segment and only treats an AttributeError/ IndexError as "past the last layer" when it comes from the segment templated by {layer} itself; any other segment's error propagates. install_streaming_experts's discovery loop now calls this instead of wrapping spec.block_of directly. test_layer_exists_reraises_unrelated_index_error covers the propagation case. Co-Authored-By: Claude Sonnet 5 --- src/edge0/moe/spec.py | 29 ++++++++++++++++++++++++++ src/edge0/streaming/install.py | 6 +----- tests/test_streaming_math.py | 38 ++++++++++++++++++++++++++++++++++ 3 files changed, 68 insertions(+), 5 deletions(-) diff --git a/src/edge0/moe/spec.py b/src/edge0/moe/spec.py index 176c43c..6d0d92c 100644 --- a/src/edge0/moe/spec.py +++ b/src/edge0/moe/spec.py @@ -110,6 +110,35 @@ def block_of(self, model, layer: int): obj = getattr(obj, part) return obj + def layer_exists(self, model, layer: int) -> bool: + """True if ``layer`` resolves at block_path's layer-index segment. + + Used by layer-count discovery (``install_streaming_experts`` with + ``num_layers=None``) to find where the layer list ends. Only an + ``AttributeError``/``IndexError`` raised while resolving the + segment templated by ``{layer}`` itself means "past the last + layer"; the same errors raised by a *different* segment further + down ``block_path`` (e.g. a fixed expert-slot index) indicate a + bug in the spec or model and are re-raised rather than read as + end-of-list. + """ + template_parts = self.block_path.split(".") + layer_pos = next( + i for i, p in enumerate(template_parts) if "{layer}" in p) + obj = model + for i, raw_part in enumerate(template_parts): + part = raw_part.format(layer=layer) + try: + if part.isdigit(): + obj = obj[int(part)] + else: + obj = getattr(obj, part) + except (AttributeError, IndexError): + if i == layer_pos: + return False + raise + return True + def layer_of(self, model, layer: int): """Resolve the decoder layer object (the block's owner). diff --git a/src/edge0/streaming/install.py b/src/edge0/streaming/install.py index acc3585..c0c84de 100644 --- a/src/edge0/streaming/install.py +++ b/src/edge0/streaming/install.py @@ -38,11 +38,7 @@ def install_streaming_experts( """ if num_layers is None: n = 0 - while True: - try: - spec.block_of(model, n) - except (AttributeError, IndexError): # past the last layer - break + while spec.layer_exists(model, n): n += 1 if n == 0: raise ValueError( diff --git a/tests/test_streaming_math.py b/tests/test_streaming_math.py index 3fb0b6f..0d2bcbd 100644 --- a/tests/test_streaming_math.py +++ b/tests/test_streaming_math.py @@ -329,3 +329,41 @@ def test_install_discovers_layer_count(tmp_path): assert twins[1] is None and twins[2] is None # dense layers assert moe.switch_mlp is twins[0] twins[0].close() + + +def test_layer_exists_stops_on_attribute_error(): + """The AttributeError half of ``layer_exists``'s except clause is + reachable independently of IndexError: a family whose layer container + is attribute-based (no list, no ``__getitem__``) runs off the end via + a plain missing attribute, not an out-of-range index.""" + from types import SimpleNamespace + + spec = MoESpec( + num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, + key_template="layer_{layer}.mlp.switch_mlp", + block_path="layer_{layer}", + ) + model = SimpleNamespace(layer_0=object(), layer_1=object()) + assert spec.layer_exists(model, 0) is True + assert spec.layer_exists(model, 1) is True + assert spec.layer_exists(model, 2) is False # no `layer_2` attribute + + +def test_layer_exists_reraises_unrelated_index_error(): + """An IndexError from a segment *other* than the ``{layer}`` slot is a + real bug (a malformed block_path or a broken model), not end-of-list, + and must propagate instead of being read as "past the last layer" -- + otherwise install_streaming_experts(num_layers=None) would silently + under-count layers instead of surfacing the break.""" + from types import SimpleNamespace + + spec = MoESpec( + num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, + key_template="layers.{layer}.experts.9.switch_mlp", + block_path="layers.{layer}.experts.9", + ) + # `layer=0` is in range for `layers`, but the fixed trailing index `9` + # is out of range for `experts` -- unrelated to layer-count discovery. + model = SimpleNamespace(layers=[SimpleNamespace(experts=[object()])]) + with pytest.raises(IndexError): + spec.layer_exists(model, 0)