# `vllm_mlx.specprefill`

SpecPrefill: Attention-based sparse prefill for MLX.

[View the complete module source at #L1-L845](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L1-L845).

## 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.specprefill
    options:
      members:
        - _AttentionCapture
        - _qwen35_extract_queries
        - _llama_extract_queries
        - _nemotron_h_extract_queries
        - _patch_attention_for_capture
        - _unpatch_attention_capture
        - _prefill_draft
        - _lookahead_decode
        - _avg_pool1d
        - _compute_importance
        - score_tokens
        - select_chunks
        - manual_rope
        - manual_rope_with_freqs
        - _PositionMappedRoPE
        - _OffsetAdjustedRoPE
        - _get_dims
        - _get_pre_scale
        - _find_attention_layers
        - _get_attn_module
        - _get_rope
        - _set_rope
        - _set_attn_module
        - _build_layer_to_cache_map
        - sparse_prefill
        - cleanup_rope
      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.specprefill._AttentionCapture" markdown="1">
<summary><code>vllm_mlx.specprefill._AttentionCapture</code> · class</summary>

```python
vllm_mlx.specprefill._AttentionCapture(original, buf_idx, query_buffer, query_extractor = None)
```

Wrapper that captures post-RoPE query vectors and delegates to original.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `original` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `buf_idx` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `query_buffer` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `query_extractor` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Constructs: `vllm_mlx.specprefill._AttentionCapture`

**Exceptions and behavior**

Class `_AttentionCapture` declares 3 direct member(s).
No direct `raise` statement appears in this definition.

[View source #L53-L73](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L53-L73).

</details>

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

```python
vllm_mlx.specprefill._AttentionCapture.__init__(original, buf_idx, query_buffer, query_extractor = None) -> not annotated
```

Method `_AttentionCapture.__init__` updates `self._original`, `self._buf_idx`, `self._query_buffer`, `self._query_extractor`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `original` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `buf_idx` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `query_buffer` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `query_extractor` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Method `_AttentionCapture.__init__` updates `self._original`, `self._buf_idx`, `self._query_buffer`, `self._query_extractor`.
No direct `raise` statement appears in this definition.

[View source #L61-L65](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L61-L65).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._AttentionCapture.__call__" markdown="1">
<summary><code>vllm_mlx.specprefill._AttentionCapture.__call__</code> · method</summary>

```python
vllm_mlx.specprefill._AttentionCapture.__call__(x, mask = None, cache = None) -> not annotated
```

Method `_AttentionCapture.__call__` calls `self._query_extractor`, `self._query_buffer[self._buf_idx].append`, `self._original`; returns `self._original(x, mask=mask, cache=cache)`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `mask` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |
| `cache` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `self._original(x, mask=mask, cache=cache)`

**Exceptions and behavior**

Method `_AttentionCapture.__call__` calls `self._query_extractor`, `self._query_buffer[self._buf_idx].append`, `self._original`; returns `self._original(x, mask=mask, cache=cache)`.
No direct `raise` statement appears in this definition.

[View source #L67-L70](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L67-L70).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._AttentionCapture.__getattr__" markdown="1">
<summary><code>vllm_mlx.specprefill._AttentionCapture.__getattr__</code> · method</summary>

```python
vllm_mlx.specprefill._AttentionCapture.__getattr__(name) -> not annotated
```

Method `_AttentionCapture.__getattr__` calls `getattr`; returns `getattr(self._original, name)`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `name` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `getattr(self._original, name)`

**Exceptions and behavior**

Method `_AttentionCapture.__getattr__` calls `getattr`; returns `getattr(self._original, name)`.
No direct `raise` statement appears in this definition.

[View source #L72-L73](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L72-L73).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._qwen35_extract_queries" markdown="1">
<summary><code>vllm_mlx.specprefill._qwen35_extract_queries</code> · function</summary>

```python
vllm_mlx.specprefill._qwen35_extract_queries(attn, x, cache = None) -> not annotated
```

Extract post-RoPE queries from Qwen3.5 attention (gate split + q_norm).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `attn` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `queries`

**Exceptions and behavior**

Function `_qwen35_extract_queries` calls `attn.q_proj`, `mx.split`, `q_out.reshape`, `attn.q_norm(queries).transpose`; returns `queries`.
No direct `raise` statement appears in this definition.

[View source #L76-L92](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L76-L92).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._llama_extract_queries" markdown="1">
<summary><code>vllm_mlx.specprefill._llama_extract_queries</code> · function</summary>

```python
vllm_mlx.specprefill._llama_extract_queries(attn, x, cache = None) -> not annotated
```

Extract post-RoPE queries from standard transformer attention.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `attn` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `queries`

**Exceptions and behavior**

Function `_llama_extract_queries` calls `getattr`, `attn.q_proj`, `queries.reshape(B, L, n_heads, -1).transpose`, `queries.reshape`; returns `queries`.
No direct `raise` statement appears in this definition.

[View source #L95-L113](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L95-L113).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._nemotron_h_extract_queries" markdown="1">
<summary><code>vllm_mlx.specprefill._nemotron_h_extract_queries</code> · function</summary>

```python
vllm_mlx.specprefill._nemotron_h_extract_queries(attn, x, cache = None) -> not annotated
```

Extract queries from Nemotron-H attention (no RoPE, no gate, no q_norm).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `attn` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `queries`

**Exceptions and behavior**

Function `_nemotron_h_extract_queries` calls `attn.q_proj(x).reshape(B, L, attn.num_heads, -1).transpose`, `attn.q_proj(x).reshape`, `attn.q_proj`; returns `queries`.
No direct `raise` statement appears in this definition.

[View source #L116-L125](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L116-L125).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._patch_attention_for_capture" markdown="1">
<summary><code>vllm_mlx.specprefill._patch_attention_for_capture</code> · function</summary>

```python
vllm_mlx.specprefill._patch_attention_for_capture(model, query_buffer, query_extractor = None) -> not annotated
```

Replace attention modules on full-attention layers with capture wrappers.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `query_buffer` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `query_extractor` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `(originals, attn_indices)`

**Exceptions and behavior**

Function `_patch_attention_for_capture` calls `_find_attention_layers`, `len`, `attn_indices.append`, `_get_attn_module`; returns `(originals, attn_indices)`.
No direct `raise` statement appears in this definition.

[View source #L128-L149](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L128-L149).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._unpatch_attention_capture" markdown="1">
<summary><code>vllm_mlx.specprefill._unpatch_attention_capture</code> · function</summary>

```python
vllm_mlx.specprefill._unpatch_attention_capture(model, originals) -> not annotated
```

Restore original attention modules after capture.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `originals` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Function `_unpatch_attention_capture` calls `_set_attn_module`.
No direct `raise` statement appears in this definition.

[View source #L152-L155](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L152-L155).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._prefill_draft" markdown="1">
<summary><code>vllm_mlx.specprefill._prefill_draft</code> · function</summary>

```python
vllm_mlx.specprefill._prefill_draft(model, tokens, cache, step_size = 2048, cancel_check = None) -> not annotated
```

Prefill prompt tokens into cache.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `tokens` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `step_size` | `not annotated` | `no` | `2048` | Optional positional or keyword input; defaults to `2048`. |
| `cancel_check` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `logits`

**Exceptions and behavior**

Function `_prefill_draft` calls `isinstance`, `mx.array`, `len`, `cancel_check`; returns `logits`.
No direct `raise` statement appears in this definition.

[View source #L158-L175](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L158-L175).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._lookahead_decode" markdown="1">
<summary><code>vllm_mlx.specprefill._lookahead_decode</code> · function</summary>

```python
vllm_mlx.specprefill._lookahead_decode(model, first_logits, cache, n_steps, temp = 0.6, top_p = 0.95, cancel_check = None) -> not annotated
```

Run n_steps autoregressive decode, returning generated token ids.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `first_logits` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `n_steps` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `temp` | `not annotated` | `no` | `0.6` | Optional positional or keyword input; defaults to `0.6`. |
| `top_p` | `not annotated` | `no` | `0.95` | Optional positional or keyword input; defaults to `0.95`. |
| `cancel_check` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `generated`

**Exceptions and behavior**

Function `_lookahead_decode` calls `make_sampler`, `cancel_check`, `sampler`, `mx.eval`; returns `generated`.
No direct `raise` statement appears in this definition.

[View source #L178-L204](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L178-L204).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._avg_pool1d" markdown="1">
<summary><code>vllm_mlx.specprefill._avg_pool1d</code> · function</summary>

```python
vllm_mlx.specprefill._avg_pool1d(x, kernel_size) -> not annotated
```

1D average pooling along last axis via prefix-sum.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `not annotated` | `yes` | `none` | (..., M) input |
| `kernel_size` | `not annotated` | `yes` | `none` | window size (odd for centered) |

**Returns**

- Type: `not annotated`
- Direct return expressions: `x`; `(prefix[..., kernel_size:] - prefix[..., :-kernel_size]) / kernel_size`

**Exceptions and behavior**

Function `_avg_pool1d` calls `mx.pad`, `mx.zeros`, `mx.concatenate`, `mx.cumsum`; has 2 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L207-L223](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L207-L223).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._compute_importance" markdown="1">
<summary><code>vllm_mlx.specprefill._compute_importance</code> · function</summary>

```python
vllm_mlx.specprefill._compute_importance(query_buffer, attn_caches, n_prompt, n_attn_heads, n_kv_heads, pool_kernel = 13) -> not annotated
```

Compute per-token importance from captured queries and cached keys.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `query_buffer` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `attn_caches` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `n_prompt` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `n_attn_heads` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `n_kv_heads` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `pool_kernel` | `not annotated` | `no` | `13` | Optional positional or keyword input; defaults to `13`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `importance`

**Exceptions and behavior**

Function `_compute_importance` calls `enumerate`, `mx.concatenate`, `mx.repeat`, `expanded_keys.transpose`; can raise `RuntimeError`; returns `importance`.
Directly raised exceptions: `RuntimeError`.

[View source #L226-L271](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L226-L271).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.score_tokens" markdown="1">
<summary><code>vllm_mlx.specprefill.score_tokens</code> · function</summary>

```python
vllm_mlx.specprefill.score_tokens(model, tokens, n_lookahead = 8, pool_kernel = 13, temp = 0.6, top_p = 0.95, prefill_step_size = 2048, query_extractor = None, cancel_check = None) -> not annotated
```

Score token importance using attention-based analysis on a draft model.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Draft model (small, fast — e.g. 4B) |
| `tokens` | `not annotated` | `yes` | `none` | list or mx.array of token IDs |
| `n_lookahead` | `not annotated` | `no` | `8` | decode steps for query capture (default 8) |
| `pool_kernel` | `not annotated` | `no` | `13` | smoothing kernel for avg_pool1d (default 13, 0=disable) |
| `temp` | `not annotated` | `no` | `0.6` | sampling temperature for lookahead (default 0.6) |
| `top_p` | `not annotated` | `no` | `0.95` | top-p for lookahead (default 0.95) |
| `prefill_step_size` | `not annotated` | `no` | `2048` | chunk size for draft prefill (default 2048) |
| `query_extractor` | `not annotated` | `no` | `None` | function(attn, x, cache) → queries tensor. |
| `cancel_check` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `importance`

**Exceptions and behavior**

Function `score_tokens` calls `isinstance`, `tokens.tolist`, `len`, `_find_attention_layers`; returns `importance`.
No direct `raise` statement appears in this definition.

[View source #L274-L396](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L274-L396).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.select_chunks" markdown="1">
<summary><code>vllm_mlx.specprefill.select_chunks</code> · function</summary>

```python
vllm_mlx.specprefill.select_chunks(importance, keep_pct = 0.3, chunk_size = 32, backbone_pct = 0.0) -> not annotated
```

Select top-k% token chunks by average importance.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `importance` | `not annotated` | `yes` | `none` | (M,) per-token importance scores |
| `keep_pct` | `not annotated` | `no` | `0.3` | fraction of chunks to keep (default 0.3) |
| `chunk_size` | `not annotated` | `no` | `32` | tokens per chunk (default 32) |
| `backbone_pct` | `not annotated` | `no` | `0.0` | fraction of chunks reserved for evenly-spaced coverage |

**Returns**

- Type: `not annotated`
- Direct return expressions: `mx.arange(M)`; `mx.array(indices)`

**Exceptions and behavior**

Function `select_chunks` calls `mx.arange`, `math.ceil`, `max`, `range`; has 2 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L399-L467](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L399-L467).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.select_chunks._selected_token_count" markdown="1">
<summary><code>vllm_mlx.specprefill.select_chunks._selected_token_count</code> · nested function</summary>

```python
vllm_mlx.specprefill.select_chunks._selected_token_count(chunks) -> not annotated
```

Nested Function `select_chunks._selected_token_count` calls `min`; returns `total`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `chunks` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `total`

**Exceptions and behavior**

Nested Function `select_chunks._selected_token_count` calls `min`; returns `total`.
No direct `raise` statement appears in this definition.

[View source #L437-L443](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L437-L443).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.manual_rope" markdown="1">
<summary><code>vllm_mlx.specprefill.manual_rope</code> · function</summary>

```python
vllm_mlx.specprefill.manual_rope(x, positions, dims, base = 10000.0, scale = 1.0) -> not annotated
```

Apply RoPE at arbitrary (non-contiguous) positions.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `not annotated` | `yes` | `none` | (B, n_heads, L, head_dim) input tensor |
| `positions` | `not annotated` | `yes` | `none` | (L,) position indices (can be non-contiguous) |
| `dims` | `not annotated` | `yes` | `none` | number of dimensions to rotate (head_dim * partial_rotary_factor) |
| `base` | `not annotated` | `no` | `10000.0` | RoPE base frequency (default 10000.0) |
| `scale` | `not annotated` | `no` | `1.0` | position scale divisor (default 1.0, higher = compressed positions) |

**Returns**

- Type: `not annotated`
- Direct return expressions: `mx.concatenate([rotated, x_pass], axis=-1)`

**Exceptions and behavior**

Function `manual_rope` calls `mx.arange`, `positions.astype`, `mx.cos`, `mx.sin`; returns `mx.concatenate([rotated, x_pass], axis=-1)`.
No direct `raise` statement appears in this definition.

[View source #L480-L508](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L480-L508).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.manual_rope_with_freqs" markdown="1">
<summary><code>vllm_mlx.specprefill.manual_rope_with_freqs</code> · function</summary>

```python
vllm_mlx.specprefill.manual_rope_with_freqs(x, positions, dims, freqs, pre_scale = 1.0) -> not annotated
```

Apply RoPE at arbitrary positions using pre-computed frequencies.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `positions` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `dims` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `freqs` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `pre_scale` | `not annotated` | `no` | `1.0` | Optional positional or keyword input; defaults to `1.0`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `mx.concatenate([rotated, x_pass], axis=-1)`

**Exceptions and behavior**

Function `manual_rope_with_freqs` calls `(1.0 / freqs).astype`, `positions[:, None].astype`, `mx.cos`, `mx.sin`; returns `mx.concatenate([rotated, x_pass], axis=-1)`.
No direct `raise` statement appears in this definition.

[View source #L511-L528](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L511-L528).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._PositionMappedRoPE" markdown="1">
<summary><code>vllm_mlx.specprefill._PositionMappedRoPE</code> · class</summary>

```python
vllm_mlx.specprefill._PositionMappedRoPE(original_rope, all_positions, cache_start = 0)
```

Wraps a RoPE module to apply rotation at non-contiguous positions.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `original_rope` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `all_positions` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache_start` | `not annotated` | `no` | `0` | Optional positional or keyword input; defaults to `0`. |

**Returns**

- Constructs: `vllm_mlx.specprefill._PositionMappedRoPE`

**Exceptions and behavior**

Class `_PositionMappedRoPE` declares 2 direct member(s).
No direct `raise` statement appears in this definition.

[View source #L536-L571](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L536-L571).

</details>

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

```python
vllm_mlx.specprefill._PositionMappedRoPE.__init__(original_rope, all_positions, cache_start = 0) -> not annotated
```

Method `_PositionMappedRoPE.__init__` updates `self._original`, `self._all_positions`, `self._cache_start`, `self._has_custom_freqs`; calls `hasattr`, `_get_dims`, `_get_pre_scale`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `original_rope` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `all_positions` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `cache_start` | `not annotated` | `no` | `0` | Optional positional or keyword input; defaults to `0`. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Method `_PositionMappedRoPE.__init__` updates `self._original`, `self._all_positions`, `self._cache_start`, `self._has_custom_freqs`; calls `hasattr`, `_get_dims`, `_get_pre_scale`.
No direct `raise` statement appears in this definition.

[View source #L547-L561](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L547-L561).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._PositionMappedRoPE.__call__" markdown="1">
<summary><code>vllm_mlx.specprefill._PositionMappedRoPE.__call__</code> · method</summary>

```python
vllm_mlx.specprefill._PositionMappedRoPE.__call__(x, offset = 0) -> not annotated
```

Method `_PositionMappedRoPE.__call__` calls `manual_rope_with_freqs`, `manual_rope`; has 2 explicit return paths.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `offset` | `not annotated` | `no` | `0` | Optional positional or keyword input; defaults to `0`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `manual_rope_with_freqs(x, positions, self._dims, self._freqs, pre_scale=self._pre_scale)`; `manual_rope(x, positions, self._dims, base=self._base, scale=self._scale)`

**Exceptions and behavior**

Method `_PositionMappedRoPE.__call__` calls `manual_rope_with_freqs`, `manual_rope`; has 2 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L563-L571](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L563-L571).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._OffsetAdjustedRoPE" markdown="1">
<summary><code>vllm_mlx.specprefill._OffsetAdjustedRoPE</code> · class</summary>

```python
vllm_mlx.specprefill._OffsetAdjustedRoPE(original_rope, adjustment)
```

Wraps a RoPE module to add a constant offset for decode after sparse prefill.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `original_rope` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `adjustment` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Constructs: `vllm_mlx.specprefill._OffsetAdjustedRoPE`

**Exceptions and behavior**

Class `_OffsetAdjustedRoPE` declares 2 direct member(s).
No direct `raise` statement appears in this definition.

[View source #L574-L590](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L574-L590).

</details>

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

```python
vllm_mlx.specprefill._OffsetAdjustedRoPE.__init__(original_rope, adjustment) -> not annotated
```

Method `_OffsetAdjustedRoPE.__init__` updates `self._original`, `self._adjustment`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `original_rope` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `adjustment` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Method `_OffsetAdjustedRoPE.__init__` updates `self._original`, `self._adjustment`.
No direct `raise` statement appears in this definition.

[View source #L585-L587](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L585-L587).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._OffsetAdjustedRoPE.__call__" markdown="1">
<summary><code>vllm_mlx.specprefill._OffsetAdjustedRoPE.__call__</code> · method</summary>

```python
vllm_mlx.specprefill._OffsetAdjustedRoPE.__call__(x, offset = 0) -> not annotated
```

Method `_OffsetAdjustedRoPE.__call__` calls `self._original`; returns `self._original(x, offset=offset + self._adjustment)`.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `x` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `offset` | `not annotated` | `no` | `0` | Optional positional or keyword input; defaults to `0`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `self._original(x, offset=offset + self._adjustment)`

**Exceptions and behavior**

Method `_OffsetAdjustedRoPE.__call__` calls `self._original`; returns `self._original(x, offset=offset + self._adjustment)`.
No direct `raise` statement appears in this definition.

[View source #L589-L590](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L589-L590).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._get_dims" markdown="1">
<summary><code>vllm_mlx.specprefill._get_dims</code> · function</summary>

```python
vllm_mlx.specprefill._get_dims(rope_module) -> not annotated
```

Extract rotary dimensions from any RoPE variant.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `rope_module` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `getattr(rope_module, attr)`

**Exceptions and behavior**

Function `_get_dims` calls `hasattr`, `getattr`, `ValueError`, `type`; can raise `ValueError`; returns `getattr(rope_module, attr)`.
Directly raised exceptions: `ValueError`.

[View source #L598-L603](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L598-L603).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._get_pre_scale" markdown="1">
<summary><code>vllm_mlx.specprefill._get_pre_scale</code> · function</summary>

```python
vllm_mlx.specprefill._get_pre_scale(rope_module) -> not annotated
```

Extract pre-scale factor from custom RoPE variants (SuScaled, Yarn).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `rope_module` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `rope_module.mscale`; `rope_module._scale`; `1.0`

**Exceptions and behavior**

Function `_get_pre_scale` calls `hasattr`; has 3 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L606-L612](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L606-L612).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._find_attention_layers" markdown="1">
<summary><code>vllm_mlx.specprefill._find_attention_layers</code> · function</summary>

```python
vllm_mlx.specprefill._find_attention_layers(model) -> not annotated
```

Find all full-attention layers across architectures.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `results`

**Exceptions and behavior**

Function `_find_attention_layers` calls `enumerate`, `hasattr`, `results.append`, `getattr`; returns `results`.
No direct `raise` statement appears in this definition.

[View source #L615-L630](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L615-L630).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._get_attn_module" markdown="1">
<summary><code>vllm_mlx.specprefill._get_attn_module</code> · function</summary>

```python
vllm_mlx.specprefill._get_attn_module(layer) -> not annotated
```

Get the attention module from a layer (self_attn or mixer).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `layer` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `layer.self_attn`; `layer.mixer`; `None`

**Exceptions and behavior**

Function `_get_attn_module` calls `hasattr`, `getattr`; has 3 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L633-L639](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L633-L639).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._get_rope" markdown="1">
<summary><code>vllm_mlx.specprefill._get_rope</code> · function</summary>

```python
vllm_mlx.specprefill._get_rope(attn) -> not annotated
```

Get the RoPE module from an attention layer, or None.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `attn` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `getattr(attn, 'rope', None) or getattr(attn, 'rotary_emb', None)`

**Exceptions and behavior**

Function `_get_rope` calls `getattr`; returns `getattr(attn, 'rope', None) or getattr(attn, 'rotary_emb', None)`.
No direct `raise` statement appears in this definition.

[View source #L642-L647](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L642-L647).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._set_rope" markdown="1">
<summary><code>vllm_mlx.specprefill._set_rope</code> · function</summary>

```python
vllm_mlx.specprefill._set_rope(attn, rope_module) -> not annotated
```

Set the RoPE module on an attention layer.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `attn` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `rope_module` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Function `_set_rope` calls `hasattr`.
No direct `raise` statement appears in this definition.

[View source #L650-L655](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L650-L655).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._set_attn_module" markdown="1">
<summary><code>vllm_mlx.specprefill._set_attn_module</code> · function</summary>

```python
vllm_mlx.specprefill._set_attn_module(layer, module) -> not annotated
```

Set the attention module on a layer (self_attn or mixer).

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `layer` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `module` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Function `_set_attn_module` calls `hasattr`, `getattr`.
No direct `raise` statement appears in this definition.

[View source #L658-L663](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L658-L663).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill._build_layer_to_cache_map" markdown="1">
<summary><code>vllm_mlx.specprefill._build_layer_to_cache_map</code> · function</summary>

```python
vllm_mlx.specprefill._build_layer_to_cache_map(model) -> not annotated
```

Build mapping from model layer index to cache index.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `{i: i for i in range(len(model.layers))}`; `layer_to_cache`

**Exceptions and behavior**

Function `_build_layer_to_cache_map` calls `any`, `hasattr`, `range`, `len`; has 2 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L666-L690](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L666-L690).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.sparse_prefill" markdown="1">
<summary><code>vllm_mlx.specprefill.sparse_prefill</code> · function</summary>

```python
vllm_mlx.specprefill.sparse_prefill(model, tokens, selected_indices, cache, step_size = 2048, position_offset = 0, cancel_check = None) -> not annotated
```

Prefill the model cache with selected tokens at their original positions.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Language model with .layers property (TextModel or VLM Model) |
| `tokens` | `not annotated` | `yes` | `none` | (M,) all prompt token IDs (mx.array or list) |
| `selected_indices` | `not annotated` | `yes` | `none` | (N,) sorted indices into tokens to keep (mx.array or list) |
| `cache` | `not annotated` | `yes` | `none` | list of KVCache/ArraysCache from make_prompt_cache() |
| `step_size` | `not annotated` | `no` | `2048` | chunk size for processing (default 2048) |
| `position_offset` | `not annotated` | `no` | `0` | added to selected_indices for RoPE positions (default 0). Use when the cache already has tokens from a prior prefill (e.g., system prompt KV cache with S tokens → position_offset=S). |
| `cancel_check` | `not annotated` | `no` | `None` | Optional positional or keyword input; defaults to `None`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `logits`

**Exceptions and behavior**

Function `sparse_prefill` calls `isinstance`, `mx.array`, `type`, `max`; returns `logits`.
No direct `raise` statement appears in this definition.

[View source #L698-L827](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L698-L827).

</details>

<details class="api-contract" id="contract-vllm_mlx.specprefill.cleanup_rope" markdown="1">
<summary><code>vllm_mlx.specprefill.cleanup_rope</code> · function</summary>

```python
vllm_mlx.specprefill.cleanup_rope(model) -> not annotated
```

Restore original RoPE on all attention layers.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Function `cleanup_rope` calls `_find_attention_layers`, `_get_attn_module`, `_get_rope`, `isinstance`.
No direct `raise` statement appears in this definition.

[View source #L830-L845](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L830-L845).

</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 |
| --- | --- | --- | --- | --- |
| [`_AttentionCapture`](#contract-vllm_mlx.specprefill._AttentionCapture) | class | `_AttentionCapture(original, buf_idx, query_buffer, query_extractor = None)` | Wrapper that captures post-RoPE query vectors and delegates to original. | [#L53-L73](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L53-L73) |
| [`_AttentionCapture.__init__`](#contract-vllm_mlx.specprefill._AttentionCapture.__init__) | method | `_AttentionCapture.__init__(original, buf_idx, query_buffer, query_extractor = None) -> not annotated` | Method `_AttentionCapture.__init__` updates `self._original`, `self._buf_idx`, `self._query_buffer`, `self._query_extractor`. | [#L61-L65](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L61-L65) |
| [`_AttentionCapture.__call__`](#contract-vllm_mlx.specprefill._AttentionCapture.__call__) | method | `_AttentionCapture.__call__(x, mask = None, cache = None) -> not annotated` | Method `_AttentionCapture.__call__` calls `self._query_extractor`, `self._query_buffer[self._buf_idx].append`, `self._original`; returns `self._original(x, mask=mask, cache=cache)`. | [#L67-L70](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L67-L70) |
| [`_AttentionCapture.__getattr__`](#contract-vllm_mlx.specprefill._AttentionCapture.__getattr__) | method | `_AttentionCapture.__getattr__(name) -> not annotated` | Method `_AttentionCapture.__getattr__` calls `getattr`; returns `getattr(self._original, name)`. | [#L72-L73](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L72-L73) |
| [`_qwen35_extract_queries`](#contract-vllm_mlx.specprefill._qwen35_extract_queries) | function | `_qwen35_extract_queries(attn, x, cache = None) -> not annotated` | Extract post-RoPE queries from Qwen3.5 attention (gate split + q_norm). | [#L76-L92](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L76-L92) |
| [`_llama_extract_queries`](#contract-vllm_mlx.specprefill._llama_extract_queries) | function | `_llama_extract_queries(attn, x, cache = None) -> not annotated` | Extract post-RoPE queries from standard transformer attention. | [#L95-L113](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L95-L113) |
| [`_nemotron_h_extract_queries`](#contract-vllm_mlx.specprefill._nemotron_h_extract_queries) | function | `_nemotron_h_extract_queries(attn, x, cache = None) -> not annotated` | Extract queries from Nemotron-H attention (no RoPE, no gate, no q_norm). | [#L116-L125](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L116-L125) |
| [`_patch_attention_for_capture`](#contract-vllm_mlx.specprefill._patch_attention_for_capture) | function | `_patch_attention_for_capture(model, query_buffer, query_extractor = None) -> not annotated` | Replace attention modules on full-attention layers with capture wrappers. | [#L128-L149](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L128-L149) |
| [`_unpatch_attention_capture`](#contract-vllm_mlx.specprefill._unpatch_attention_capture) | function | `_unpatch_attention_capture(model, originals) -> not annotated` | Restore original attention modules after capture. | [#L152-L155](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L152-L155) |
| [`_prefill_draft`](#contract-vllm_mlx.specprefill._prefill_draft) | function | `_prefill_draft(model, tokens, cache, step_size = 2048, cancel_check = None) -> not annotated` | Prefill prompt tokens into cache. | [#L158-L175](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L158-L175) |
| [`_lookahead_decode`](#contract-vllm_mlx.specprefill._lookahead_decode) | function | `_lookahead_decode(model, first_logits, cache, n_steps, temp = 0.6, top_p = 0.95, cancel_check = None) -> not annotated` | Run n_steps autoregressive decode, returning generated token ids. | [#L178-L204](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L178-L204) |
| [`_avg_pool1d`](#contract-vllm_mlx.specprefill._avg_pool1d) | function | `_avg_pool1d(x, kernel_size) -> not annotated` | 1D average pooling along last axis via prefix-sum. | [#L207-L223](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L207-L223) |
| [`_compute_importance`](#contract-vllm_mlx.specprefill._compute_importance) | function | `_compute_importance(query_buffer, attn_caches, n_prompt, n_attn_heads, n_kv_heads, pool_kernel = 13) -> not annotated` | Compute per-token importance from captured queries and cached keys. | [#L226-L271](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L226-L271) |
| [`score_tokens`](#contract-vllm_mlx.specprefill.score_tokens) | function | `score_tokens(model, tokens, n_lookahead = 8, pool_kernel = 13, temp = 0.6, top_p = 0.95, prefill_step_size = 2048, query_extractor = None, cancel_check = None) -> not annotated` | Score token importance using attention-based analysis on a draft model. | [#L274-L396](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L274-L396) |
| [`select_chunks`](#contract-vllm_mlx.specprefill.select_chunks) | function | `select_chunks(importance, keep_pct = 0.3, chunk_size = 32, backbone_pct = 0.0) -> not annotated` | Select top-k% token chunks by average importance. | [#L399-L467](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L399-L467) |
| [`select_chunks._selected_token_count`](#contract-vllm_mlx.specprefill.select_chunks._selected_token_count) | nested function | `select_chunks._selected_token_count(chunks) -> not annotated` | Nested Function `select_chunks._selected_token_count` calls `min`; returns `total`. | [#L437-L443](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L437-L443) |
| [`manual_rope`](#contract-vllm_mlx.specprefill.manual_rope) | function | `manual_rope(x, positions, dims, base = 10000.0, scale = 1.0) -> not annotated` | Apply RoPE at arbitrary (non-contiguous) positions. | [#L480-L508](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L480-L508) |
| [`manual_rope_with_freqs`](#contract-vllm_mlx.specprefill.manual_rope_with_freqs) | function | `manual_rope_with_freqs(x, positions, dims, freqs, pre_scale = 1.0) -> not annotated` | Apply RoPE at arbitrary positions using pre-computed frequencies. | [#L511-L528](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L511-L528) |
| [`_PositionMappedRoPE`](#contract-vllm_mlx.specprefill._PositionMappedRoPE) | class | `_PositionMappedRoPE(original_rope, all_positions, cache_start = 0)` | Wraps a RoPE module to apply rotation at non-contiguous positions. | [#L536-L571](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L536-L571) |
| [`_PositionMappedRoPE.__init__`](#contract-vllm_mlx.specprefill._PositionMappedRoPE.__init__) | method | `_PositionMappedRoPE.__init__(original_rope, all_positions, cache_start = 0) -> not annotated` | Method `_PositionMappedRoPE.__init__` updates `self._original`, `self._all_positions`, `self._cache_start`, `self._has_custom_freqs`; calls `hasattr`, `_get_dims`, `_get_pre_scale`. | [#L547-L561](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L547-L561) |
| [`_PositionMappedRoPE.__call__`](#contract-vllm_mlx.specprefill._PositionMappedRoPE.__call__) | method | `_PositionMappedRoPE.__call__(x, offset = 0) -> not annotated` | Method `_PositionMappedRoPE.__call__` calls `manual_rope_with_freqs`, `manual_rope`; has 2 explicit return paths. | [#L563-L571](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L563-L571) |
| [`_OffsetAdjustedRoPE`](#contract-vllm_mlx.specprefill._OffsetAdjustedRoPE) | class | `_OffsetAdjustedRoPE(original_rope, adjustment)` | Wraps a RoPE module to add a constant offset for decode after sparse prefill. | [#L574-L590](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L574-L590) |
| [`_OffsetAdjustedRoPE.__init__`](#contract-vllm_mlx.specprefill._OffsetAdjustedRoPE.__init__) | method | `_OffsetAdjustedRoPE.__init__(original_rope, adjustment) -> not annotated` | Method `_OffsetAdjustedRoPE.__init__` updates `self._original`, `self._adjustment`. | [#L585-L587](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L585-L587) |
| [`_OffsetAdjustedRoPE.__call__`](#contract-vllm_mlx.specprefill._OffsetAdjustedRoPE.__call__) | method | `_OffsetAdjustedRoPE.__call__(x, offset = 0) -> not annotated` | Method `_OffsetAdjustedRoPE.__call__` calls `self._original`; returns `self._original(x, offset=offset + self._adjustment)`. | [#L589-L590](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L589-L590) |
| [`_get_dims`](#contract-vllm_mlx.specprefill._get_dims) | function | `_get_dims(rope_module) -> not annotated` | Extract rotary dimensions from any RoPE variant. | [#L598-L603](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L598-L603) |
| [`_get_pre_scale`](#contract-vllm_mlx.specprefill._get_pre_scale) | function | `_get_pre_scale(rope_module) -> not annotated` | Extract pre-scale factor from custom RoPE variants (SuScaled, Yarn). | [#L606-L612](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L606-L612) |
| [`_find_attention_layers`](#contract-vllm_mlx.specprefill._find_attention_layers) | function | `_find_attention_layers(model) -> not annotated` | Find all full-attention layers across architectures. | [#L615-L630](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L615-L630) |
| [`_get_attn_module`](#contract-vllm_mlx.specprefill._get_attn_module) | function | `_get_attn_module(layer) -> not annotated` | Get the attention module from a layer (self_attn or mixer). | [#L633-L639](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L633-L639) |
| [`_get_rope`](#contract-vllm_mlx.specprefill._get_rope) | function | `_get_rope(attn) -> not annotated` | Get the RoPE module from an attention layer, or None. | [#L642-L647](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L642-L647) |
| [`_set_rope`](#contract-vllm_mlx.specprefill._set_rope) | function | `_set_rope(attn, rope_module) -> not annotated` | Set the RoPE module on an attention layer. | [#L650-L655](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L650-L655) |
| [`_set_attn_module`](#contract-vllm_mlx.specprefill._set_attn_module) | function | `_set_attn_module(layer, module) -> not annotated` | Set the attention module on a layer (self_attn or mixer). | [#L658-L663](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L658-L663) |
| [`_build_layer_to_cache_map`](#contract-vllm_mlx.specprefill._build_layer_to_cache_map) | function | `_build_layer_to_cache_map(model) -> not annotated` | Build mapping from model layer index to cache index. | [#L666-L690](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L666-L690) |
| [`sparse_prefill`](#contract-vllm_mlx.specprefill.sparse_prefill) | function | `sparse_prefill(model, tokens, selected_indices, cache, step_size = 2048, position_offset = 0, cancel_check = None) -> not annotated` | Prefill the model cache with selected tokens at their original positions. | [#L698-L827](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L698-L827) |
| [`cleanup_rope`](#contract-vllm_mlx.specprefill.cleanup_rope) | function | `cleanup_rope(model) -> not annotated` | Restore original RoPE on all attention layers. | [#L830-L845](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/vllm_mlx/specprefill.py#L830-L845) |