vllm_mlx.specprefill

SpecPrefill: Attention-based sparse prefill for MLX.

View the complete module source at #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

SpecPrefill: Attention-based sparse prefill for MLX.

Full pipeline for reducing TTFT on long prompts

Step 1 (score_tokens): Use a small draft model to identify important tokens Step 2 (sparse_prefill): Prefill target model with only selected tokens, preserving original positional encoding via manual RoPE

Usage

from specprefill import score_tokens, select_chunks, sparse_prefill, cleanup_rope

1. Score with draft model

importance = score_tokens(draft_model, tokens)

2. Select important token chunks

selected = select_chunks(importance, keep_pct=0.3)

3. Sparse prefill on target model

target_cache = make_prompt_cache(target_model) logits = sparse_prefill(target_model, tokens, selected, target_cache)

4. Generate normally using target_cache...

5. Cleanup

cleanup_rope(target_model)

Design notes

• RoPE is relative: Q_m @ K_p^T depends only on (m - p). Selected keys stored contiguously in the cache buffer with correct RoPE angles produce correct attention during decode.
• After sparse prefill of N tokens from a total prompt of M, cache.offset = N but decode RoPE needs position M. The _OffsetAdjustedRoPE adds (M - N) to each RoPE offset call, so decode position = N + i + (M - N) = M + i.
• GatedDeltaNet (linear attention) layers process sparse tokens through their conv/SSM state normally. This is lossy but acceptable per the SpecPrefill paper — attention layers are the primary long-range mechanism.

Reference: arxiv.org/abs/2502.02789 (SpecPrefill: Speculative Prefilling)

vllm_mlx.specprefill._AttentionCapture

_AttentionCapture(original, buf_idx, query_buffer, query_extractor=None)

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

Installed on attention layers during lookahead decode to capture query vectors for importance scoring. Supports multiple architectures via query_extractor callback.

Source code in vllm_mlx/specprefill.py

def __init__(self, original, buf_idx, query_buffer, query_extractor=None):
    self._original = original
    self._buf_idx = buf_idx
    self._query_buffer = query_buffer
    self._query_extractor = query_extractor or _qwen35_extract_queries

vllm_mlx.specprefill._AttentionCapture._original instance-attribute

_original = original

vllm_mlx.specprefill._AttentionCapture._buf_idx instance-attribute

_buf_idx = buf_idx

vllm_mlx.specprefill._AttentionCapture._query_buffer instance-attribute

_query_buffer = query_buffer

vllm_mlx.specprefill._AttentionCapture._query_extractor instance-attribute

_query_extractor = query_extractor or _qwen35_extract_queries

vllm_mlx.specprefill._AttentionCapture.__call__

__call__(x, mask=None, cache=None)

Source code in vllm_mlx/specprefill.py

def __call__(self, x, mask=None, cache=None):
    queries = self._query_extractor(self._original, x, cache)
    self._query_buffer[self._buf_idx].append(queries)
    return self._original(x, mask=mask, cache=cache)

vllm_mlx.specprefill._AttentionCapture.__getattr__

__getattr__(name)

Source code in vllm_mlx/specprefill.py

def __getattr__(self, name):
    return getattr(self._original, name)

vllm_mlx.specprefill._PositionMappedRoPE

_PositionMappedRoPE(original_rope, all_positions, cache_start=0)

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

Used during sparse prefill. The offset parameter from the cache tells us which slice of the position array to use for the current chunk: positions = all_positions[(offset - cache_start) : (offset - cache_start) + L]

When composing with a pre-populated cache (e.g., system KV cache), cache_start is the initial cache offset so indexing into the position array is correct.

Source code in vllm_mlx/specprefill.py

def __init__(self, original_rope, all_positions, cache_start=0):
    self._original = original_rope
    self._all_positions = all_positions
    self._cache_start = cache_start
    self._has_custom_freqs = hasattr(original_rope, "_freqs")

    if self._has_custom_freqs:
        self._freqs = original_rope._freqs
        self._dims = _get_dims(original_rope)
        self._pre_scale = _get_pre_scale(original_rope)
    else:
        # Standard nn.RoPE: attributes are dims, base, scale (no underscore)
        self._dims = original_rope.dims
        self._base = original_rope.base
        self._scale = original_rope.scale

vllm_mlx.specprefill._PositionMappedRoPE._original instance-attribute

_original = original_rope

vllm_mlx.specprefill._PositionMappedRoPE._all_positions instance-attribute

_all_positions = all_positions

vllm_mlx.specprefill._PositionMappedRoPE._cache_start instance-attribute

_cache_start = cache_start

vllm_mlx.specprefill._PositionMappedRoPE._has_custom_freqs instance-attribute

_has_custom_freqs = hasattr(original_rope, '_freqs')

vllm_mlx.specprefill._PositionMappedRoPE._freqs instance-attribute

_freqs = original_rope._freqs

vllm_mlx.specprefill._PositionMappedRoPE._dims instance-attribute

_dims = _get_dims(original_rope)

vllm_mlx.specprefill._PositionMappedRoPE._pre_scale instance-attribute

_pre_scale = _get_pre_scale(original_rope)

vllm_mlx.specprefill._PositionMappedRoPE._base instance-attribute

_base = original_rope.base

vllm_mlx.specprefill._PositionMappedRoPE._scale instance-attribute

_scale = original_rope.scale

vllm_mlx.specprefill._PositionMappedRoPE.__call__

__call__(x, offset=0)

Source code in vllm_mlx/specprefill.py

def __call__(self, x, offset=0):
    L = x.shape[2]
    idx = offset - self._cache_start
    positions = self._all_positions[idx : idx + L]
    if self._has_custom_freqs:
        return manual_rope_with_freqs(
            x, positions, self._dims, self._freqs, pre_scale=self._pre_scale
        )
    return manual_rope(x, positions, self._dims, base=self._base, scale=self._scale)

vllm_mlx.specprefill._OffsetAdjustedRoPE

_OffsetAdjustedRoPE(original_rope, adjustment)

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

After sparse prefill of N tokens from a prompt of M total tokens

cache.offset = N + i (i = decode step) desired RoPE position = M + i adjustment = M - N

So: RoPE(x, offset = cache.offset + adjustment) = RoPE(x, M + i)

Source code in vllm_mlx/specprefill.py

def __init__(self, original_rope, adjustment):
    self._original = original_rope
    self._adjustment = adjustment

vllm_mlx.specprefill._OffsetAdjustedRoPE._original instance-attribute

_original = original_rope

vllm_mlx.specprefill._OffsetAdjustedRoPE._adjustment instance-attribute

_adjustment = adjustment

vllm_mlx.specprefill._OffsetAdjustedRoPE.__call__

__call__(x, offset=0)

Source code in vllm_mlx/specprefill.py

def __call__(self, x, offset=0):
    return self._original(x, offset=offset + self._adjustment)

vllm_mlx.specprefill._qwen35_extract_queries

_qwen35_extract_queries(attn, x, cache=None)

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

Qwen3.5 q_proj output is 2x wider: [queries, gate]. We split, normalize, then apply RoPE.

Source code in vllm_mlx/specprefill.py

def _qwen35_extract_queries(attn, x, cache=None):
    """Extract post-RoPE queries from Qwen3.5 attention (gate split + q_norm).

    Qwen3.5 q_proj output is 2x wider: [queries, gate]. We split, normalize,
    then apply RoPE.
    """
    B, L, D = x.shape
    q_out = attn.q_proj(x)
    queries, _gate = mx.split(
        q_out.reshape(B, L, attn.num_attention_heads, -1), 2, axis=-1
    )
    queries = attn.q_norm(queries).transpose(0, 2, 1, 3)
    if cache is not None:
        queries = attn.rope(queries, offset=cache.offset)
    else:
        queries = attn.rope(queries)
    return queries

vllm_mlx.specprefill._llama_extract_queries

_llama_extract_queries(attn, x, cache=None)

Extract post-RoPE queries from standard transformer attention.

Standard architecture: q_proj → reshape → RoPE. No gate, no q_norm. Works for Llama 3.x, Mistral, Gemma, GPT-OSS, and other GQA models.

Source code in vllm_mlx/specprefill.py

def _llama_extract_queries(attn, x, cache=None):
    """Extract post-RoPE queries from standard transformer attention.

    Standard architecture: q_proj → reshape → RoPE. No gate, no q_norm.
    Works for Llama 3.x, Mistral, Gemma, GPT-OSS, and other GQA models.
    """
    B, L, D = x.shape
    n_heads = getattr(
        attn,
        "num_attention_heads",
        getattr(attn, "n_heads", getattr(attn, "num_heads", None)),
    )
    queries = attn.q_proj(x)
    queries = queries.reshape(B, L, n_heads, -1).transpose(0, 2, 1, 3)
    if cache is not None:
        queries = attn.rope(queries, offset=cache.offset)
    else:
        queries = attn.rope(queries)
    return queries

vllm_mlx.specprefill._nemotron_h_extract_queries

_nemotron_h_extract_queries(attn, x, cache=None)

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

Nemotron-H attention layers have NO positional encoding — RoPE is absent. Positional modeling comes from Mamba2 layers. Attention is content-based only.

Source code in vllm_mlx/specprefill.py

def _nemotron_h_extract_queries(attn, x, cache=None):
    """Extract queries from Nemotron-H attention (no RoPE, no gate, no q_norm).

    Nemotron-H attention layers have NO positional encoding — RoPE is absent.
    Positional modeling comes from Mamba2 layers. Attention is content-based only.
    """
    B, L, D = x.shape
    queries = attn.q_proj(x).reshape(B, L, attn.num_heads, -1).transpose(0, 2, 1, 3)
    # No RoPE to apply — queries are used as-is for content-based scoring
    return queries

vllm_mlx.specprefill._patch_attention_for_capture

_patch_attention_for_capture(model, query_buffer, query_extractor=None)

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

Supports both self_attn (Qwen3.5/Llama/GPT-OSS) and mixer (Nemotron-H block_type="*") attribute conventions.

Returns (originals, attn_layer_indices) for cleanup.

Source code in vllm_mlx/specprefill.py

def _patch_attention_for_capture(model, query_buffer, query_extractor=None):
    """Replace attention modules on full-attention layers with capture wrappers.

    Supports both `self_attn` (Qwen3.5/Llama/GPT-OSS) and `mixer`
    (Nemotron-H block_type="*") attribute conventions.

    Returns (originals, attn_layer_indices) for cleanup.
    """
    originals = []
    attn_indices = []
    for layer_idx, layer in _find_attention_layers(model):
        buf_idx = len(attn_indices)
        attn_indices.append(layer_idx)
        orig = _get_attn_module(layer)
        _set_attn_module(
            layer,
            _AttentionCapture(
                orig, buf_idx, query_buffer, query_extractor=query_extractor
            ),
        )
        originals.append((layer_idx, orig))
    return originals, attn_indices

vllm_mlx.specprefill._unpatch_attention_capture

_unpatch_attention_capture(model, originals)

Restore original attention modules after capture.

Source code in vllm_mlx/specprefill.py

def _unpatch_attention_capture(model, originals):
    """Restore original attention modules after capture."""
    for layer_idx, orig in originals:
        _set_attn_module(model.layers[layer_idx], orig)

vllm_mlx.specprefill._prefill_draft

_prefill_draft(model, tokens, cache, step_size=2048, cancel_check=None)

Prefill prompt tokens into cache. Returns logits from last token.

Source code in vllm_mlx/specprefill.py

def _prefill_draft(model, tokens, cache, step_size=2048, cancel_check=None):
    """Prefill prompt tokens into cache. Returns logits from last token."""
    prompt = mx.array(tokens) if not isinstance(tokens, mx.array) else tokens
    n = len(tokens)
    processed = 0
    while n - processed > 1:
        if cancel_check is not None:
            cancel_check()
        chunk = min(step_size, n - processed - 1)
        model(prompt[processed : processed + chunk][None], cache=cache)
        mx.eval([c.state for c in cache])
        processed += chunk
        mx.clear_cache()
    if cancel_check is not None:
        cancel_check()
    logits = model(prompt[processed:][None], cache=cache)
    mx.eval(logits)
    return logits

vllm_mlx.specprefill._lookahead_decode

_lookahead_decode(model, first_logits, cache, n_steps, temp=0.6, top_p=0.95, cancel_check=None)

Run n_steps autoregressive decode, returning generated token ids.

Query vectors are captured by the monkey-patched attention layers.

Source code in vllm_mlx/specprefill.py

def _lookahead_decode(
    model,
    first_logits,
    cache,
    n_steps,
    temp=0.6,
    top_p=0.95,
    cancel_check=None,
):
    """Run n_steps autoregressive decode, returning generated token ids.

    Query vectors are captured by the monkey-patched attention layers.
    """
    sampler = make_sampler(temp=temp, top_p=top_p)
    if cancel_check is not None:
        cancel_check()
    y = sampler(first_logits[:, -1, :])
    mx.eval(y)
    generated = [y.item()]
    for _ in range(n_steps):
        if cancel_check is not None:
            cancel_check()
        logits = model(y.reshape(1, -1), cache=cache)
        y = sampler(logits[:, -1, :])
        mx.eval(y)
        generated.append(y.item())
    return generated

vllm_mlx.specprefill._avg_pool1d

_avg_pool1d(x, kernel_size)

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

Parameters:

• x –

  (..., M) input

• kernel_size –

  window size (odd for centered)

Returns:

• –

  (..., M) pooled (same size, zero-padded at edges)

Source code in vllm_mlx/specprefill.py

def _avg_pool1d(x, kernel_size):
    """1D average pooling along last axis via prefix-sum.

    Args:
        x: (..., M) input
        kernel_size: window size (odd for centered)

    Returns:
        (..., M) pooled (same size, zero-padded at edges)
    """
    if kernel_size <= 1:
        return x
    pad = kernel_size // 2
    padded = mx.pad(x, [(0, 0)] * (x.ndim - 1) + [(pad, pad)])
    zeros = mx.zeros(x.shape[:-1] + (1,), dtype=x.dtype)
    prefix = mx.concatenate([zeros, mx.cumsum(padded, axis=-1)], axis=-1)
    return (prefix[..., kernel_size:] - prefix[..., :-kernel_size]) / kernel_size

vllm_mlx.specprefill._compute_importance

_compute_importance(query_buffer, attn_caches, n_prompt, n_attn_heads, n_kv_heads, pool_kernel=13)

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

Aggregation (SpecPrefill paper): 1. softmax(Q @ K^T / sqrt(d)) per head, per layer, per lookahead token 2. avg_pool1d smoothing 3. max across (layers × heads) 4. mean across lookahead tokens

Returns: (n_prompt,) importance scores.

Source code in vllm_mlx/specprefill.py

def _compute_importance(
    query_buffer, attn_caches, n_prompt, n_attn_heads, n_kv_heads, pool_kernel=13
):
    """Compute per-token importance from captured queries and cached keys.

    Aggregation (SpecPrefill paper):
      1. softmax(Q @ K^T / sqrt(d)) per head, per layer, per lookahead token
      2. avg_pool1d smoothing
      3. max across (layers × heads)
      4. mean across lookahead tokens

    Returns: (n_prompt,) importance scores.
    """
    heads_per_group = n_attn_heads // n_kv_heads
    all_scores = []

    for layer_i, captures in enumerate(query_buffer):
        if not captures:
            continue
        cache = attn_caches[layer_i]
        prompt_keys = cache.keys[..., :n_prompt, :]
        # Skip layers with windowed/rotating caches that don't span
        # the full prompt (e.g., GPT-OSS sliding_attention with 128-token window).
        # These lack global context and would produce mismatched score shapes.
        if prompt_keys.shape[-2] < n_prompt:
            continue
        head_dim = prompt_keys.shape[-1]
        q_stack = mx.concatenate(captures, axis=2)
        if heads_per_group > 1:
            expanded_keys = mx.repeat(prompt_keys, heads_per_group, axis=1)
        else:
            expanded_keys = prompt_keys
        scale = head_dim**-0.5
        scores = (q_stack @ expanded_keys.transpose(0, 1, 3, 2)) * scale
        weights = mx.softmax(scores.astype(mx.float32), axis=-1)
        all_scores.append(weights.squeeze(0))

    if not all_scores:
        raise RuntimeError("No attention scores captured — check model/patching")

    combined = mx.concatenate(all_scores, axis=0)
    if pool_kernel and pool_kernel > 1:
        combined = _avg_pool1d(combined, pool_kernel)
    max_scores = mx.max(combined, axis=0)
    importance = mx.mean(max_scores, axis=0)
    return importance

vllm_mlx.specprefill.score_tokens

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)

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

Runs the full scoring pipeline

1. Prefill the draft model with all tokens
2. N lookahead decode steps, capturing query vectors from attention layers
3. Compute importance: Q_lookahead @ K_prompt^T, aggregated across heads/layers

The draft model's cache is created internally and discarded after scoring.

Parameters:

• model –

  Draft model (small, fast — e.g. 4B)

• tokens –

  list or mx.array of token IDs

• n_lookahead –

  decode steps for query capture (default 8)

• pool_kernel –

  smoothing kernel for avg_pool1d (default 13, 0=disable)

• temp –

  sampling temperature for lookahead (default 0.6)

• top_p –

  top-p for lookahead (default 0.95)

• prefill_step_size –

  chunk size for draft prefill (default 2048)

• query_extractor –

  function(attn, x, cache) → queries tensor. Default: _qwen35_extract_queries. Use _llama_extract_queries for standard Llama/Mistral/Gemma models.

Returns:

• importance –

  (M,) mx.array of per-token importance scores

Source code in vllm_mlx/specprefill.py

def 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,
):
    """Score token importance using attention-based analysis on a draft model.

    Runs the full scoring pipeline:
      1. Prefill the draft model with all tokens
      2. N lookahead decode steps, capturing query vectors from attention layers
      3. Compute importance: Q_lookahead @ K_prompt^T, aggregated across heads/layers

    The draft model's cache is created internally and discarded after scoring.

    Args:
        model: Draft model (small, fast — e.g. 4B)
        tokens: list or mx.array of token IDs
        n_lookahead: decode steps for query capture (default 8)
        pool_kernel: smoothing kernel for avg_pool1d (default 13, 0=disable)
        temp: sampling temperature for lookahead (default 0.6)
        top_p: top-p for lookahead (default 0.95)
        prefill_step_size: chunk size for draft prefill (default 2048)
        query_extractor: function(attn, x, cache) → queries tensor.
            Default: _qwen35_extract_queries. Use _llama_extract_queries for
            standard Llama/Mistral/Gemma models.

    Returns:
        importance: (M,) mx.array of per-token importance scores
    """
    if isinstance(tokens, mx.array):
        tokens = tokens.tolist()
    n_prompt = len(tokens)

    # Model topology — detect attribute names across architectures
    attn_layers = _find_attention_layers(model)
    n_attn_layers = len(attn_layers)
    attn_obj = _get_attn_module(attn_layers[0][1])
    # Attribute names vary: num_attention_heads (Qwen3.5), n_heads (Llama),
    # num_heads (Nemotron-H)
    n_attn_heads = getattr(
        attn_obj,
        "num_attention_heads",
        getattr(attn_obj, "n_heads", getattr(attn_obj, "num_heads", None)),
    )
    n_kv_heads = getattr(
        attn_obj, "num_key_value_heads", getattr(attn_obj, "n_kv_heads", None)
    )

    # Auto-detect query extractor from model_type (explicit registry, not
    # attribute sniffing -- avoids silent misclassification on new models).
    if query_extractor is None:
        model_type = getattr(getattr(model, "config", None), "model_type", "")
        _EXTRACTOR_REGISTRY = {
            "qwen3_5": _qwen35_extract_queries,
            "qwen3_5_moe": _qwen35_extract_queries,
            "qwen3_vl": _qwen35_extract_queries,
            "qwen3_vl_moe": _qwen35_extract_queries,
            "nemotron_h": _nemotron_h_extract_queries,
        }
        query_extractor = _EXTRACTOR_REGISTRY.get(model_type)
        if query_extractor is None:
            if _get_rope(attn_obj) is not None:
                query_extractor = _llama_extract_queries
            else:
                query_extractor = _nemotron_h_extract_queries

    # Phase 1: Prefill
    cache = make_prompt_cache(model)
    logits = _prefill_draft(
        model,
        tokens,
        cache,
        step_size=prefill_step_size,
        cancel_check=cancel_check,
    )

    # Phase 2: Lookahead decode with query capture
    query_buffer = [[] for _ in range(n_attn_layers)]
    patches, attn_indices = _patch_attention_for_capture(
        model, query_buffer, query_extractor=query_extractor
    )
    try:
        _lookahead_decode(
            model,
            logits,
            cache,
            n_lookahead,
            temp=temp,
            top_p=top_p,
            cancel_check=cancel_check,
        )
        mx.eval(query_buffer)
    finally:
        _unpatch_attention_capture(model, patches)

    # Phase 3: Compute importance
    # Map layer indices to cache indices (identity for standard models,
    # compacted for Nemotron-H where only M/* layers have cache entries)
    layer_to_cache = _build_layer_to_cache_map(model)
    attn_caches = [cache[layer_to_cache[i]] for i in attn_indices]
    if cancel_check is not None:
        cancel_check()
    importance = _compute_importance(
        query_buffer,
        attn_caches,
        n_prompt,
        n_attn_heads,
        n_kv_heads,
        pool_kernel=pool_kernel if pool_kernel > 0 else None,
    )
    mx.eval(importance)

    # Draft cache is no longer needed — let GC reclaim it
    del cache, logits, query_buffer, attn_caches
    mx.clear_cache()

    return importance

vllm_mlx.specprefill.select_chunks

select_chunks(importance, keep_pct=0.3, chunk_size=32, backbone_pct=0.0)

Select top-k% token chunks by average importance.

Parameters:

• importance –

  (M,) per-token importance scores

• keep_pct –

  fraction of chunks to keep (default 0.3)

• chunk_size –

  tokens per chunk (default 32)

• backbone_pct –

  fraction of chunks reserved for evenly-spaced coverage

Returns:

• –

  sorted mx.array of kept token indices

Source code in vllm_mlx/specprefill.py

def select_chunks(importance, keep_pct=0.3, chunk_size=32, backbone_pct=0.0):
    """Select top-k% token chunks by average importance.

    Args:
        importance: (M,) per-token importance scores
        keep_pct: fraction of chunks to keep (default 0.3)
        chunk_size: tokens per chunk (default 32)
        backbone_pct: fraction of chunks reserved for evenly-spaced coverage

    Returns:
        sorted mx.array of kept token indices
    """
    M = importance.shape[0]
    if keep_pct >= 1.0:
        return mx.arange(M)

    n_chunks = math.ceil(M / chunk_size)
    target_tokens = max(1, math.ceil(M * keep_pct))
    keep_n = max(1, math.ceil(n_chunks * keep_pct))
    backbone_n = max(0, math.ceil(n_chunks * backbone_pct)) if backbone_pct > 0 else 0
    top_n = max(0, keep_n - backbone_n)

    chunk_scores = []
    for i in range(n_chunks):
        start = i * chunk_size
        end = min(start + chunk_size, M)
        chunk_scores.append(mx.mean(importance[start:end]).item())

    selected_chunks = set(
        sorted(range(n_chunks), key=lambda i: chunk_scores[i], reverse=True)[:top_n]
    )
    if backbone_n > 0:
        if backbone_n >= n_chunks:
            selected_chunks.update(range(n_chunks))
        else:
            for i in range(backbone_n):
                selected_chunks.add(round(i * (n_chunks - 1) / max(1, backbone_n - 1)))

    def _selected_token_count(chunks):
        total = 0
        for chunk_idx in chunks:
            start = chunk_idx * chunk_size
            end = min(start + chunk_size, M)
            total += end - start
        return total

    if (
        len(selected_chunks) < keep_n
        or _selected_token_count(selected_chunks) < target_tokens
    ):
        for chunk_idx in sorted(
            range(n_chunks), key=lambda i: chunk_scores[i], reverse=True
        ):
            selected_chunks.add(chunk_idx)
            if (
                len(selected_chunks) >= keep_n
                and _selected_token_count(selected_chunks) >= target_tokens
            ):
                break

    top_chunks = sorted(selected_chunks)

    indices = []
    for ci in top_chunks:
        start = ci * chunk_size
        end = min(start + chunk_size, M)
        indices.extend(range(start, end))

    return mx.array(indices)

vllm_mlx.specprefill.manual_rope

manual_rope(x, positions, dims, base=10000.0, scale=1.0)

Apply RoPE at arbitrary (non-contiguous) positions.

Uses non-traditional (interleaved) layout matching Qwen3.5: rotates first dims dimensions as pairs [0,half), [half,dims), passes through [dims:] unchanged.

Parameters:

• x –

  (B, n_heads, L, head_dim) input tensor

• positions –

  (L,) position indices (can be non-contiguous)

• dims –

  number of dimensions to rotate (head_dim * partial_rotary_factor)

• base –

  RoPE base frequency (default 10000.0)

• scale –

  position scale divisor (default 1.0, higher = compressed positions)

Returns:

• –

  (B, n_heads, L, head_dim) with RoPE applied

Source code in vllm_mlx/specprefill.py

def manual_rope(x, positions, dims, base=10000.0, scale=1.0):
    """Apply RoPE at arbitrary (non-contiguous) positions.

    Uses non-traditional (interleaved) layout matching Qwen3.5:
    rotates first `dims` dimensions as pairs [0,half), [half,dims),
    passes through [dims:] unchanged.

    Args:
        x: (B, n_heads, L, head_dim) input tensor
        positions: (L,) position indices (can be non-contiguous)
        dims: number of dimensions to rotate (head_dim * partial_rotary_factor)
        base: RoPE base frequency (default 10000.0)
        scale: position scale divisor (default 1.0, higher = compressed positions)

    Returns:
        (B, n_heads, L, head_dim) with RoPE applied
    """
    half = dims // 2
    inv_freq = 1.0 / (base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims))
    scaled_pos = positions.astype(mx.float32) / scale
    angles = scaled_pos[:, None] * inv_freq[None, :]  # (L, half)
    cos_a = mx.cos(angles)[None, None, :, :]  # (1, 1, L, half)
    sin_a = mx.sin(angles)[None, None, :, :]
    x_rot, x_pass = x[..., :dims], x[..., dims:]
    x1, x2 = x_rot[..., :half], x_rot[..., half:]
    rotated = mx.concatenate(
        [x1 * cos_a - x2 * sin_a, x1 * sin_a + x2 * cos_a], axis=-1
    )
    return mx.concatenate([rotated, x_pass], axis=-1)

vllm_mlx.specprefill.manual_rope_with_freqs

manual_rope_with_freqs(x, positions, dims, freqs, pre_scale=1.0)

Apply RoPE at arbitrary positions using pre-computed frequencies.

For custom RoPE variants (Llama3, Yarn, SuScaled) that store _freqs.

Source code in vllm_mlx/specprefill.py

def manual_rope_with_freqs(x, positions, dims, freqs, pre_scale=1.0):
    """Apply RoPE at arbitrary positions using pre-computed frequencies.

    For custom RoPE variants (Llama3, Yarn, SuScaled) that store _freqs.
    """
    half = dims // 2
    inv_freq = (1.0 / freqs).astype(mx.float32)
    angles = positions[:, None].astype(mx.float32) * inv_freq[None, :]
    cos_a = mx.cos(angles)[None, None, :, :]
    sin_a = mx.sin(angles)[None, None, :, :]
    x_rot, x_pass = x[..., :dims], x[..., dims:]
    if pre_scale != 1.0:
        x_rot = pre_scale * x_rot
    x1, x2 = x_rot[..., :half], x_rot[..., half:]
    rotated = mx.concatenate(
        [x1 * cos_a - x2 * sin_a, x1 * sin_a + x2 * cos_a], axis=-1
    )
    return mx.concatenate([rotated, x_pass], axis=-1)

vllm_mlx.specprefill._get_dims

_get_dims(rope_module)

Extract rotary dimensions from any RoPE variant.

Source code in vllm_mlx/specprefill.py

def _get_dims(rope_module):
    """Extract rotary dimensions from any RoPE variant."""
    for attr in ("_dims", "dim", "dims"):
        if hasattr(rope_module, attr):
            return getattr(rope_module, attr)
    raise ValueError(f"Cannot determine dims from {type(rope_module)}")

vllm_mlx.specprefill._get_pre_scale

_get_pre_scale(rope_module)

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

Source code in vllm_mlx/specprefill.py

def _get_pre_scale(rope_module):
    """Extract pre-scale factor from custom RoPE variants (SuScaled, Yarn)."""
    if hasattr(rope_module, "mscale"):
        return rope_module.mscale
    if hasattr(rope_module, "_scale") and hasattr(rope_module, "dim"):
        return rope_module._scale
    return 1.0

vllm_mlx.specprefill._find_attention_layers

_find_attention_layers(model)

Find all full-attention layers across architectures.

Supports

• Qwen3.5 / Llama / GPT-OSS: layers with self_attn attribute
• Nemotron-H: layers with block_type == "*" (attention blocks use mixer)

Returns list of (layer_idx, layer) tuples.

Source code in vllm_mlx/specprefill.py

def _find_attention_layers(model):
    """Find all full-attention layers across architectures.

    Supports:
      - Qwen3.5 / Llama / GPT-OSS: layers with `self_attn` attribute
      - Nemotron-H: layers with `block_type == "*"` (attention blocks use `mixer`)

    Returns list of (layer_idx, layer) tuples.
    """
    results = []
    for idx, layer in enumerate(model.layers):
        if hasattr(layer, "self_attn"):
            results.append((idx, layer))
        elif getattr(layer, "block_type", None) == "*":
            results.append((idx, layer))
    return results

vllm_mlx.specprefill._get_attn_module

_get_attn_module(layer)

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

Source code in vllm_mlx/specprefill.py

def _get_attn_module(layer):
    """Get the attention module from a layer (self_attn or mixer)."""
    if hasattr(layer, "self_attn"):
        return layer.self_attn
    if getattr(layer, "block_type", None) == "*":
        return layer.mixer
    return None

vllm_mlx.specprefill._get_rope

_get_rope(attn)

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

mlx_lm models use self.rope; mlx_vlm models use self.rotary_emb.

Source code in vllm_mlx/specprefill.py

def _get_rope(attn):
    """Get the RoPE module from an attention layer, or None.

    mlx_lm models use ``self.rope``; mlx_vlm models use ``self.rotary_emb``.
    """
    return getattr(attn, "rope", None) or getattr(attn, "rotary_emb", None)

vllm_mlx.specprefill._set_rope

_set_rope(attn, rope_module)

Set the RoPE module on an attention layer.

Source code in vllm_mlx/specprefill.py

def _set_rope(attn, rope_module):
    """Set the RoPE module on an attention layer."""
    if hasattr(attn, "rope"):
        attn.rope = rope_module
    elif hasattr(attn, "rotary_emb"):
        attn.rotary_emb = rope_module

vllm_mlx.specprefill._set_attn_module

_set_attn_module(layer, module)

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

Source code in vllm_mlx/specprefill.py

def _set_attn_module(layer, module):
    """Set the attention module on a layer (self_attn or mixer)."""
    if hasattr(layer, "self_attn"):
        layer.self_attn = module
    elif getattr(layer, "block_type", None) == "*":
        layer.mixer = module

vllm_mlx.specprefill._build_layer_to_cache_map

_build_layer_to_cache_map(model)

Build mapping from model layer index to cache index.

Standard models (Qwen3.5, Llama, GPT-OSS): one cache entry per layer, so the mapping is identity (layer_idx → layer_idx).

Nemotron-H: only M (Mamba2) and * (attention) layers have cache entries. MLP (-) and MoE (E) layers get no cache. The mapping is compacted.

Returns dict {layer_idx: cache_idx}.

Source code in vllm_mlx/specprefill.py

def _build_layer_to_cache_map(model):
    """Build mapping from model layer index to cache index.

    Standard models (Qwen3.5, Llama, GPT-OSS): one cache entry per layer,
    so the mapping is identity (layer_idx → layer_idx).

    Nemotron-H: only M (Mamba2) and * (attention) layers have cache entries.
    MLP (-) and MoE (E) layers get no cache. The mapping is compacted.

    Returns dict {layer_idx: cache_idx}.
    """
    has_block_type = any(hasattr(layer, "block_type") for layer in model.layers)
    if not has_block_type:
        # Standard model: identity mapping
        return {i: i for i in range(len(model.layers))}

    # Nemotron-H style: count cache entries for M/* layers
    layer_to_cache = {}
    cache_idx = 0
    for layer_idx, layer in enumerate(model.layers):
        bt = getattr(layer, "block_type", None)
        if bt in ("M", "*"):
            layer_to_cache[layer_idx] = cache_idx
            cache_idx += 1
    return layer_to_cache

vllm_mlx.specprefill.sparse_prefill

sparse_prefill(model, tokens, selected_indices, cache, step_size=2048, position_offset=0, cancel_check=None)

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

Runs the model forward on only the selected tokens while preserving their original positional encoding via manual RoPE. After this call, the cache contains KV entries with correct RoPE positions, and attention layers have _OffsetAdjustedRoPE installed for correct decode positioning.

Parameters:

• model –

  Language model with .layers property (TextModel or VLM Model)

• tokens –

  (M,) all prompt token IDs (mx.array or list)

• selected_indices –

  (N,) sorted indices into tokens to keep (mx.array or list)

• cache –

  list of KVCache/ArraysCache from make_prompt_cache()

• step_size –

  chunk size for processing (default 2048)

• position_offset –

  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).

Returns:

• logits –

  (1, 1, vocab_size) from the last selected token

Side effects

• Populates cache with KV for selected tokens
• Installs _OffsetAdjustedRoPE on attention layers for decode
• Call cleanup_rope(model) after generation to restore original RoPE

Source code in vllm_mlx/specprefill.py

def sparse_prefill(
    model,
    tokens,
    selected_indices,
    cache,
    step_size=2048,
    position_offset=0,
    cancel_check=None,
):
    """Prefill the model cache with selected tokens at their original positions.

    Runs the model forward on only the selected tokens while preserving their
    original positional encoding via manual RoPE. After this call, the cache
    contains KV entries with correct RoPE positions, and attention layers have
    _OffsetAdjustedRoPE installed for correct decode positioning.

    Args:
        model: Language model with .layers property (TextModel or VLM Model)
        tokens: (M,) all prompt token IDs (mx.array or list)
        selected_indices: (N,) sorted indices into tokens to keep (mx.array or list)
        cache: list of KVCache/ArraysCache from make_prompt_cache()
        step_size: chunk size for processing (default 2048)
        position_offset: 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).

    Returns:
        logits: (1, 1, vocab_size) from the last selected token

    Side effects:
        - Populates cache with KV for selected tokens
        - Installs _OffsetAdjustedRoPE on attention layers for decode
        - Call cleanup_rope(model) after generation to restore original RoPE
    """
    if not isinstance(tokens, mx.array):
        tokens = mx.array(tokens)
    if not isinstance(selected_indices, mx.array):
        selected_indices = mx.array(selected_indices)

    M = tokens.shape[0]

    # Detect RotatingKVCache and ensure tail tokens are included only when the
    # prompt actually exceeds the live cache window. If the full prompt still
    # fits inside ``max_size`` there is no eviction yet, so forcing the entire
    # tail back in would collapse sparse prefill into dense work.
    max_rotating_size = 0
    for c in cache:
        if type(c).__name__ == "RotatingKVCache":
            max_rotating_size = max(max_rotating_size, getattr(c, "max_size", 0))
    if max_rotating_size > 0 and M > max_rotating_size:
        tail_start = max(0, M - max_rotating_size)
        tail_indices = set(range(tail_start, M))
        existing = set(selected_indices.tolist())
        merged = sorted(existing | tail_indices)
        selected_indices = mx.array(merged)

    # RoPE positions: absolute positions accounting for any prefix
    selected_positions = selected_indices.astype(mx.int32) + position_offset
    selected_tokens = tokens[selected_indices]
    N = selected_tokens.shape[0]

    # Determine initial cache offset (non-zero when system KV cache is restored)
    attn_layers = _find_attention_layers(model)
    layer_to_cache = _build_layer_to_cache_map(model)
    first_attn_layer_idx = attn_layers[0][0]
    first_attn_cache_idx = layer_to_cache[first_attn_layer_idx]
    cache_start = (
        cache[first_attn_cache_idx].offset
        if hasattr(cache[first_attn_cache_idx], "offset")
        else 0
    )

    # Check if attention layers use RoPE (Nemotron-H has none)
    first_attn = _get_attn_module(attn_layers[0][1])
    has_rope = _get_rope(first_attn) is not None

    # Patch RoPE on attention layers for position-mapped prefill
    # (skipped for architectures without RoPE, e.g. Nemotron-H)
    original_ropes = {}
    if has_rope:
        for layer_idx, layer in attn_layers:
            attn = _get_attn_module(layer)
            rope = _get_rope(attn)
            original_ropes[layer_idx] = (attn, rope)
            _set_rope(
                attn,
                _PositionMappedRoPE(rope, selected_positions, cache_start=cache_start),
            )

    try:
        prompt = selected_tokens
        n = int(N)
        processed = 0

        while n - processed > 1:
            if cancel_check is not None:
                cancel_check()
            chunk = min(step_size, n - processed - 1)
            model(prompt[processed : processed + chunk][None], cache=cache)
            mx.eval([c.state for c in cache])
            processed += chunk
            mx.clear_cache()

        # Last token → logits
        if cancel_check is not None:
            cancel_check()
        logits = model(prompt[processed:][None], cache=cache)
        mx.eval(logits)

    finally:
        # Replace position-mapped RoPE with offset-adjusted RoPE for decode.
        # Skipped for architectures without RoPE (e.g. Nemotron-H).
        #
        # Total prompt length = position_offset + M (prefix + current tokens).
        # After prefill, cache offset = cache_start + N.
        # Decode needs RoPE position = total_len + i, cache gives offset = cache_start + N + i.
        # Adjustment = total_len - (cache_start + N) = position_offset + M - cache_start - N.
        # When cache_start == position_offset (normal case): adjustment = M - N.
        if has_rope:
            total_prompt_len = position_offset + M
            final_cache_offset = cache_start + N
            adjustment = int(total_prompt_len) - int(final_cache_offset)
            for layer_idx, layer in attn_layers:
                attn, original = original_ropes[layer_idx]
                if adjustment > 0:
                    _set_rope(attn, _OffsetAdjustedRoPE(original, adjustment))
                else:
                    _set_rope(attn, original)

    return logits

vllm_mlx.specprefill.cleanup_rope

cleanup_rope(model)

Restore original RoPE on all attention layers.

Call this after generation is complete to remove _OffsetAdjustedRoPE wrappers installed by sparse_prefill(). No-op for architectures without RoPE (e.g. Nemotron-H).

Source code in vllm_mlx/specprefill.py

def cleanup_rope(model):
    """Restore original RoPE on all attention layers.

    Call this after generation is complete to remove _OffsetAdjustedRoPE
    wrappers installed by sparse_prefill(). No-op for architectures
    without RoPE (e.g. Nemotron-H).
    """
    for _, layer in _find_attention_layers(model):
        attn = _get_attn_module(layer)
        if attn is None:
            continue
        rope = _get_rope(attn)
        if rope is None:
            continue
        if isinstance(rope, (_OffsetAdjustedRoPE, _PositionMappedRoPE)):
            _set_rope(attn, rope._original)

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.specprefill._AttentionCapture · class

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.

vllm_mlx.specprefill._AttentionCapture.__init__ · method

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.

vllm_mlx.specprefill._AttentionCapture.__call__ · method

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.

vllm_mlx.specprefill._AttentionCapture.__getattr__ · method

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.

vllm_mlx.specprefill._qwen35_extract_queries · function

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.

vllm_mlx.specprefill._llama_extract_queries · function

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.

vllm_mlx.specprefill._nemotron_h_extract_queries · function

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.

vllm_mlx.specprefill._patch_attention_for_capture · function

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.

vllm_mlx.specprefill._unpatch_attention_capture · function

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.

vllm_mlx.specprefill._prefill_draft · function

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.

vllm_mlx.specprefill._lookahead_decode · function

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.

vllm_mlx.specprefill._avg_pool1d · function

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.

vllm_mlx.specprefill._compute_importance · function

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.

vllm_mlx.specprefill.score_tokens · function

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.

vllm_mlx.specprefill.select_chunks · function

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.

vllm_mlx.specprefill.select_chunks._selected_token_count · nested function

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.

vllm_mlx.specprefill.manual_rope · function

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.

vllm_mlx.specprefill.manual_rope_with_freqs · function

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.

vllm_mlx.specprefill._PositionMappedRoPE · class

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.

vllm_mlx.specprefill._PositionMappedRoPE.__init__ · method

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.

vllm_mlx.specprefill._PositionMappedRoPE.__call__ · method

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.

vllm_mlx.specprefill._OffsetAdjustedRoPE · class

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.

vllm_mlx.specprefill._OffsetAdjustedRoPE.__init__ · method

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.

vllm_mlx.specprefill._OffsetAdjustedRoPE.__call__ · method

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.

vllm_mlx.specprefill._get_dims · function

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.

vllm_mlx.specprefill._get_pre_scale · function

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.

vllm_mlx.specprefill._find_attention_layers · function

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.

vllm_mlx.specprefill._get_attn_module · function

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.

vllm_mlx.specprefill._get_rope · function

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.

vllm_mlx.specprefill._set_rope · function

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.

vllm_mlx.specprefill._set_attn_module · function

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.

vllm_mlx.specprefill._build_layer_to_cache_map · function

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.

vllm_mlx.specprefill.sparse_prefill · function

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.

vllm_mlx.specprefill.cleanup_rope · function

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.

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	class	_AttentionCapture(original, buf_idx, query_buffer, query_extractor = None)	Wrapper that captures post-RoPE query vectors and delegates to original.	#L53-L73
_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
_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
_AttentionCapture.__getattr__	method	_AttentionCapture.__getattr__(name) -> not annotated	Method _AttentionCapture.__getattr__ calls getattr; returns getattr(self._original, name).	#L72-L73
_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
_llama_extract_queries	function	_llama_extract_queries(attn, x, cache = None) -> not annotated	Extract post-RoPE queries from standard transformer attention.	#L95-L113
_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
_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
_unpatch_attention_capture	function	_unpatch_attention_capture(model, originals) -> not annotated	Restore original attention modules after capture.	#L152-L155
_prefill_draft	function	_prefill_draft(model, tokens, cache, step_size = 2048, cancel_check = None) -> not annotated	Prefill prompt tokens into cache.	#L158-L175
_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
_avg_pool1d	function	_avg_pool1d(x, kernel_size) -> not annotated	1D average pooling along last axis via prefix-sum.	#L207-L223
_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
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
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
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
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
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
_PositionMappedRoPE	class	_PositionMappedRoPE(original_rope, all_positions, cache_start = 0)	Wraps a RoPE module to apply rotation at non-contiguous positions.	#L536-L571
_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
_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
_OffsetAdjustedRoPE	class	_OffsetAdjustedRoPE(original_rope, adjustment)	Wraps a RoPE module to add a constant offset for decode after sparse prefill.	#L574-L590
_OffsetAdjustedRoPE.__init__	method	_OffsetAdjustedRoPE.__init__(original_rope, adjustment) -> not annotated	Method _OffsetAdjustedRoPE.__init__ updates self._original, self._adjustment.	#L585-L587
_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
_get_dims	function	_get_dims(rope_module) -> not annotated	Extract rotary dimensions from any RoPE variant.	#L598-L603
_get_pre_scale	function	_get_pre_scale(rope_module) -> not annotated	Extract pre-scale factor from custom RoPE variants (SuScaled, Yarn).	#L606-L612
_find_attention_layers	function	_find_attention_layers(model) -> not annotated	Find all full-attention layers across architectures.	#L615-L630
_get_attn_module	function	_get_attn_module(layer) -> not annotated	Get the attention module from a layer (self_attn or mixer).	#L633-L639
_get_rope	function	_get_rope(attn) -> not annotated	Get the RoPE module from an attention layer, or None.	#L642-L647
_set_rope	function	_set_rope(attn, rope_module) -> not annotated	Set the RoPE module on an attention layer.	#L650-L655
_set_attn_module	function	_set_attn_module(layer, module) -> not annotated	Set the attention module on a layer (self_attn or mixer).	#L658-L663
_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
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
cleanup_rope	function	cleanup_rope(model) -> not annotated	Restore original RoPE on all attention layers.	#L830-L845