`vllm_mlx.mlx_streams`¶

Helpers for binding MLX generation streams to worker threads.

View the complete module source at #L1-L39.

API details¶

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.mlx_streams ¶

Helpers for binding MLX generation streams to worker threads.

vllm_mlx.mlx_streams._STREAM_REBIND_LOCK `module-attribute` ¶

_STREAM_REBIND_LOCK = threading.Lock()

vllm_mlx.mlx_streams.bind_generation_streams ¶

bind_generation_streams(module_names: Iterable[str] = ('mlx_lm.generate', 'mlx_vlm.generate')) -> object

Bind mlx-lm/mlx-vlm generation streams to the current thread.

MLX streams are thread-local. If a model is loaded on one thread and generation runs on another, module-level generation streams created during import can point at a stream that does not exist in the worker thread.

This intentionally creates a fresh stream for the current worker call and replaces module-level generation_stream handles under a process-local lock. It is an admission/ownership fix, not a batching optimization; callers should invoke it at worker-entry boundaries rather than inside token loops.

Source code in vllm_mlx/mlx_streams.py

def bind_generation_streams(
    module_names: Iterable[str] = ("mlx_lm.generate", "mlx_vlm.generate"),
) -> object:
    """Bind mlx-lm/mlx-vlm generation streams to the current thread.

    MLX streams are thread-local. If a model is loaded on one thread and
    generation runs on another, module-level generation streams created during
    import can point at a stream that does not exist in the worker thread.

    This intentionally creates a fresh stream for the current worker call and
    replaces module-level generation_stream handles under a process-local lock.
    It is an admission/ownership fix, not a batching optimization; callers
    should invoke it at worker-entry boundaries rather than inside token loops.
    """
    with _STREAM_REBIND_LOCK:
        default_stream = mx.new_stream(mx.default_device())
        mx.set_default_stream(default_stream)
        for module_name in module_names:
            try:
                module = importlib.import_module(module_name)
            except ImportError:
                continue
            if hasattr(module, "generation_stream"):
                setattr(module, "generation_stream", default_stream)
        return default_stream

Complete contract reference¶

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.mlx_streams.bind_generation_streams · function

vllm_mlx.mlx_streams.bind_generation_streams(module_names: Iterable[str] = ('mlx_lm.generate', 'mlx_vlm.generate')) -> object

Bind mlx-lm/mlx-vlm generation streams to the current thread.

Parameters

Name	Type	Required	Default	Description
`module_names`	`Iterable[str]`	`no`	`('mlx_lm.generate', 'mlx_vlm.generate')`	Optional positional or keyword input; defaults to `('mlx_lm.generate', 'mlx_vlm.generate')`.

Returns

Type: object
Direct return expressions: default_stream

Exceptions and behavior

Function bind_generation_streams calls mx.new_stream, mx.default_device, mx.set_default_stream, importlib.import_module; returns default_stream. No direct raise statement appears in this definition.

View source #L15-L39.

Complete symbol map¶

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
`bind_generation_streams`	function	`bind_generation_streams(module_names: Iterable[str] = ('mlx_lm.generate', 'mlx_vlm.generate')) -> object`	Bind mlx-lm/mlx-vlm generation streams to the current thread.	#L15-L39

vllm_mlx.mlx_streams¶