vllm_mlx.rerank_forward

MLX forward pass for BERT-family sequence classification models.

View the complete module source at #L1-L265.

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

MLX forward pass for BERT-family sequence classification models.

Implements a from-weights forward pass for cross-encoder rerankers that use the standard BERT/XLM-RoBERTa architecture with a classification head. This avoids pulling in the full transformers modeling stack at inference time — only the tokenizer is needed from transformers.

vllm_mlx.rerank_forward.classifier_forward

classifier_forward(input_ids: array, attention_mask: array, weights: dict[str, array], config: dict) -> array

Run a BERT-family classifier forward pass on MLX.

Parameters:

• input_ids (array) –

  (batch, seq_len) token IDs.

• attention_mask (array) –

  (batch, seq_len) attention mask (1=attend, 0=pad).

• weights (dict[str, array]) –

  Dict mapping weight name -> mx.array.

• config (dict) –

  Model config dict (from config.json).

Returns:

• logits ( array ) –

  (batch, num_labels) classification logits.

Source code in vllm_mlx/rerank_forward.py

def classifier_forward(
    input_ids: mx.array,
    attention_mask: mx.array,
    weights: dict[str, mx.array],
    config: dict,
) -> mx.array:
    """
    Run a BERT-family classifier forward pass on MLX.

    Args:
        input_ids: (batch, seq_len) token IDs.
        attention_mask: (batch, seq_len) attention mask (1=attend, 0=pad).
        weights: Dict mapping weight name -> mx.array.
        config: Model config dict (from config.json).

    Returns:
        logits: (batch, num_labels) classification logits.
    """
    hidden_size = config["hidden_size"]
    num_heads = config["num_attention_heads"]
    num_layers = config["num_hidden_layers"]
    num_labels = config.get("num_labels", 1)
    eps = config.get("layer_norm_eps", 1e-12)

    head_dim = hidden_size // num_heads

    # Detect weight prefix (bert.* vs roberta.* vs xlm-roberta.*)
    prefix = _detect_prefix(weights)

    # --- Embeddings ---
    word_emb = weights[f"{prefix}.embeddings.word_embeddings.weight"]
    pos_emb = weights[f"{prefix}.embeddings.position_embeddings.weight"]
    tok_type_emb = weights[f"{prefix}.embeddings.token_type_embeddings.weight"]
    ln_w = weights[f"{prefix}.embeddings.LayerNorm.weight"]
    ln_b = weights[f"{prefix}.embeddings.LayerNorm.bias"]

    batch_size, seq_len = input_ids.shape
    position_ids = _position_ids_for_config(config, input_ids, attention_mask)
    token_type_ids = mx.zeros_like(input_ids)

    hidden = word_emb[input_ids] + pos_emb[position_ids] + tok_type_emb[token_type_ids]
    hidden = _layer_norm(hidden, ln_w, ln_b, eps)

    # --- Encoder layers ---
    # Build causal-free attention mask: (batch, 1, 1, seq_len)
    if attention_mask is not None:
        ext_mask = attention_mask[:, None, None, :].astype(mx.float32)
        ext_mask = (1.0 - ext_mask) * -1e9
    else:
        ext_mask = None

    for i in range(num_layers):
        lp = f"{prefix}.encoder.layer.{i}"
        hidden = _encoder_layer(
            hidden, ext_mask, weights, lp, num_heads, head_dim, eps, config
        )

    # --- Pooler (CLS token) ---
    cls_hidden = hidden[:, 0, :]  # (batch, hidden_size)
    pooler_w = weights.get(f"{prefix}.pooler.dense.weight")
    pooler_b = weights.get(f"{prefix}.pooler.dense.bias")
    if pooler_w is not None:
        pooled = mx.tanh(cls_hidden @ pooler_w.T + pooler_b)
    else:
        pooled = cls_hidden

    # --- Classifier head ---
    logits = _classification_head_forward(pooled, weights)

    return logits

vllm_mlx.rerank_forward._position_ids_for_config

_position_ids_for_config(config: dict, input_ids: array, attention_mask: array | None) -> array

Build BERT or RoBERTa-family absolute position IDs.

Source code in vllm_mlx/rerank_forward.py

def _position_ids_for_config(
    config: dict,
    input_ids: mx.array,
    attention_mask: mx.array | None,
) -> mx.array:
    """Build BERT or RoBERTa-family absolute position IDs."""
    _, seq_len = input_ids.shape
    model_type = str(config.get("model_type", "")).lower()
    if model_type not in {"roberta", "xlm-roberta", "xlm_roberta"}:
        return mx.arange(seq_len)[None, :]

    padding_idx = int(config.get("pad_token_id", 1))
    if attention_mask is None:
        return mx.arange(padding_idx + 1, seq_len + padding_idx + 1)[None, :]

    mask = attention_mask.astype(mx.int32)
    positions = mx.cumsum(mask, axis=1) * mask + padding_idx
    return positions.astype(mx.int32)

vllm_mlx.rerank_forward._classification_head_forward

_classification_head_forward(pooled: array, weights: dict[str, array]) -> array

Run BERT flat or XLM-RoBERTa two-layer sequence-classification head.

Source code in vllm_mlx/rerank_forward.py

def _classification_head_forward(
    pooled: mx.array,
    weights: dict[str, mx.array],
) -> mx.array:
    """Run BERT flat or XLM-RoBERTa two-layer sequence-classification head."""
    if "classifier.dense.weight" in weights:
        hidden = pooled @ weights["classifier.dense.weight"].T
        hidden = hidden + weights["classifier.dense.bias"]
        hidden = mx.tanh(hidden)
        return (
            hidden @ weights["classifier.out_proj.weight"].T
            + weights["classifier.out_proj.bias"]
        )

    return pooled @ weights["classifier.weight"].T + weights["classifier.bias"]

vllm_mlx.rerank_forward._detect_prefix

_detect_prefix(weights: dict) -> str

Detect the model weight prefix (bert, roberta, xlm-roberta).

Source code in vllm_mlx/rerank_forward.py

def _detect_prefix(weights: dict) -> str:
    """Detect the model weight prefix (bert, roberta, xlm-roberta)."""
    for key in weights:
        if key.startswith("bert."):
            return "bert"
        if key.startswith("roberta."):
            return "roberta"
        if key.startswith("xlm-roberta."):
            return "xlm-roberta"
    # Default to bert
    return "bert"

vllm_mlx.rerank_forward._layer_norm

_layer_norm(x: array, weight: array, bias: array, eps: float) -> array

Apply layer normalization.

Source code in vllm_mlx/rerank_forward.py

def _layer_norm(x: mx.array, weight: mx.array, bias: mx.array, eps: float) -> mx.array:
    """Apply layer normalization."""
    mean = mx.mean(x, axis=-1, keepdims=True)
    var = mx.var(x, axis=-1, keepdims=True)
    return weight * (x - mean) / mx.sqrt(var + eps) + bias

vllm_mlx.rerank_forward._encoder_layer

_encoder_layer(hidden: array, ext_mask: array | None, weights: dict, prefix: str, num_heads: int, head_dim: int, eps: float, config: dict) -> array

Run one BERT encoder layer (self-attention + FFN).

Source code in vllm_mlx/rerank_forward.py

def _encoder_layer(
    hidden: mx.array,
    ext_mask: mx.array | None,
    weights: dict,
    prefix: str,
    num_heads: int,
    head_dim: int,
    eps: float,
    config: dict,
) -> mx.array:
    """Run one BERT encoder layer (self-attention + FFN)."""
    hidden_size = num_heads * head_dim

    # --- Self-attention ---
    q_w = weights[f"{prefix}.attention.self.query.weight"]
    q_b = weights[f"{prefix}.attention.self.query.bias"]
    k_w = weights[f"{prefix}.attention.self.key.weight"]
    k_b = weights[f"{prefix}.attention.self.key.bias"]
    v_w = weights[f"{prefix}.attention.self.value.weight"]
    v_b = weights[f"{prefix}.attention.self.value.bias"]

    batch_size, seq_len, _ = hidden.shape

    q = (
        (hidden @ q_w.T + q_b)
        .reshape(batch_size, seq_len, num_heads, head_dim)
        .transpose(0, 2, 1, 3)
    )
    k = (
        (hidden @ k_w.T + k_b)
        .reshape(batch_size, seq_len, num_heads, head_dim)
        .transpose(0, 2, 1, 3)
    )
    v = (
        (hidden @ v_w.T + v_b)
        .reshape(batch_size, seq_len, num_heads, head_dim)
        .transpose(0, 2, 1, 3)
    )

    scale = head_dim**-0.5
    attn_scores = (q @ k.transpose(0, 1, 3, 2)) * scale  # (batch, heads, seq, seq)

    if ext_mask is not None:
        attn_scores = attn_scores + ext_mask

    attn_probs = mx.softmax(attn_scores, axis=-1)
    attn_out = (
        (attn_probs @ v).transpose(0, 2, 1, 3).reshape(batch_size, seq_len, hidden_size)
    )

    # Attention output projection + residual + LayerNorm
    ao_w = weights[f"{prefix}.attention.output.dense.weight"]
    ao_b = weights[f"{prefix}.attention.output.dense.bias"]
    ao_ln_w = weights[f"{prefix}.attention.output.LayerNorm.weight"]
    ao_ln_b = weights[f"{prefix}.attention.output.LayerNorm.bias"]

    attn_out = attn_out @ ao_w.T + ao_b
    hidden = _layer_norm(hidden + attn_out, ao_ln_w, ao_ln_b, eps)

    # --- FFN ---
    inter_w = weights[f"{prefix}.intermediate.dense.weight"]
    inter_b = weights[f"{prefix}.intermediate.dense.bias"]
    out_w = weights[f"{prefix}.output.dense.weight"]
    out_b = weights[f"{prefix}.output.dense.bias"]
    out_ln_w = weights[f"{prefix}.output.LayerNorm.weight"]
    out_ln_b = weights[f"{prefix}.output.LayerNorm.bias"]

    intermediate = hidden @ inter_w.T + inter_b
    intermediate = _apply_hidden_activation(intermediate, config)
    ffn_out = intermediate @ out_w.T + out_b
    hidden = _layer_norm(hidden + ffn_out, out_ln_w, out_ln_b, eps)

    return hidden

vllm_mlx.rerank_forward._gelu

_gelu(x: array) -> array

GELU activation (exact form).

Source code in vllm_mlx/rerank_forward.py

def _gelu(x: mx.array) -> mx.array:
    """GELU activation (exact form)."""
    return nn.gelu(x)

vllm_mlx.rerank_forward._gelu_new

_gelu_new(x: array) -> array

BERT GELU approximation used by transformers gelu_new.

Source code in vllm_mlx/rerank_forward.py

def _gelu_new(x: mx.array) -> mx.array:
    """BERT GELU approximation used by transformers gelu_new."""
    return 0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x**3)))

vllm_mlx.rerank_forward._relu

_relu(x: array) -> array

ReLU activation.

Source code in vllm_mlx/rerank_forward.py

def _relu(x: mx.array) -> mx.array:
    """ReLU activation."""
    return mx.maximum(x, 0)

vllm_mlx.rerank_forward._silu

_silu(x: array) -> array

SiLU/swish activation.

Source code in vllm_mlx/rerank_forward.py

def _silu(x: mx.array) -> mx.array:
    """SiLU/swish activation."""
    return x * mx.sigmoid(x)

vllm_mlx.rerank_forward._apply_hidden_activation

_apply_hidden_activation(x: array, config: dict) -> array

Apply the configured encoder hidden activation.

The MLX reranker forward pass targets standard BERT/XLM-RoBERTa-style sequence classifiers. Configs that request an activation outside that supported contract fail explicitly instead of silently using GELU.

Source code in vllm_mlx/rerank_forward.py

def _apply_hidden_activation(x: mx.array, config: dict) -> mx.array:
    """
    Apply the configured encoder hidden activation.

    The MLX reranker forward pass targets standard BERT/XLM-RoBERTa-style
    sequence classifiers. Configs that request an activation outside that
    supported contract fail explicitly instead of silently using GELU.
    """
    hidden_act = config.get("hidden_act", "gelu")
    if isinstance(hidden_act, dict):
        hidden_act = hidden_act.get("type", "gelu")
    hidden_act = str(hidden_act).lower()

    if hidden_act == "gelu":
        return _gelu(x)
    if hidden_act in {"gelu_new", "gelu_fast"}:
        return _gelu_new(x)
    if hidden_act == "relu":
        return _relu(x)
    if hidden_act in {"silu", "swish"}:
        return _silu(x)

    raise ValueError(
        f"Unsupported reranker hidden_act '{hidden_act}'. "
        "Supported activations are gelu, gelu_new/gelu_fast, relu, and silu/swish."
    )

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.rerank_forward.classifier_forward · function

vllm_mlx.rerank_forward.classifier_forward(input_ids: mx.array, attention_mask: mx.array, weights: dict[str, mx.array], config: dict) -> mx.array

Run a BERT-family classifier forward pass on MLX.

Parameters

Name	Type	Required	Default	Description
input_ids	mx.array	yes	none	(batch, seq_len) token IDs.
attention_mask	mx.array	yes	none	(batch, seq_len) attention mask (1=attend, 0=pad).
weights	dict[str, mx.array]	yes	none	Dict mapping weight name -> mx.array.
config	dict	yes	none	Model config dict (from config.json).

Returns

• Type: mx.array
• Direct return expressions: logits

Exceptions and behavior

Function classifier_forward calls config.get, _detect_prefix, _position_ids_for_config, mx.zeros_like; returns logits. No direct raise statement appears in this definition.

View source #L16-L85.

vllm_mlx.rerank_forward._position_ids_for_config · function

vllm_mlx.rerank_forward._position_ids_for_config(config: dict, input_ids: mx.array, attention_mask: mx.array | None) -> mx.array

Build BERT or RoBERTa-family absolute position IDs.

Parameters

Name	Type	Required	Default	Description
config	dict	yes	none	Required positional or keyword input.
input_ids	mx.array	yes	none	Required positional or keyword input.
attention_mask	mx.array \| None	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: mx.arange(seq_len)[None, :]; mx.arange(padding_idx + 1, seq_len + padding_idx + 1)[None, :]; positions.astype(mx.int32)

Exceptions and behavior

Function _position_ids_for_config calls str(config.get('model_type', '')).lower, str, config.get, mx.arange; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L88-L105.

vllm_mlx.rerank_forward._classification_head_forward · function

vllm_mlx.rerank_forward._classification_head_forward(pooled: mx.array, weights: dict[str, mx.array]) -> mx.array

Run BERT flat or XLM-RoBERTa two-layer sequence-classification head.

Parameters

Name	Type	Required	Default	Description
pooled	mx.array	yes	none	Required positional or keyword input.
weights	dict[str, mx.array]	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: hidden @ weights['classifier.out_proj.weight'].T + weights['classifier.out_proj.bias']; pooled @ weights['classifier.weight'].T + weights['classifier.bias']

Exceptions and behavior

Function _classification_head_forward calls mx.tanh; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L108-L122.

vllm_mlx.rerank_forward._detect_prefix · function

vllm_mlx.rerank_forward._detect_prefix(weights: dict) -> str

Detect the model weight prefix (bert, roberta, xlm-roberta).

Parameters

Name	Type	Required	Default	Description
weights	dict	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: 'bert'; 'roberta'; 'xlm-roberta'

Exceptions and behavior

Function _detect_prefix calls key.startswith; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L125-L135.

vllm_mlx.rerank_forward._layer_norm · function

vllm_mlx.rerank_forward._layer_norm(x: mx.array, weight: mx.array, bias: mx.array, eps: float) -> mx.array

Apply layer normalization.

Parameters

Name	Type	Required	Default	Description
x	mx.array	yes	none	Required positional or keyword input.
weight	mx.array	yes	none	Required positional or keyword input.
bias	mx.array	yes	none	Required positional or keyword input.
eps	float	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: weight * (x - mean) / mx.sqrt(var + eps) + bias

Exceptions and behavior

Function _layer_norm calls mx.mean, mx.var, mx.sqrt; returns weight * (x - mean) / mx.sqrt(var + eps) + bias. No direct raise statement appears in this definition.

View source #L138-L142.

vllm_mlx.rerank_forward._encoder_layer · function

vllm_mlx.rerank_forward._encoder_layer(hidden: mx.array, ext_mask: mx.array | None, weights: dict, prefix: str, num_heads: int, head_dim: int, eps: float, config: dict) -> mx.array

Run one BERT encoder layer (self-attention + FFN).

Parameters

Name	Type	Required	Default	Description
hidden	mx.array	yes	none	Required positional or keyword input.
ext_mask	mx.array \| None	yes	none	Required positional or keyword input.
weights	dict	yes	none	Required positional or keyword input.
prefix	str	yes	none	Required positional or keyword input.
num_heads	int	yes	none	Required positional or keyword input.
head_dim	int	yes	none	Required positional or keyword input.
eps	float	yes	none	Required positional or keyword input.
config	dict	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: hidden

Exceptions and behavior

Function _encoder_layer calls (hidden @ q_w.T + q_b).reshape(batch_size, seq_len, num_heads, head_dim).transpose, (hidden @ q_w.T + q_b).reshape, (hidden @ k_w.T + k_b).reshape(batch_size, seq_len, num_heads, head_dim).transpose, (hidden @ k_w.T + k_b).reshape; returns hidden. No direct raise statement appears in this definition.

View source #L145-L217.

vllm_mlx.rerank_forward._gelu · function

vllm_mlx.rerank_forward._gelu(x: mx.array) -> mx.array

GELU activation (exact form).

Parameters

Name	Type	Required	Default	Description
x	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: nn.gelu(x)

Exceptions and behavior

Function _gelu calls nn.gelu; returns nn.gelu(x). No direct raise statement appears in this definition.

View source #L220-L222.

vllm_mlx.rerank_forward._gelu_new · function

vllm_mlx.rerank_forward._gelu_new(x: mx.array) -> mx.array

BERT GELU approximation used by transformers gelu_new.

Parameters

Name	Type	Required	Default	Description
x	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: 0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x ** 3)))

Exceptions and behavior

Function _gelu_new calls mx.tanh; returns 0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x ** 3))). No direct raise statement appears in this definition.

View source #L225-L227.

vllm_mlx.rerank_forward._relu · function

vllm_mlx.rerank_forward._relu(x: mx.array) -> mx.array

ReLU activation.

Parameters

Name	Type	Required	Default	Description
x	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: mx.maximum(x, 0)

Exceptions and behavior

Function _relu calls mx.maximum; returns mx.maximum(x, 0). No direct raise statement appears in this definition.

View source #L230-L232.

vllm_mlx.rerank_forward._silu · function

vllm_mlx.rerank_forward._silu(x: mx.array) -> mx.array

SiLU/swish activation.

Parameters

Name	Type	Required	Default	Description
x	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: x * mx.sigmoid(x)

Exceptions and behavior

Function _silu calls mx.sigmoid; returns x * mx.sigmoid(x). No direct raise statement appears in this definition.

View source #L235-L237.

vllm_mlx.rerank_forward._apply_hidden_activation · function

vllm_mlx.rerank_forward._apply_hidden_activation(x: mx.array, config: dict) -> mx.array

Apply the configured encoder hidden activation.

Parameters

Name	Type	Required	Default	Description
x	mx.array	yes	none	Required positional or keyword input.
config	dict	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: _gelu(x); _gelu_new(x); _relu(x); _silu(x)

Exceptions and behavior

Function _apply_hidden_activation calls config.get, isinstance, hidden_act.get, str(hidden_act).lower; can raise ValueError; has 4 explicit return paths. Directly raised exceptions: ValueError.

View source #L240-L265.

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
classifier_forward	function	classifier_forward(input_ids: mx.array, attention_mask: mx.array, weights: dict[str, mx.array], config: dict) -> mx.array	Run a BERT-family classifier forward pass on MLX.	#L16-L85
_position_ids_for_config	function	_position_ids_for_config(config: dict, input_ids: mx.array, attention_mask: mx.array \| None) -> mx.array	Build BERT or RoBERTa-family absolute position IDs.	#L88-L105
_classification_head_forward	function	_classification_head_forward(pooled: mx.array, weights: dict[str, mx.array]) -> mx.array	Run BERT flat or XLM-RoBERTa two-layer sequence-classification head.	#L108-L122
_detect_prefix	function	_detect_prefix(weights: dict) -> str	Detect the model weight prefix (bert, roberta, xlm-roberta).	#L125-L135
_layer_norm	function	_layer_norm(x: mx.array, weight: mx.array, bias: mx.array, eps: float) -> mx.array	Apply layer normalization.	#L138-L142
_encoder_layer	function	_encoder_layer(hidden: mx.array, ext_mask: mx.array \| None, weights: dict, prefix: str, num_heads: int, head_dim: int, eps: float, config: dict) -> mx.array	Run one BERT encoder layer (self-attention + FFN).	#L145-L217
_gelu	function	_gelu(x: mx.array) -> mx.array	GELU activation (exact form).	#L220-L222
_gelu_new	function	_gelu_new(x: mx.array) -> mx.array	BERT GELU approximation used by transformers gelu_new.	#L225-L227
_relu	function	_relu(x: mx.array) -> mx.array	ReLU activation.	#L230-L232
_silu	function	_silu(x: mx.array) -> mx.array	SiLU/swish activation.	#L235-L237
_apply_hidden_activation	function	_apply_hidden_activation(x: mx.array, config: dict) -> mx.array	Apply the configured encoder hidden activation.	#L240-L265