mlx-dfloat 0.1.0__py3-none-any.whl

mlx_dfloat/integrate/names.py

@@ -0,0 +1,121 @@
+"""Name maps: how a checkpoint's matrix names land on a module's attributes."""
+
+import re
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import mlx.nn as nn
+
+from mlx_dfloat.errors import DFloatIntegrationError
+
+BlockShapes = dict[str, tuple[int, ...]]
+Shapes = dict[str, BlockShapes]
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class Placement:
+    """Where one checkpoint matrix lands: the block name and the dotted attribute path of its module."""
+
+    block: str
+    attr: str
+
+
+class NameMap(Protocol):
+    """The adapter-specific naming an integration needs; everything else in this package is generic."""
+
+    kinds: tuple[str, ...]
+
+    def attrs_of(self, kind: str) -> tuple[str, ...]:
+        """Attribute paths of every matrix module of a block of ``kind``, in the map's own order.
+
+        The order is the map's construction order, not the checkpoint's: a checkpoint group's own
+        concatenation order lives in its ``matrix_names``, recovered per matrix through ``place()``.
+        """
+        ...
+
+    def place(self, matrix_name: str) -> Placement:
+        """The block and attribute path of a checkpoint matrix name."""
+        ...
+
+    def kind_of(self, block_name: str) -> str:
+        """The kind of a block name."""
+        ...
+
+    def param_name(self, checkpoint_name: str) -> str:
+        """The module parameter name of a non-matrix checkpoint tensor (biases, norms, embedders)."""
+        ...
+
+    def is_matrix_module(self, module: Any) -> bool:
+        """Whether ``module`` is one the seam swaps weights on."""
+        ...
+
+
+class StaticNameMap:
+    """A name map from explicit tables: kind -> {checkpoint sub-path: attribute path}.
+
+    Block names are ``<kind>.<index>``; a matrix name is ``<block>.<sub-path>.weight``; matrix modules are
+    ``nn.Linear``.
+    """
+
+    def __init__(self, tables: Mapping[str, Mapping[str, str]]) -> None:
+        """Keep one table per kind, checkpoint sub-path -> attribute path, in the given order."""
+        self._tables = {k: dict(v) for k, v in tables.items()}
+        self.kinds = tuple(self._tables)
+        # `[0-9]`, not `\d`: `\d` matches any Unicode digit, and int("٣") == 3 would alias block 3.
+        self._block = re.compile(
+            rf"^({'|'.join(re.escape(k) for k in self.kinds)})\.([0-9]+)\.(.+)$"
+        )
+
+    def attrs_of(self, kind: str) -> tuple[str, ...]:
+        """Attribute paths of every matrix module of a block of ``kind``, in table order."""
+        return tuple(self._tables[kind].values())
+
+    def _split(self, name: str) -> tuple[str, int, str] | None:
+        match = self._block.match(name)
+        return None if match is None else (match.group(1), int(match.group(2)), match.group(3))
+
+    def place(self, matrix_name: str) -> Placement:
+        """The block and attribute path of a checkpoint matrix name.
+
+        Raises:
+            DFloatIntegrationError: The name is not a ``<kind>.<index>.<sub-path>.weight`` name the
+                table covers.
+        """
+        parsed = self._split(matrix_name)
+        if parsed is None or not parsed[2].endswith(".weight"):
+            raise DFloatIntegrationError(f"{matrix_name!r} is not a block matrix name")
+        kind, idx, rest = parsed
+        attr = self._tables[kind].get(rest.removesuffix(".weight"))
+        if attr is None:
+            raise DFloatIntegrationError(f"{matrix_name!r}: not a {kind} matrix the map covers")
+        return Placement(block=f"{kind}.{idx}", attr=attr)
+
+    def kind_of(self, block_name: str) -> str:
+        """The kind of a ``<kind>.<index>`` block name.
+
+        Raises:
+            DFloatIntegrationError: ``block_name`` is not a block name of a known kind.
+        """
+        kind, _dot, idx = block_name.partition(".")
+        if kind not in self._tables or not (idx.isascii() and idx.isdigit()):
+            raise DFloatIntegrationError(f"{block_name!r} is not a block name")
+        return kind
+
+    def param_name(self, checkpoint_name: str) -> str:
+        """The module parameter name of a non-matrix checkpoint tensor.
+
+        A block extra is renamed through the same table as its block's matrices; every other name
+        (already the module's own) is returned unchanged.
+        """
+        parsed = self._split(checkpoint_name)
+        if parsed is None:
+            return checkpoint_name
+        kind, idx, rest = parsed
+        head, _dot, leaf = rest.rpartition(".")
+        mapped = self._tables[kind].get(head)
+        return checkpoint_name if mapped is None else f"{kind}.{idx}.{mapped}.{leaf}"
+
+    def is_matrix_module(self, module: Any) -> bool:
+        """Whether ``module`` is an ``nn.Linear``, the only matrix module this map swaps weights on."""
+        return isinstance(module, nn.Linear)  # type: ignore[attr-defined]

mlx_dfloat/integrate/placeholders.py

@@ -0,0 +1,72 @@
+"""Zero-size placeholders in place of every block matrix, and the shapes they had."""
+
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import mlx.core as mx
+
+from mlx_dfloat.errors import DFloatIntegrationError
+from mlx_dfloat.integrate.names import BlockShapes, NameMap, Shapes
+
+PLACEHOLDER = mx.zeros((0,), dtype=mx.bfloat16)
+
+
+def _is_index(part: str) -> bool:
+    """An ASCII-digit component (``"0"``); ``"²".isdigit()`` is true too, but ``int("²")`` raises."""
+    return part.isascii() and part.isdigit()
+
+
+def get_attr_path(module: Any, path: str) -> Any:
+    """Resolve a dotted path on a module; a digit component indexes a list (``attn.to_out.0``).
+
+    Raises:
+        DFloatIntegrationError: A component does not exist.
+    """
+    node = module
+    for part in path.split("."):
+        try:
+            node = node[int(part)] if _is_index(part) else getattr(node, part)
+        except (AttributeError, IndexError, KeyError, TypeError) as exc:
+            raise DFloatIntegrationError(f"{path!r}: no {part!r} on {type(node).__name__}") from exc
+    return node
+
+
+def set_attr_path(module: Any, path: str, value: Any) -> None:
+    """Assign ``value`` at a dotted path (see ``get_attr_path``)."""
+    head, _dot, leaf = path.rpartition(".")
+    parent = get_attr_path(module, head) if head else module
+    if _is_index(leaf):
+        parent[int(leaf)] = value
+    else:
+        setattr(parent, leaf, value)
+
+
+def install_placeholders(
+    block_lists: Iterable[tuple[str, Sequence[Any]]], name_map: NameMap
+) -> Shapes:
+    """Replace every block matrix weight with ``PLACEHOLDER``; return block name -> attr -> shape.
+
+    Every mapped attribute must be a matrix module of the block and every matrix module of the block must
+    be mapped, so an added or renamed layer is caught here, not at the first matmul.
+
+    Raises:
+        DFloatIntegrationError: A block's matrix-module set differs from the map.
+    """
+    shapes: Shapes = {}
+    for kind, blocks in block_lists:
+        mapped = set(name_map.attrs_of(kind))
+        for idx, block in enumerate(blocks):
+            name = f"{kind}.{idx}"
+            present = {p for p, m in block.named_modules() if name_map.is_matrix_module(m)}
+            if present != mapped:
+                raise DFloatIntegrationError(
+                    f"{name}: mapped matrices missing from the block: {sorted(mapped - present)}; "
+                    f"Linear layers the map does not cover: {sorted(present - mapped)}"
+                )
+            per: BlockShapes = {}
+            for attr in name_map.attrs_of(kind):
+                module = get_attr_path(block, attr)
+                per[attr] = tuple(int(d) for d in module.weight.shape)
+                module.weight = PLACEHOLDER
+            shapes[name] = per
+    return shapes

mlx_dfloat/integrate/providers.py

@@ -0,0 +1,271 @@
+"""Weight providers: what the seam asks for one block's matrices."""
+
+from collections.abc import Callable, Mapping, Sequence
+from functools import partial
+from pathlib import Path
+from typing import Protocol, cast
+
+import mlx.core as mx
+import numpy as np
+
+from mlx_dfloat._safetensors import TensorInfo, read_array
+from mlx_dfloat.decode import DecodeResult, check_status, decode_group, split_matrices
+from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError
+from mlx_dfloat.format import MxGroup
+from mlx_dfloat.integrate.names import BlockShapes, NameMap
+
+# "per-block": evaluate each block's output. "depth2": async_eval it, then eval the previous
+# block's (hides host-side allocation while the GPU runs the block; not a decode/compute overlap).
+# "none": no evaluation inside the step (a non-launching provider only, or a few blocks).
+EVAL_POLICIES: tuple[str, ...] = ("per-block", "depth2", "none")
+
+
+class WeightProvider(Protocol):
+    """Hands the seam one block's matrices as ``{attribute path: bf16 array of the block's shape}``."""
+
+    launches: int
+    launching: bool
+    policies: tuple[str, ...]
+
+    def weights_for(self, block_name: str, shapes: BlockShapes) -> dict[str, mx.array]:
+        """Weights for ``block_name``; ``shapes`` is that block's entry from ``install_placeholders``."""
+        ...
+
+    def verify(self) -> None:
+        """Check what ``weights_for`` deferred (decode status words); call it after the step's eval."""
+        ...
+
+    def reset(self) -> None:
+        """Drop per-step state after a step that raised."""
+        ...
+
+
+def read_bf16(path: Path, info: TensorInfo) -> mx.array:
+    """One BF16 tensor of a safetensors file as a bf16 array (a bit view, never a cast).
+
+    Raises:
+        DFloatFormatError: The tensor is not BF16.
+    """
+    if info.dtype != "BF16":
+        raise DFloatFormatError(f"{info.name}: tensor is {info.dtype}, expected BF16")
+    return mx.array(np.ascontiguousarray(read_array(path, info))).view(mx.bfloat16)
+
+
+class StreamingBF16Provider:
+    """Reads each block's BF16 matrices from safetensors shards as the block runs; nothing stays resident.
+
+    ``index`` maps a checkpoint matrix name to ``(shard path, tensor info)``: a base repository's
+    weight index plus the shards' headers. The reference side of an image identity check, and the
+    seam's way of streaming an uncompressed model.
+    """
+
+    launching = False
+    # Each block reads fresh arrays from the shards; under "none" the whole BF16 transformer would
+    # stay alive in one step's lazy graph, so only the per-block boundaries are offered.
+    policies: tuple[str, ...] = ("per-block", "depth2")
+
+    def __init__(
+        self,
+        index: Mapping[str, tuple[Path, TensorInfo]],
+        matrix_names: Mapping[str, Sequence[str]],
+        name_map: NameMap,
+    ) -> None:
+        """Hold the index, each block's matrix names (the checkpoint's order) and the map that places them."""
+        self._index = index
+        self._matrix_names = matrix_names
+        self._name_map = name_map
+        self.launches = 0
+        self.reads = 0
+
+    def verify(self) -> None:
+        """Nothing deferred: no decode happened."""
+
+    def reset(self) -> None:
+        """No per-step state."""
+
+    def weights_for(self, block_name: str, shapes: BlockShapes) -> dict[str, mx.array]:
+        """Read the block's matrices from their shards (fresh arrays; the caller's step frees them).
+
+        Raises:
+            DFloatIntegrationError: No matrix names for the block, a name missing from the index, a
+                name that is not a matrix of this block, or a tensor of the wrong shape.
+            DFloatFormatError: A tensor that is not BF16.
+        """
+        names = self._matrix_names.get(block_name)
+        if names is None:
+            raise DFloatIntegrationError(f"{block_name}: no matrix names")
+        weights: dict[str, mx.array] = {}
+        for matrix_name in names:
+            entry = self._index.get(matrix_name)
+            if entry is None:
+                raise DFloatIntegrationError(f"{matrix_name}: not in the BF16 index")
+            placement = self._name_map.place(matrix_name)
+            if placement.block != block_name or placement.attr not in shapes:
+                raise DFloatIntegrationError(
+                    f"{matrix_name}: not a matrix of {block_name} with a known shape"
+                )
+            array = read_bf16(*entry)
+            self.reads += 1
+            if tuple(array.shape) != tuple(shapes[placement.attr]):
+                raise DFloatIntegrationError(
+                    f"{matrix_name}: shard tensor has shape {tuple(array.shape)}, "
+                    f"the block needs {shapes[placement.attr]}"
+                )
+            weights[placement.attr] = array
+        return weights
+
+
+class DF11Provider:
+    """Decodes each block's DF11 group when asked: one decode launch per call.
+
+    The status words are not read inside ``weights_for`` (a host read there would sync before every
+    block); they queue in ``pending`` and ``verify`` reads them after the step's final eval.
+    """
+
+    launching = True
+    policies = EVAL_POLICIES
+
+    def __init__(
+        self,
+        resident: Mapping[str, MxGroup],
+        matrix_names: Mapping[str, Sequence[str]],
+        name_map: NameMap,
+        *,
+        decode: Callable[[MxGroup], DecodeResult] | None = None,
+    ) -> None:
+        """Hold the resident groups, each group's matrix names, and the name map that places them.
+
+        ``decode`` defaults to the Metal backend; tests inject a counting reference decode.
+        """
+        self._resident = resident
+        self._matrix_names = matrix_names
+        self._name_map = name_map
+        self._decode = decode if decode is not None else partial(decode_group, backend="metal")
+        self.launches = 0
+        self.pending: list[tuple[str, mx.array]] = []
+
+    def verify(self) -> None:
+        """Read every pending status word on the host and clear the list.
+
+        Raises:
+            DFloatFormatError: A block's decode reported an error; the message names the block.
+        """
+        pending, self.pending = self.pending, []
+        for block_name, status in pending:
+            check_status(status, name=block_name)
+
+    def reset(self) -> None:
+        """Forget queued status words (their decodes belong to a step that did not finish)."""
+        self.pending = []
+
+    def weights_for(self, block_name: str, shapes: BlockShapes) -> dict[str, mx.array]:
+        """Decode the block's group and cut it into views of the block's shapes (no copies, no host read).
+
+        Raises:
+            DFloatIntegrationError: The block is not resident or has no matrix names (both refused before
+                any decode), or the names do not fit the decoded group: their count, their block, or a
+                matrix's size against its shape.
+        """
+        group = self._resident.get(block_name)
+        if group is None:
+            raise DFloatIntegrationError(f"{block_name}: no resident DF11 group")
+        names = self._matrix_names.get(block_name)
+        if names is None:
+            raise DFloatIntegrationError(f"{block_name}: no matrix names for its DF11 group")
+        result = self._decode(group)
+        self.pending.append((block_name, result.status))
+        parts = split_matrices(result.bits, group.split_positions)
+        if len(names) != len(parts):
+            raise DFloatIntegrationError(
+                f"{block_name}: {len(parts)} matrices decoded for {len(names)} names"
+            )
+        weights: dict[str, mx.array] = {}
+        for matrix_name, part in zip(names, parts, strict=True):
+            placement = self._name_map.place(matrix_name)
+            if placement.block != block_name or placement.attr not in shapes:
+                raise DFloatIntegrationError(
+                    f"{matrix_name}: not a matrix of {block_name} with a known shape"
+                )
+            shape = shapes[placement.attr]
+            n = 1
+            for d in shape:
+                n *= d
+            if part.size != n:
+                raise DFloatIntegrationError(
+                    f"{matrix_name}: decoded {part.size} elements, expected {shape} = {n}"
+                )
+            weights[placement.attr] = part.view(mx.bfloat16).reshape(shape)
+        self.launches += 1
+        return weights
+
+
+class ReuseProvider:
+    """One pre-decoded dict per block kind, returned for every block of that kind; never launches."""
+
+    launching = False
+    policies = EVAL_POLICIES
+
+    def __init__(self, per_kind: Mapping[str, Mapping[str, mx.array]], name_map: NameMap) -> None:
+        """Keep one dict per kind, handed back as-is (by identity) for every block of that kind."""
+        self._per_kind: dict[str, Mapping[str, mx.array]] = dict(per_kind)
+        self._name_map = name_map
+        self.launches = 0
+
+    def verify(self) -> None:
+        """Nothing deferred: no decode happened."""
+
+    def reset(self) -> None:
+        """No per-step state."""
+
+    def weights_for(self, block_name: str, shapes: BlockShapes) -> dict[str, mx.array]:
+        """The kind's dict, checked against the block's shapes.
+
+        Raises:
+            DFloatIntegrationError: No dict for the block's kind, or a shape differs.
+        """
+        kind = self._name_map.kind_of(block_name)
+        weights = self._per_kind.get(kind)
+        if weights is None:
+            raise DFloatIntegrationError(f"{block_name}: no reusable block of kind {kind!r}")
+        _check_shapes(block_name, weights, shapes)
+        return cast("dict[str, mx.array]", weights)
+
+
+class ResidentProvider:
+    """Per-block pre-decoded dicts (every block resident at once); never launches."""
+
+    launching = False
+    policies = EVAL_POLICIES
+
+    def __init__(self, per_block: Mapping[str, Mapping[str, mx.array]]) -> None:
+        """Keep the per-block dicts, handed back as-is (by identity) for their own block."""
+        self._per_block: dict[str, Mapping[str, mx.array]] = dict(per_block)
+        self.launches = 0
+
+    def verify(self) -> None:
+        """Nothing deferred: no decode happened."""
+
+    def reset(self) -> None:
+        """No per-step state."""
+
+    def weights_for(self, block_name: str, shapes: BlockShapes) -> dict[str, mx.array]:
+        """The block's own dict, checked against its shapes.
+
+        Raises:
+            DFloatIntegrationError: The block is not resident, or a shape differs.
+        """
+        weights = self._per_block.get(block_name)
+        if weights is None:
+            raise DFloatIntegrationError(f"{block_name}: not resident")
+        _check_shapes(block_name, weights, shapes)
+        return cast("dict[str, mx.array]", weights)
+
+
+def _check_shapes(block_name: str, weights: Mapping[str, mx.array], shapes: BlockShapes) -> None:
+    for attr, shape in shapes.items():
+        if attr not in weights:
+            raise DFloatIntegrationError(f"{block_name}: no weight for {attr!r}")
+        if tuple(weights[attr].shape) != tuple(shape):
+            raise DFloatIntegrationError(
+                f"{block_name}.{attr}: weight has shape {tuple(weights[attr].shape)}, the block needs {shape}"
+            )

mlx_dfloat/integrate/seam.py

@@ -0,0 +1,196 @@
+"""The block seam: assign one block's weights, run it, evaluate per policy, restore the placeholders."""
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import mlx.core as mx
+
+from mlx_dfloat.errors import DFloatIntegrationError
+from mlx_dfloat.integrate.names import Shapes
+from mlx_dfloat.integrate.placeholders import PLACEHOLDER, get_attr_path
+from mlx_dfloat.integrate.providers import EVAL_POLICIES, WeightProvider
+
+# "none" with a launching provider keeps every decoded group alive
+MAX_NONE_POLICY_LAUNCHING_BLOCKS = 2
+
+# The seam calls MLX through these names so a test can record the evaluation order.
+_eval = mx.eval
+_async_eval = mx.async_eval
+
+
+class TraceHook(Protocol):
+    """Receives one event per block per step and the step boundaries; four clock reads per block, no sync."""
+
+    def begin_step(self) -> None:
+        """Open a new step (called once, before its first block runs)."""
+        ...
+
+    def end_step(self) -> None:
+        """Close the step (called once, after its last block's eval)."""
+        ...
+
+    def record(
+        self,
+        block: str,
+        *,
+        decode_start: float,
+        decode_end: float,
+        encode_end: float,
+        eval_end: float,
+    ) -> None:
+        """Record one block's four clock reads: decode, the mflux call, and the policy's eval."""
+        ...
+
+
+@dataclass(slots=True)
+class SeamState:
+    """Mutable per-attachment state the seam threads through a step: never rebuilt mid-step."""
+
+    provider: WeightProvider
+    shapes: Shapes
+    policy: str
+    tracer: TraceHook | None = None
+    verify_in_call: bool = False
+    prev: Any = None  # depth2: the previous block's output, evaluated after the next is queued
+
+
+def attach_state(
+    provider: WeightProvider,
+    shapes: Shapes,
+    *,
+    eval_policy: str = "per-block",
+    tracer: TraceHook | None = None,
+    verify_in_call: bool = False,
+) -> SeamState:
+    """Validate the policy for this provider and build the seam's state.
+
+    Raises:
+        DFloatIntegrationError: Unknown policy; a policy the provider does not run under; ``"none"`` with a
+            launching provider over more than ``MAX_NONE_POLICY_LAUNCHING_BLOCKS`` blocks; or
+            ``verify_in_call`` with ``"none"`` and a launching provider (reading the status words at the
+            end of the call would force every decode before the caller's own eval).
+    """
+    if eval_policy not in EVAL_POLICIES:
+        raise DFloatIntegrationError(
+            f"unknown eval policy {eval_policy!r}; choose from {EVAL_POLICIES}"
+        )
+    if eval_policy not in provider.policies:
+        raise DFloatIntegrationError(
+            f"{type(provider).__name__} runs only under {provider.policies}, not {eval_policy!r}"
+        )
+    if (
+        eval_policy == "none"
+        and provider.launching
+        and len(shapes) > MAX_NONE_POLICY_LAUNCHING_BLOCKS
+    ):
+        raise DFloatIntegrationError(
+            f"eval policy 'none' with a launching provider over {len(shapes)} blocks would keep every decoded "
+            f"group resident until the final eval; 'none' is for non-launching providers (or at most "
+            f"{MAX_NONE_POLICY_LAUNCHING_BLOCKS} blocks)"
+        )
+    if verify_in_call and eval_policy == "none" and provider.launching:
+        raise DFloatIntegrationError(
+            "verify_in_call with eval policy 'none' and a launching provider would read the status "
+            "words before the caller's eval, forcing every decode of the step; verify_step() after "
+            "the step's eval instead"
+        )
+    return SeamState(
+        provider=provider,
+        shapes=shapes,
+        policy=eval_policy,
+        tracer=tracer,
+        verify_in_call=verify_in_call,
+    )
+
+
+def begin_step(state: SeamState) -> None:
+    """Refuse a step while the last one's status words are unchecked; open the tracer's step.
+
+    Raises:
+        DFloatIntegrationError: The provider still holds status words from an earlier step (nobody
+            called ``verify_step()`` after it); they are kept, so a ``verify_step()`` can still check them.
+    """
+    pending = getattr(state.provider, "pending", None)
+    if pending:
+        raise DFloatIntegrationError(
+            f"{len(pending)} decode status words from an earlier step are unchecked; call "
+            "verify_step() after each step's eval, or attach with verify_in_call=True"
+        )
+    if state.tracer is not None:
+        state.tracer.begin_step()
+
+
+def end_step(state: SeamState, out: Any) -> Any:
+    """Drain depth-2's tail, verify when asked, close the trace; returns ``out``."""
+    if state.prev is not None:
+        _eval(state.prev)
+        state.prev = None
+    if state.verify_in_call:
+        state.provider.verify()
+    if state.tracer is not None:
+        state.tracer.end_step()
+    return out
+
+
+def abort_step(state: SeamState) -> None:
+    """After a step that raised: no stale look-ahead, no stale status words, no depth-2 tail."""
+    state.prev = None
+    state.provider.reset()
+
+
+def verify_step(state: SeamState) -> None:
+    """Run the provider's deferred checks (an explicit call, for a step run without ``verify_in_call``)."""
+    state.provider.verify()
+
+
+def _seam_eval(state: SeamState, out: Any) -> None:
+    if state.policy == "per-block":
+        _eval(out)
+    elif state.policy == "depth2":
+        _async_eval(out)  # type: ignore[no-untyped-call]  # mlx's async_eval stub carries no annotations
+        if state.prev is not None:
+            _eval(state.prev)
+        state.prev = out
+
+
+def run_block(state: SeamState, block_name: str, block: Any, run: Callable[[], Any]) -> Any:
+    """Assign the block's weights, run it, evaluate per policy, restore the placeholders (also on a raise).
+
+    Raises:
+        DFloatIntegrationError: The provider's dict does not cover the block's matrices, or a weight's
+            shape does not match the block.
+    """
+    shapes = state.shapes[block_name]
+    t_decode_start = time.perf_counter()
+    weights = state.provider.weights_for(block_name, shapes)
+    t_decode_end = time.perf_counter()
+    if weights.keys() != shapes.keys():
+        raise DFloatIntegrationError(
+            f"{block_name}: provider returned {sorted(weights)}; the block needs {sorted(shapes)} "
+            f"(missing {sorted(shapes.keys() - weights.keys())})"
+        )
+    for attr, shape in shapes.items():
+        if tuple(weights[attr].shape) != tuple(shape):
+            raise DFloatIntegrationError(
+                f"{block_name}.{attr}: weight has shape {tuple(weights[attr].shape)}, the block needs {shape}"
+            )
+    try:
+        for attr, weight in weights.items():
+            get_attr_path(block, attr).weight = weight
+        out = run()
+        t_encode_end = time.perf_counter()
+        _seam_eval(state, out)
+        if state.tracer is not None:
+            state.tracer.record(
+                block_name,
+                decode_start=t_decode_start,
+                decode_end=t_decode_end,
+                encode_end=t_encode_end,
+                eval_end=time.perf_counter(),
+            )
+        return out
+    finally:
+        for attr in shapes:
+            get_attr_path(block, attr).weight = PLACEHOLDER

mlx_dfloat/mflux/__init__.py

@@ -0,0 +1,31 @@
+"""mflux adapters. Everything here imports mflux lazily; ``import mlx_dfloat`` never needs it."""
+
+import importlib.util
+from typing import Any
+
+from mlx_dfloat.errors import DFloatDependencyError
+
+
+def require_mflux() -> None:
+    """Raise a package-rooted error when the optional ``mflux`` extra is missing.
+
+    Raises:
+        DFloatDependencyError: ``mflux`` cannot be imported.
+    """
+    if importlib.util.find_spec("mflux") is None:
+        raise DFloatDependencyError(
+            "the mflux adapters need the optional extra: install mlx-dfloat[mflux]"
+        )
+
+
+def __getattr__(name: str) -> Any:
+    """``DFloatFlux1`` is imported on first use, so ``import mlx_dfloat.mflux`` never needs mflux."""
+    if name == "DFloatFlux1":
+        from mlx_dfloat.mflux.flux1.model import DFloatFlux1
+
+        return DFloatFlux1
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+# DFloatFlux1 stays out of __all__: a star-import must not need mflux; `mlx_dfloat.mflux.DFloatFlux1` still works.
+__all__ = ["require_mflux"]