ci: pyright seam-drift guard on static upstream importers

.github/workflows/test.yml

@@ -89,6 +89,17 @@ jobs:
         # renaming a model_type - the weekly cron catches those) must not leave
         # the committed doc lying about what loads.
         run: python scripts/check-coverage.py --check --strict
+      - name: Seam drift check (pyright)
+        # Pyright infers upstream signatures from the installed mlx-vlm /
+        # mlx-lm source, so it needs this job's environment, not the
+        # deps-free lint job. One interpreter is enough: the checked
+        # surface does not vary by Python version. On the weekly cron a
+        # new upstream release can fail this step with no repo change;
+        # that is the drift signal, handled per docs/internals/upstream-upgrades.md.
+        if: matrix.python-version == '3.12'
+        run: |
+          pip install "pyright[nodejs]==1.1.414"
+          pyright
       - name: Run CPU logic tests
         env:
           # Hosted runners have no dependable Metal device; the conftest flips

CONTRIBUTING.md

@@ -48,8 +48,20 @@ tests, in
 ruff check .
 python scripts/check-docs.py   # docs style and link check, also a CI step
 pre-commit install             # optional, runs ruff on each commit
+pip install "pyright[nodejs]==1.1.414" && pyright   # seam drift check, needs the deps installed
 ```
 
+pyright is a guard against upstream symbol and signature drift, not a type
+checker for the tree. It covers the files in `[tool.pyright].include`, all
+of them at zero errors; the other files that import mlx_vlm or mlx_lm
+statically are listed in `[tool.gmlx.pyright].backlog` and join the gate
+one file at a time as each reaches zero. `tests/test_pyright_scope.py`
+holds both lists to the tree, so a new static import lands in one list or
+the other. A `# pyright: ignore[rule]` is allowed only where the false
+positive comes from upstream typing (a stub that omits the attribute, a
+scalar union, a field added to an upstream dataclass), with a comment
+saying which; keep the count under twenty repo-wide.
+
 ## Things to know before you patch
 
 - The serving stack is stock mlx-vlm with late-bound patches over its

docs/internals/upstream-upgrades.md

@@ -23,6 +23,7 @@ environment inside those bounds:
 | declared versions | `pyproject.toml` | the exact mlx-vlm pin and the mlx-lm and mlx-kquant floors described above |
 | seam contract | `tests/upstream/test_upstream_seams.py` | every patched symbol is pinned to a source fingerprint. Drift fails CI naming the seam |
 | runtime gate | `check_upstream_versions`, at CLI entry | mlx, mlx-lm or mlx-vlm below its floor refuses to run with an upgrade message, and newer than the qualified set warns once. `gmlx doctor` is exempt |
+| static-import check | `pyright`, in the macOS CI job | the files in `[tool.pyright].include` are checked against the installed upstream source; a moved symbol or changed signature there fails CI |
 
 ## Watching upstream releases
 

gmlx/cache/apc_manager.py

@@ -81,7 +81,7 @@ class GmlxAPCManager(_apc.APCManager):
     override defers to the stock store instead.
     """
 
-    def autosize(self, model, budget_fraction: float = None) -> None:
+    def autosize(self, model, budget_fraction: float | None = None) -> None:
         """Size the caches to the box post-load. Pool: raise the block
         cap to a working-budget share when APC_NUM_BLOCKS is unset
         (blocks allocate lazily, so the cap costs nothing until
@@ -166,7 +166,8 @@ def entry_bytes(e):
                 total -= sizes.pop(k, 0)
             # Mirror of pool_bytes for the exact tier: lets harnesses
             # separate budgeted, evictable retention from real residue.
-            self.stats.exact_bytes = int(total)
+            # gmlx side counter on the stock stats dataclass.
+            self.stats.exact_bytes = int(total)  # pyright: ignore[reportAttributeAccessIssue]
 
     def stats_snapshot(self) -> dict:
         """Stock snapshot plus the gmlx ckpt-tier side counters (pure
@@ -524,7 +525,7 @@ def _flush_pending(force=False):
                         "APC disk save scheduling failed: %s", e)
             self.stats.pool_used = sum(
                 1 for x in self.pool if x.block_hash is not None)
-            self.stats.pool_bytes = int(
+            self.stats.pool_bytes = int(  # pyright: ignore[reportAttributeAccessIssue]
                 self.stats.pool_used * self.block_size
                 * getattr(self, "_pool_per_token_bytes", 0))
             return new_blocks

gmlx/load/modules.py

@@ -1594,6 +1594,9 @@ class LoRAKQuantLinear(nn.Module):
     with nothing published the static ``scale`` applies to every row.
     """
 
+    _kq_tables: dict
+    _kq_extra: list
+
     def __init__(self, base: nn.Module, a: mx.array, b: mx.array, scale: float,
                  slot: int = 0):
         super().__init__()

gmlx/serve/bridge_vlm.py

@@ -48,7 +48,9 @@
 import logging
 import os
 import sys
+from collections.abc import Callable
 from contextvars import ContextVar
+from typing import Any
 
 from mlx_vlm import tokenizer_utils as _mlxvlm_tok
 from gmlx.models.vlm_text_only import Model as TextOnlyModel
@@ -233,11 +235,11 @@ def __getattr__(self, name):
         return getattr(self._wrapper, name)
 
 
-def _as_dict(config) -> dict:
+def _as_dict(config: Any) -> dict:
     if isinstance(config, dict):
         return config
     for attr in ("to_dict", "__dict__"):
-        value = getattr(config, attr, None)
+        value: Callable[..., Any] | dict | None = getattr(config, attr, None)
         if callable(value):
             return dict(value())
         if isinstance(value, dict):
@@ -372,9 +374,10 @@ def _load_serveable_vlm(
     """
     from gmlx.load.vlm import load_vlm_model
 
-    model, _config_dict, processor = load_vlm_model(
+    loaded = load_vlm_model(
         gguf_path, mmproj_path, hf_source=hf_source, verbose=False
     )
+    model, processor = loaded[0], loaded[2]
     # Return the model's own dataclass config (what stock load_model_resources
     # returns as the 3rd element), not the synthesized dict.
     return model, processor, model.config
@@ -387,7 +390,7 @@ def _make_text_processor(tokenizer) -> "_GgufServerProcessor":
     detokenizer. (The VLM path gets an engine-ready processor from the loader and
     does not use this.)"""
     backend = getattr(tokenizer, "_tokenizer", tokenizer)
-    eos = getattr(tokenizer, "eos_token_ids", None) or getattr(
+    eos: Any = getattr(tokenizer, "eos_token_ids", None) or getattr(
         tokenizer, "eos_token_id", None
     )
     # StoppingCriteria.add_eos_token_ids() mutates this list in place, so it must
@@ -662,7 +665,7 @@ def _reject_unwired(base_kind: str, *, streamable: bool = False,
                 )
         moe_experts = moe_expert_mass = None
         moe_miss_shed = moe_layer_shed = moe_prestage = None
-    _levers = dict(
+    _levers: dict[str, Any] = dict(
         moe_experts=moe_experts, moe_expert_mass=moe_expert_mass,
         moe_miss_shed=moe_miss_shed, moe_prestage=moe_prestage,
         moe_layer_shed=moe_layer_shed)
@@ -761,7 +764,7 @@ def load_drafter(path_or_repo, kind=None, **kwargs):
         _apply_draft_block_size_override(result)
         return result
 
-    drafters.load_drafter = load_drafter
+    setattr(drafters, "load_drafter", load_drafter)
     setattr(drafters, _DRAFTER_PATCH_FLAG, True)
 
 
@@ -825,7 +828,7 @@ class _DrafterSourceFilter(logging.Filter):
     def filter(self, record):
         if str(record.msg).startswith("Loading speculative drafter"):
             args = record.args or ()
-            path = args[-1] if args else None
+            path = args[-1] if isinstance(args, tuple) and args else None
             if isinstance(path, str) and os.path.abspath(path) in _MTP_DRAFTER_STASH:
                 return False
         return True
@@ -951,7 +954,7 @@ def _bridge_load(model_path, adapter_path=None):
             # `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(
+            feeders: dict[str, Any] = 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),
@@ -1053,7 +1056,7 @@ def load_model_resources(model_path, adapter_path=None):
             _raise_if_first_party_import(e)
             raise
 
-    generation.load_model_resources = load_model_resources
+    setattr(generation, "load_model_resources", load_model_resources)
     setattr(generation, _BRIDGE_FLAG, True)
     # generation.py logs on the parent "mlx_vlm.server" logger.
     engine_log = logging.getLogger("mlx_vlm.server")
@@ -1077,7 +1080,7 @@ def load_model_resources(model_path, adapter_path=None):
 # finds it, and exposes the resolved spec for *this* request through a ContextVar
 # (mirroring residency's ``_active_entry`` discipline).
 
-_RESOLVED_MODELS: dict[str, "object"] = {}     # id -> ResolvedModel
+_RESOLVED_MODELS: dict[str, Any] = {}          # id -> ResolvedModel
 _PATH_TO_IDS: dict[str, list[str]] = {}        # abspath -> [id, ...]
 _SERVER_CFG = None                             # the live ServerCfg (for re-resolve)
 # The ResolvedModel for the request in flight - set at the residency seam, read at

gmlx/spec/dflash_drafter.py

@@ -47,7 +47,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -501,6 +501,9 @@ def make_cache(self, left_padding: Optional[List[int]] = None) -> list:
         caches = []
         for layer_type in self.config.layer_types:
             if layer_type == "sliding_attention":
+                if self.config.sliding_window is None:
+                    raise ValueError(
+                        "sliding_attention layer without a sliding_window")
                 # Temporal and slack-backed: the draft path reads cache.state
                 # directly and indexes its rows as time.
                 caches.append(BufferedRotatingKVCache(
@@ -692,7 +695,15 @@ def draft_block(
 
 # --- target side --------------------------------------------------------------
 
-class DFlashCaptureHooks:
+if TYPE_CHECKING:
+    # The hooks are mixed in ahead of this class; giving pyright the real
+    # base resolves the super() calls against upstream's method set.
+    from mlx_vlm.models.qwen3_5.language import LanguageModel as _CaptureBase
+else:
+    _CaptureBase = object
+
+
+class DFlashCaptureHooks(_CaptureBase):
     """Packed-hidden capture for owned qwen3.5 LanguageModels.
 
     While armed, every hidden the engine sees is ``[trunk | cap ...]``: the

gmlx/spec/helpers.py

@@ -48,13 +48,19 @@ def _generation_stream():
 
 # --- draft/target sampler RNG coupling -------------------------------------
 
+def _rng_state() -> list[mx.array]:
+    # The mlx.core.random stub does not type ``state``.
+    return mx.random.state  # pyright: ignore[reportReturnType]
+
+
 def _copy_rng_state() -> list[mx.array]:
-    return [mx.array(state) for state in mx.random.state]
+    return [mx.array(state) for state in _rng_state()]
 
 
 def _restore_rng_state(state: list[mx.array]) -> None:
+    live = _rng_state()
     for i, value in enumerate(state):
-        mx.random.state[i] = value
+        live[i] = value
 
 
 def _append_arrays(value: Any, arrays: list[mx.array]) -> None:
@@ -91,8 +97,10 @@ class _SpeculativeSamplerRNG:
     def __init__(self, draft_model: nn.Module, *, enabled: bool):
         self.draft_model = draft_model
         self.enabled = bool(enabled)
-        self._target_rng_state = _copy_rng_state() if self.enabled else None
-        self._draft_rng_state = _copy_rng_state() if self.enabled else None
+        self._target_rng_state: list[mx.array] = (
+            _copy_rng_state() if self.enabled else [])
+        self._draft_rng_state: list[mx.array] = (
+            _copy_rng_state() if self.enabled else [])
 
     def draft_call(self, fn: Callable, *args, **kwargs):
         if not self.enabled:
@@ -109,7 +117,7 @@ def draft_call(self, fn: Callable, *args, **kwargs):
         result = fn(*args, **kwargs)
 
         arrays = _draft_sampler_state_arrays(self.draft_model)
-        arrays.extend(mx.random.state)
+        arrays.extend(_rng_state())
         if arrays:
             mx.async_eval(*arrays)
 
@@ -133,7 +141,7 @@ def draft_tokens(self, fn: Callable, *args, **kwargs):
         arrays = []
         _append_arrays(result, arrays)
         arrays.extend(_draft_sampler_state_arrays(self.draft_model))
-        arrays.extend(mx.random.state)
+        arrays.extend(_rng_state())
         if arrays:
             mx.async_eval(*arrays)
 
@@ -239,10 +247,11 @@ def _mtp_next_block_size(
         return min(budget, native)
     if getattr(draft_model, "prefer_requested_block_size", False):
         return budget
+    accept_lens: Any = draft_model.accept_lens
     return _effective_mtp_block_size(
         requested_block_total,
         configured_block_total,
-        draft_model.accept_lens,
+        accept_lens,
         remaining_budget,
     )
 
@@ -291,7 +300,7 @@ def _mtp_cache_offset(prompt_cache: list[Any]) -> Any:
 
 def _mtp_cache_offset_max(prompt_cache: list[Any]) -> int:
     offset = _mtp_cache_offset(prompt_cache)
-    return int(offset.max().item()) if isinstance(offset, mx.array) else int(offset)
+    return int(offset.max()) if isinstance(offset, mx.array) else int(offset)
 
 
 def _mtp_draft_position(kv_valid_len: Any) -> Any:
@@ -350,7 +359,7 @@ class _MTPVerifyResult:
 
 
 def _mtp_draft_hidden(lm: nn.Module, hidden: mx.array) -> mx.array:
-    prepare = getattr(lm, "speculative_draft_hidden", None)
+    prepare: Callable[..., Any] | None = getattr(lm, "speculative_draft_hidden", None)
     return prepare(hidden) if callable(prepare) else hidden
 
 
@@ -399,7 +408,8 @@ def _mtp_verify_without_logits(
     verify_input: mx.array,
     prompt_cache: list[Any],
 ) -> _MTPVerifyResult | None:
-    verify_hidden = getattr(lm, "speculative_verify_hidden", None)
+    verify_hidden: Callable[..., Any] | None = getattr(
+        lm, "speculative_verify_hidden", None)
     if callable(verify_hidden):
         _note_verify_branch("hook:speculative_verify_hidden", lm)
         result = verify_hidden(verify_input, prompt_cache)
@@ -481,7 +491,8 @@ def _mtp_verify_target(
     sample_target_tokens: bool = True,
 ) -> _MTPVerifyResult:
     if sample_target_tokens:
-        argmax_from_hidden = getattr(lm, "speculative_argmax_from_hidden", None)
+        argmax_from_hidden: Callable[..., Any] | None = getattr(
+            lm, "speculative_argmax_from_hidden", None)
         if callable(argmax_from_hidden):
             result = _mtp_verify_without_logits(lm, verify_input, prompt_cache)
             if result is not None:

gmlx/spec/mtp_drafter.py

@@ -83,7 +83,7 @@ def _cache_offset(caches) -> int:
         return 0
     off = getattr(caches[0], "offset", 0)
     if isinstance(off, mx.array):
-        return int(off.max().item()) if off.size else 0
+        return int(off.max()) if off.size else 0
     return int(off)
 
 
@@ -459,7 +459,7 @@ def accept_verified_tokens_batch(
             for cache in self._cache:
                 cache.trim(self._round_appended)
 
-        draft_rows = draft_tokens.tolist()
+        draft_rows: Any = draft_tokens.tolist()
         row_tokens: list[list[int]] = []
         row_hiddens: list[list[mx.array]] = []
         for row, accepted_i in enumerate(accepted):