# `vllm_mlx.utils.mamba_cache`

BatchMambaCache implementation for continuous batching with Mamba models.

[View the complete module source at #L1-L215](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#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
    options:
      members:
        - logger
        - BatchMambaCache
        - patch_mlx_lm_for_mamba
        - _patched
        - ensure_mamba_support
      filters: []
      show_if_no_docstring: 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.

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.BatchMambaCache" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.BatchMambaCache</code> · class</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L24-L96).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.BatchMambaCache.__init__" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.BatchMambaCache.__init__</code> · method</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L32-L43).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.BatchMambaCache.extract" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.BatchMambaCache.extract</code> · method</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L45-L63).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.BatchMambaCache.merge" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.BatchMambaCache.merge</code> · method</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L66-L96).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba</code> · function</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L99-L194).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache</code> · nested function</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L126-L166).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache</code> · nested function</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L137-L155).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_merge_caches" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba._patched_merge_caches</code> · nested function</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L174-L190).

</details>

<details class="api-contract" id="contract-vllm_mlx.utils.mamba_cache.ensure_mamba_support" markdown="1">
<summary><code>vllm_mlx.utils.mamba_cache.ensure_mamba_support</code> · function</summary>

```python
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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L201-L215).

</details>

## 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`](#contract-vllm_mlx.utils.mamba_cache.BatchMambaCache) | class | `BatchMambaCache(left_padding: Optional[List[int]] = None, size: int = 2)` | Batch-aware MambaCache for continuous batching. | [#L24-L96](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L24-L96) |
| [`BatchMambaCache.__init__`](#contract-vllm_mlx.utils.mamba_cache.BatchMambaCache.__init__) | method | `BatchMambaCache.__init__(left_padding: Optional[List[int]] = None, size: int = 2) -> not annotated` | Initialize BatchMambaCache. | [#L32-L43](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L32-L43) |
| [`BatchMambaCache.extract`](#contract-vllm_mlx.utils.mamba_cache.BatchMambaCache.extract) | method | `BatchMambaCache.extract(idx: int) -> MambaCache` | Extract a single cache from the batch. | [#L45-L63](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L45-L63) |
| [`BatchMambaCache.merge`](#contract-vllm_mlx.utils.mamba_cache.BatchMambaCache.merge) | method | `BatchMambaCache.merge(caches: List[MambaCache]) -> 'BatchMambaCache'` | Merge multiple MambaCache objects into a BatchMambaCache. | [#L66-L96](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L66-L96) |
| [`patch_mlx_lm_for_mamba`](#contract-vllm_mlx.utils.mamba_cache.patch_mlx_lm_for_mamba) | function | `patch_mlx_lm_for_mamba() -> not annotated` | Patch mlx-lm to support MambaCache in BatchGenerator. | [#L99-L194](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L99-L194) |
| [`patch_mlx_lm_for_mamba._patched_make_cache`](#contract-vllm_mlx.utils.mamba_cache.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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L126-L166) |
| [`patch_mlx_lm_for_mamba._patched_make_cache.to_batch_cache`](#contract-vllm_mlx.utils.mamba_cache.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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L137-L155) |
| [`patch_mlx_lm_for_mamba._patched_merge_caches`](#contract-vllm_mlx.utils.mamba_cache.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](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L174-L190) |
| [`ensure_mamba_support`](#contract-vllm_mlx.utils.mamba_cache.ensure_mamba_support) | function | `ensure_mamba_support() -> not annotated` | Ensure MambaCache batching support is enabled. | [#L201-L215](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/utils/mamba_cache.py#L201-L215) |