vllm_mlx.utils.mamba_cache

BatchMambaCache implementation for continuous batching with Mamba models.

View the complete module source at #L1-L215.

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.utils.mamba_cache

BatchMambaCache implementation for continuous batching with Mamba models.

mlx-lm's BatchGenerator requires cache objects to have an extract method, but MambaCache (which extends ArraysCache) doesn't have one. This module provides a BatchMambaCache wrapper that adds batching support.

vllm_mlx.utils.mamba_cache.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.utils.mamba_cache._patched module-attribute

_patched = False

vllm_mlx.utils.mamba_cache.BatchMambaCache

BatchMambaCache(left_padding: Optional[List[int]] = None, size: int = 2)

Bases: ArraysCache

Batch-aware MambaCache for continuous batching.

This extends MambaCache to support batch operations required by mlx-lm's BatchGenerator, specifically the extract method.

Initialize BatchMambaCache.

Parameters:

• left_padding (Optional[List[int]], default: None ) –

  Amount of left padding for each sequence in batch

• size (int, default: 2 ) –

  Number of state arrays (default 2 for Mamba models)

Source code in vllm_mlx/utils/mamba_cache.py

def __init__(self, left_padding: Optional[List[int]] = None, size: int = 2):
    """
    Initialize BatchMambaCache.

    Args:
        left_padding: Amount of left padding for each sequence in batch
        size: Number of state arrays (default 2 for Mamba models)
    """
    # Always pass size - ArraysCache requires it, and MambaCache
    # (if it exists) inherits from ArraysCache
    super().__init__(size=size, left_padding=left_padding)
    self._batch_size = len(left_padding) if left_padding else 0

vllm_mlx.utils.mamba_cache.BatchMambaCache._batch_size instance-attribute

_batch_size = len(left_padding) if left_padding else 0

vllm_mlx.utils.mamba_cache.BatchMambaCache.extract

extract(idx: int) -> ArraysCache

Extract a single cache from the batch.

Parameters:

• idx (int) –

  Index of the sequence to extract

Returns:

• ArraysCache –

  A new MambaCache with the extracted state

Source code in vllm_mlx/utils/mamba_cache.py

def extract(self, idx: int) -> MambaCache:
    """
    Extract a single cache from the batch.

    Args:
        idx: Index of the sequence to extract

    Returns:
        A new MambaCache with the extracted state
    """
    size = len(self.cache)
    cache = MambaCache(size=size)
    # Extract the state arrays for this index
    cache.cache = [
        mx.contiguous(c[idx : idx + 1]) if c is not None else None
        for c in self.cache
    ]
    cache.left_padding = None  # Single sequence, no batch padding
    return cache

vllm_mlx.utils.mamba_cache.BatchMambaCache.merge classmethod

merge(caches: List[ArraysCache]) -> BatchMambaCache

Merge multiple MambaCache objects into a BatchMambaCache.

Parameters:

• caches (List[ArraysCache]) –

  List of MambaCache objects to merge

Returns:

• BatchMambaCache –

  A new BatchMambaCache containing all caches

Source code in vllm_mlx/utils/mamba_cache.py

@classmethod
def merge(cls, caches: List[MambaCache]) -> "BatchMambaCache":
    """
    Merge multiple MambaCache objects into a BatchMambaCache.

    Args:
        caches: List of MambaCache objects to merge

    Returns:
        A new BatchMambaCache containing all caches
    """
    if not caches:
        return cls([])

    # Get the structure from the first cache
    batch_size = len(caches)

    # MambaCache stores 2 arrays (size=2 in ArraysCache.__init__)
    merged_cache = cls([0] * batch_size)

    # Merge each array in the cache
    num_arrays = len(caches[0].cache)
    merged_cache.cache = []

    for i in range(num_arrays):
        arrays = [c.cache[i] for c in caches if c.cache[i] is not None]
        if arrays:
            merged_cache.cache.append(mx.concatenate(arrays, axis=0))
        else:
            merged_cache.cache.append(None)

    return merged_cache

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba

patch_mlx_lm_for_mamba()

Patch mlx-lm to support MambaCache in BatchGenerator.

This modifies the _make_cache function to handle MambaCache by converting it to BatchMambaCache.

Source code in vllm_mlx/utils/mamba_cache.py

def patch_mlx_lm_for_mamba():
    """
    Patch mlx-lm to support MambaCache in BatchGenerator.

    This modifies the _make_cache function to handle MambaCache by
    converting it to BatchMambaCache.
    """
    import importlib

    gen_module = importlib.import_module("mlx_lm.generate")
    from mlx_lm.models.cache import (
        KVCache,
        ArraysCache,
        RotatingKVCache,
        CacheList,
    )

    # MambaCache was removed in mlx-lm 0.30.6
    try:
        from mlx_lm.models.cache import MambaCache as OrigMambaCache
    except ImportError:
        OrigMambaCache = ArraysCache  # Fallback
    from mlx_lm.generate import BatchKVCache, BatchRotatingKVCache

    # Store original function
    _original_make_cache = gen_module._make_cache

    def _patched_make_cache(model, left_padding, max_kv_size=None):
        """
        Convert a list of regular caches into their corresponding
        batch-aware caches, with support for MambaCache.

        Args:
            model: The model to create cache for
            left_padding: Left padding for batch
            max_kv_size: Maximum KV cache size (mlx-lm 0.30.6+)
        """

        def to_batch_cache(c):
            if isinstance(c, KVCache):
                return BatchKVCache(left_padding)
            elif isinstance(c, OrigMambaCache):
                # Handle MambaCache -> BatchMambaCache
                return BatchMambaCache(left_padding)
            elif isinstance(c, ArraysCache):
                c.left_padding = mx.array(left_padding)
                return c
            elif isinstance(c, RotatingKVCache):
                if c.keep > 0:
                    raise ValueError(
                        "RotatingKVCache with keep tokens is not supported."
                    )
                return BatchRotatingKVCache(c.max_size, left_padding)
            elif isinstance(c, CacheList):
                return CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))
            else:
                raise ValueError(f"{type(c)} does not yet support batching")

        if hasattr(model, "make_cache"):
            cache = model.make_cache()
            return [to_batch_cache(c) for c in cache]
        elif max_kv_size is not None:
            # mlx-lm 0.30.6+: Use rotating cache with max_kv_size
            return [
                BatchRotatingKVCache(max_kv_size, left_padding) for _ in model.layers
            ]
        else:
            return [BatchKVCache(left_padding) for _ in model.layers]

    # Patch the module
    gen_module._make_cache = _patched_make_cache

    # Also patch _merge_caches to handle BatchMambaCache
    _original_merge_caches = gen_module._merge_caches

    def _patched_merge_caches(caches):
        """Merge caches with MambaCache support."""
        batch_cache = []
        for i in range(len(caches[0])):
            cache = None
            if isinstance(caches[0][i], KVCache):
                cache = BatchKVCache.merge([c[i] for c in caches])
            elif isinstance(caches[0][i], RotatingKVCache):
                cache = BatchRotatingKVCache.merge([c[i] for c in caches])
            elif isinstance(caches[0][i], (OrigMambaCache, BatchMambaCache)):
                cache = BatchMambaCache.merge([c[i] for c in caches])
            else:
                raise ValueError(
                    f"{type(caches[0][i])} does not yet support batching with history"
                )
            batch_cache.append(cache)
        return batch_cache

    gen_module._merge_caches = _patched_merge_caches

    logger.info("Patched mlx-lm for MambaCache batching support")

vllm_mlx.utils.mamba_cache.ensure_mamba_support

ensure_mamba_support()

Ensure MambaCache batching support is enabled.

NOTE: Disabled for mlx-lm >= 0.30.6 where ArraysCache natively supports all batch operations (extract, merge, filter, prepare). The old patch replaced ArraysCache with BatchMambaCache, which broke hybrid models (Qwen3.5) that mix ArraysCache + KVCache layers.

Source code in vllm_mlx/utils/mamba_cache.py

def ensure_mamba_support():
    """Ensure MambaCache batching support is enabled.

    NOTE: Disabled for mlx-lm >= 0.30.6 where ArraysCache natively supports
    all batch operations (extract, merge, filter, prepare).  The old patch
    replaced ArraysCache with BatchMambaCache, which broke hybrid models
    (Qwen3.5) that mix ArraysCache + KVCache layers.
    """
    global _patched
    if not _patched:
        logger.info(
            "[MambaCache] Skipping _make_cache patch — "
            "mlx-lm ArraysCache has native batching support"
        )
        _patched = True

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.utils.mamba_cache.BatchMambaCache · class

vllm_mlx.utils.mamba_cache.BatchMambaCache(left_padding: Optional[List[int]] = None, size: int = 2)

Batch-aware MambaCache for continuous batching.

Parameters

Name	Type	Required	Default	Description
left_padding	Optional[List[int]]	no	None	Amount of left padding for each sequence in batch
size	int	no	2	Number of state arrays (default 2 for Mamba models)

Returns

• Constructs: vllm_mlx.utils.mamba_cache.BatchMambaCache

Exceptions and behavior

Class BatchMambaCache derives from MambaCache and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L24-L96.

vllm_mlx.utils.mamba_cache.BatchMambaCache.__init__ · method

vllm_mlx.utils.mamba_cache.BatchMambaCache.__init__(left_padding: Optional[List[int]] = None, size: int = 2) -> not annotated

Initialize BatchMambaCache.

Parameters

Name	Type	Required	Default	Description
left_padding	Optional[List[int]]	no	None	Amount of left padding for each sequence in batch
size	int	no	2	Number of state arrays (default 2 for Mamba models)

Returns

• Type: not annotated

Exceptions and behavior

Method BatchMambaCache.__init__ updates self._batch_size; calls super().__init__, super, len. No direct raise statement appears in this definition.

View source #L32-L43.

vllm_mlx.utils.mamba_cache.BatchMambaCache.extract · method

vllm_mlx.utils.mamba_cache.BatchMambaCache.extract(idx: int) -> MambaCache

Extract a single cache from the batch.

Parameters

Name	Type	Required	Default	Description
idx	int	yes	none	Index of the sequence to extract

Returns

• Type: MambaCache
• Direct return expressions: cache

Exceptions and behavior

Method BatchMambaCache.extract calls len, MambaCache, mx.contiguous; returns cache. No direct raise statement appears in this definition.

View source #L45-L63.

vllm_mlx.utils.mamba_cache.BatchMambaCache.merge · method

vllm_mlx.utils.mamba_cache.BatchMambaCache.merge(caches: List[MambaCache]) -> 'BatchMambaCache'

Merge multiple MambaCache objects into a BatchMambaCache.

Parameters

Name	Type	Required	Default	Description
caches	List[MambaCache]	yes	none	List of MambaCache objects to merge

Returns

• Type: 'BatchMambaCache'
• Direct return expressions: cls([]); merged_cache

Exceptions and behavior

Method BatchMambaCache.merge calls cls, len, range, merged_cache.cache.append; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L66-L96.

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba · function

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba() -> not annotated

Patch mlx-lm to support MambaCache in BatchGenerator.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Function patch_mlx_lm_for_mamba calls importlib.import_module, logger.info. No direct raise statement appears in this definition.

View source #L99-L194.

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache · nested function

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache(model, left_padding, max_kv_size = None) -> not annotated

Convert a list of regular caches into their corresponding batch-aware caches, with support for MambaCache.

Parameters

Name	Type	Required	Default	Description
model	not annotated	yes	none	The model to create cache for
left_padding	not annotated	yes	none	Left padding for batch
max_kv_size	not annotated	no	None	Maximum KV cache size (mlx-lm 0.30.6+)

Returns

• Type: not annotated
• Direct return expressions: [to_batch_cache(c) for c in cache]; [BatchRotatingKVCache(max_kv_size, left_padding) for _ in model.layers]; [BatchKVCache(left_padding) for _ in model.layers]

Exceptions and behavior

Nested Function patch_mlx_lm_for_mamba._patched_make_cache calls hasattr, model.make_cache, to_batch_cache, BatchRotatingKVCache; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L126-L166.

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache · nested function

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache(c) -> not annotated

Nested Function patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache calls isinstance, BatchKVCache, BatchMambaCache, mx.array; can raise ValueError; has 5 explicit return paths.

Parameters

Name	Type	Required	Default	Description
c	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: BatchKVCache(left_padding); BatchMambaCache(left_padding); c; BatchRotatingKVCache(c.max_size, left_padding); CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))

Exceptions and behavior

Nested Function patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache calls isinstance, BatchKVCache, BatchMambaCache, mx.array; can raise ValueError; has 5 explicit return paths. Directly raised exceptions: ValueError.

View source #L137-L155.

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_merge_caches · nested function

vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_merge_caches(caches) -> not annotated

Merge caches with MambaCache support.

Parameters

Name	Type	Required	Default	Description
caches	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: batch_cache

Exceptions and behavior

Nested Function patch_mlx_lm_for_mamba._patched_merge_caches calls range, len, isinstance, BatchKVCache.merge; can raise ValueError; returns batch_cache. Directly raised exceptions: ValueError.

View source #L174-L190.

vllm_mlx.utils.mamba_cache.ensure_mamba_support · function

vllm_mlx.utils.mamba_cache.ensure_mamba_support() -> not annotated

Ensure MambaCache batching support is enabled.

Parameters

This callable has no explicit inputs.

Returns

• Type: not annotated

Exceptions and behavior

Function ensure_mamba_support calls logger.info. No direct raise statement appears in this definition.

View source #L201-L215.

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
BatchMambaCache	class	BatchMambaCache(left_padding: Optional[List[int]] = None, size: int = 2)	Batch-aware MambaCache for continuous batching.	#L24-L96
BatchMambaCache.__init__	method	BatchMambaCache.__init__(left_padding: Optional[List[int]] = None, size: int = 2) -> not annotated	Initialize BatchMambaCache.	#L32-L43
BatchMambaCache.extract	method	BatchMambaCache.extract(idx: int) -> MambaCache	Extract a single cache from the batch.	#L45-L63
BatchMambaCache.merge	method	BatchMambaCache.merge(caches: List[MambaCache]) -> 'BatchMambaCache'	Merge multiple MambaCache objects into a BatchMambaCache.	#L66-L96
patch_mlx_lm_for_mamba	function	patch_mlx_lm_for_mamba() -> not annotated	Patch mlx-lm to support MambaCache in BatchGenerator.	#L99-L194
patch_mlx_lm_for_mamba._patched_make_cache	nested function	patch_mlx_lm_for_mamba._patched_make_cache(model, left_padding, max_kv_size = None) -> not annotated	Convert a list of regular caches into their corresponding batch-aware caches, with support for MambaCache.	#L126-L166
patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache	nested function	patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache(c) -> not annotated	Nested Function patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache calls isinstance, BatchKVCache, BatchMambaCache, mx.array; can raise ValueError; has 5 explicit return paths.	#L137-L155
patch_mlx_lm_for_mamba._patched_merge_caches	nested function	patch_mlx_lm_for_mamba._patched_merge_caches(caches) -> not annotated	Merge caches with MambaCache support.	#L174-L190
ensure_mamba_support	function	ensure_mamba_support() -> not annotated	Ensure MambaCache batching support is enabled.	#L201-L215