Streaming¶

streaming ¶

Block streaming for low-RAM transformer inference on Apple Silicon.

LoraFuser `module-attribute` ¶

LoraFuser = Callable[[list[tuple[str, mx.array]], int, list['BlockLoraSource']], list[tuple[str, mx.array]]]

Callable signature for fusing LoRA deltas into a block's weight list.

Implementations receive the bound weights as [(param_name, array), ...], the block index being bound, and the list of :class:BlockLoraSource to fuse (already filtered by :meth:BlockLoraSource.has_block). They return the fused weight list (same shape; arrays may be replaced).

BlockLoraSource ¶

BlockLoraSource(lora_path: str | Path, block_prefix: str, strength: float = 1.0, key_mapper: Callable[[str], str | None] | None = None)

Per-block LoRA A/B matrices indexed by block index.

Streams LoRA A/B pairs from a memory-mapped safetensors file and indexes them by block index so a fuser can apply the matching delta to a streamed block on the fly.

Parameters:

Name	Type	Description	Default
`lora_path`	`str \| Path`	Path to LoRA safetensors.	required
`block_prefix`	`str`	Same prefix as the streamer (e.g. `"transformer.transformer_blocks."`). LoRA keys after optional remapping must start with this prefix and be of the form `f"{block_prefix}{idx}.{param}.lora_A.weight"`.	required
`strength`	`float`	LoRA fusion strength (default 1.0).	`1.0`
`key_mapper`	`Callable[[str], str \| None] \| None`	Optional callable that remaps raw safetensors keys to the model's state-dict naming (e.g. ComfyUI/diffusers → MLX). Returning `None` drops the key.	`None`

has_block ¶

has_block(block_idx: int) -> bool

True iff at least one matched A/B pair exists for block_idx.

get_block_lora_dict ¶

get_block_lora_dict(block_idx: int) -> dict[str, array]

Per-block LoRA dict shaped like {param.lora_A.weight: array, ...}.

Keys use the param name relative to the block (the same keys a :class:BlockStreamer would emit).

close ¶

close() -> None

Release the mmap'd LoRA dict. After this the source is unusable.

BlockStreamer ¶

BlockStreamer(weight_paths: str | Path | Iterable[str | Path], block_prefix: str, lora_fuser: LoraFuser | None = None)

Stream transformer block weights from mmap'd safetensors.

Parameters:

Name	Type	Description	Default
`weight_paths`	`str \| Path \| Iterable[str \| Path]`	One or more `.safetensors` paths whose union contains every key for the streamed blocks. Loaded via :func:`mlx.core.load` which memory-maps the file.	required
`block_prefix`	`str`	State-dict key prefix that identifies block weights, e.g. `"transformer.transformer_blocks."`. Keys of the form `f"{block_prefix}{i}.{rest}"` (where `i` is an integer) are treated as block `i`'s weights.	required
`lora_fuser`	`LoraFuser \| None`	Optional callable to fuse LoRA deltas into the bound weights. See :data:`LoraFuser`. If `None`, the `lora_sources` argument to :meth:`bind` is ignored.	`None`

Notes

After construction, :attr:block_count returns the number of distinct block indices found, and :meth:block_keys(i) lists the per-block parameter names that will be bound by :meth:bind. The mmap'd dict is held until :meth:close is called.

block_count `property` ¶

block_count: int

Number of distinct block indices discovered in the safetensors.

block_prefix `property` ¶

block_prefix: str

State-dict key prefix identifying streamed block weights.

block_keys ¶

block_keys(idx: int) -> list[str]

Per-block parameter names (without the {prefix}{idx}. part).

bind ¶

bind(block: Module, idx: int, evict_previous: int | None = None, lora_sources: list[BlockLoraSource] | None = None) -> None

Load block idx's weights into block in-place.

After this returns, block's parameters reference the safetensors-mapped arrays for index idx. A subsequent :meth:bind to a different idx rebinds them, releasing the previous arrays from the block's parameter tree.

Parameters:

Name	Type	Description	Default
`block`	`Module`	Target module. Its parameter tree must match the keys returned by :meth:`block_keys`.	required
`idx`	`int`	Block index to load.	required
`evict_previous`	`int \| None`	If given, drop the cached array references for that block index from the internal weight dict before binding. The streamer holds refs to every block's weights; without eviction, those refs prevent GC even after the bound block is replaced. For streaming inference, pass the previously-bound index so peak resident memory stays at ~one block.	`None`
`lora_sources`	`list[BlockLoraSource] \| None`	Optional LoRA sources to fuse into the bound weights. Requires a `lora_fuser` at construction.	`None`

close ¶

close() -> None

Release the mmap'd dict. After this the streamer is unusable.

block_streaming ¶

Block streaming for low-RAM inference on Apple Silicon.

Stream transformer block weights from a memory-mapped safetensors file into a single shared block module, so peak resident memory stays at ~1 block instead of ~num_blocks.

The pattern fits any model whose state-dict has block-indexed keys of the form f"{prefix}{i}.{rest}" (e.g. "transformer.transformer_blocks.0.attn.q_proj.weight"). It works because:

Apple Silicon has unified memory, so there is no host-to-device copy.
mx.load(path) memory-maps safetensors lazily — opening a 20 GB file costs ~40 MB RSS until individual arrays are touched.
MLX has a single command queue, so per-block mx.clear_cache() + dropping references keeps the resident set bounded without explicit stream/event coordination.

Typical memory profile for a 22 B bf16 transformer with 48 blocks: without streaming ~22 GB resident; with streaming ~1 block (~460 MB) + non-block params + mmap metadata ≈ ~1 GB.

LoRA fusion is decoupled: pass a lora_fuser callable to :class:BlockStreamer and it will be invoked per block. The fuser receives the bound weights and a list of :class:BlockLoraSource objects, and returns the fused weight list. This keeps quantization- aware fusion strategies out of arsenal.

LoraFuser `module-attribute` ¶

LoraFuser = Callable[[list[tuple[str, mx.array]], int, list['BlockLoraSource']], list[tuple[str, mx.array]]]

Callable signature for fusing LoRA deltas into a block's weight list.

Implementations receive the bound weights as [(param_name, array), ...], the block index being bound, and the list of :class:BlockLoraSource to fuse (already filtered by :meth:BlockLoraSource.has_block). They return the fused weight list (same shape; arrays may be replaced).

BlockLoraSource ¶

BlockLoraSource(lora_path: str | Path, block_prefix: str, strength: float = 1.0, key_mapper: Callable[[str], str | None] | None = None)

Per-block LoRA A/B matrices indexed by block index.

Streams LoRA A/B pairs from a memory-mapped safetensors file and indexes them by block index so a fuser can apply the matching delta to a streamed block on the fly.

Parameters:

Name	Type	Description	Default
`lora_path`	`str \| Path`	Path to LoRA safetensors.	required
`block_prefix`	`str`	Same prefix as the streamer (e.g. `"transformer.transformer_blocks."`). LoRA keys after optional remapping must start with this prefix and be of the form `f"{block_prefix}{idx}.{param}.lora_A.weight"`.	required
`strength`	`float`	LoRA fusion strength (default 1.0).	`1.0`
`key_mapper`	`Callable[[str], str \| None] \| None`	Optional callable that remaps raw safetensors keys to the model's state-dict naming (e.g. ComfyUI/diffusers → MLX). Returning `None` drops the key.	`None`

has_block ¶

has_block(block_idx: int) -> bool

True iff at least one matched A/B pair exists for block_idx.

get_block_lora_dict ¶

get_block_lora_dict(block_idx: int) -> dict[str, array]

Per-block LoRA dict shaped like {param.lora_A.weight: array, ...}.

Keys use the param name relative to the block (the same keys a :class:BlockStreamer would emit).

close ¶

close() -> None

Release the mmap'd LoRA dict. After this the source is unusable.

BlockStreamer ¶

BlockStreamer(weight_paths: str | Path | Iterable[str | Path], block_prefix: str, lora_fuser: LoraFuser | None = None)

Stream transformer block weights from mmap'd safetensors.

Parameters:

Name	Type	Description	Default
`weight_paths`	`str \| Path \| Iterable[str \| Path]`	One or more `.safetensors` paths whose union contains every key for the streamed blocks. Loaded via :func:`mlx.core.load` which memory-maps the file.	required
`block_prefix`	`str`	State-dict key prefix that identifies block weights, e.g. `"transformer.transformer_blocks."`. Keys of the form `f"{block_prefix}{i}.{rest}"` (where `i` is an integer) are treated as block `i`'s weights.	required
`lora_fuser`	`LoraFuser \| None`	Optional callable to fuse LoRA deltas into the bound weights. See :data:`LoraFuser`. If `None`, the `lora_sources` argument to :meth:`bind` is ignored.	`None`

Notes

After construction, :attr:block_count returns the number of distinct block indices found, and :meth:block_keys(i) lists the per-block parameter names that will be bound by :meth:bind. The mmap'd dict is held until :meth:close is called.

block_count `property` ¶

block_count: int

Number of distinct block indices discovered in the safetensors.

block_prefix `property` ¶

block_prefix: str

State-dict key prefix identifying streamed block weights.

block_keys ¶

block_keys(idx: int) -> list[str]

Per-block parameter names (without the {prefix}{idx}. part).

bind ¶

bind(block: Module, idx: int, evict_previous: int | None = None, lora_sources: list[BlockLoraSource] | None = None) -> None

Load block idx's weights into block in-place.

After this returns, block's parameters reference the safetensors-mapped arrays for index idx. A subsequent :meth:bind to a different idx rebinds them, releasing the previous arrays from the block's parameter tree.

Parameters:

Name	Type	Description	Default
`block`	`Module`	Target module. Its parameter tree must match the keys returned by :meth:`block_keys`.	required
`idx`	`int`	Block index to load.	required
`evict_previous`	`int \| None`	If given, drop the cached array references for that block index from the internal weight dict before binding. The streamer holds refs to every block's weights; without eviction, those refs prevent GC even after the bound block is replaced. For streaming inference, pass the previously-bound index so peak resident memory stays at ~one block.	`None`
`lora_sources`	`list[BlockLoraSource] \| None`	Optional LoRA sources to fuse into the bound weights. Requires a `lora_fuser` at construction.	`None`

close ¶

close() -> None

Release the mmap'd dict. After this the streamer is unusable.

Streaming¶

streaming ¶

LoraFuser module-attribute ¶

BlockLoraSource ¶

has_block ¶

get_block_lora_dict ¶

close ¶

BlockStreamer ¶

block_count property ¶

block_prefix property ¶

block_keys ¶

bind ¶

close ¶

block_streaming ¶

LoraFuser module-attribute ¶

BlockLoraSource ¶

has_block ¶

get_block_lora_dict ¶

close ¶

BlockStreamer ¶

block_count property ¶

block_prefix property ¶

block_keys ¶

bind ¶

close ¶

LoraFuser `module-attribute` ¶

block_count `property` ¶

block_prefix `property` ¶

LoraFuser `module-attribute` ¶

block_count `property` ¶

block_prefix `property` ¶