`vllm_mlx.patches.gemma4_mllm`¶

Runtime patch for mlx-vlm Gemma 4 Attention to trim oversized masks.

View the complete module source at #L1-L98.

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.patches.gemma4_mllm ¶

Runtime patch for mlx-vlm Gemma 4 Attention to trim oversized masks.

mlx-vlm 0.5.0's stock Gemma 4 attention assumes the mask's last dim matches keys.shape[-2] exactly. vllm-mlx's BatchedEngine MLLM path (continuous batching) sometimes passes a mask sized for the max sequence in the batch while a specific layer's keys end up shorter — sliding-window layers cap keys at window=512, the mask is built once for the full prompt. Without a trim, scaled_dot_product_attention sees a shape mismatch.

That mask trim is the only behavior this patch adds; everything else mirrors mlx-vlm 0.5.0 verbatim (signature, return shape, offset handling). The previous reason for this patch — BatchKVCache's in-place += on cache.offset corrupting RoPE — is now handled upstream: mlx-vlm 0.5.0 line 223 does offset = mx.array(cache.offset) if cache is not None else 0 which is a defensive copy. (Confirmed in review of PR #564.)

vllm_mlx.patches.gemma4_mllm.logger `module-attribute` ¶

logger = logging.getLogger(__name__)

vllm_mlx.patches.gemma4_mllm.patch_gemma4_attention_for_batching ¶

patch_gemma4_attention_for_batching() -> bool

Patch Gemma 4 Attention.call to trim oversized masks.

Otherwise mirrors mlx-vlm 0.5.0 upstream verbatim. Returns True if applied, False if mlx-vlm is not installed or Gemma 4 unavailable.

Source code in vllm_mlx/patches/gemma4_mllm.py

def patch_gemma4_attention_for_batching() -> bool:
    """Patch Gemma 4 Attention.__call__ to trim oversized masks.

    Otherwise mirrors mlx-vlm 0.5.0 upstream verbatim. Returns True if
    applied, False if mlx-vlm is not installed or Gemma 4 unavailable.
    """
    try:
        from mlx_vlm.models.gemma4.language import Attention as Gemma4Attention
        from mlx_vlm.models.base import scaled_dot_product_attention
    except ImportError:
        logger.debug("[Gemma4 patch] mlx-vlm Gemma4 module not available")
        return False

    if getattr(Gemma4Attention, "_batch_patched", False):
        logger.debug("[Gemma4 patch] Already patched")
        return True

    def _patched_call(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
        shared_kv: Optional[tuple] = None,
        offset: Optional[Any] = None,
    ) -> Any:
        B, L, _ = x.shape

        queries = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim)
        queries = self.q_norm(queries)

        if shared_kv is not None:
            keys, values = shared_kv
        else:
            keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)

            if self.use_k_eq_v:
                values = keys
            else:
                values = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)

            offset = mx.array(cache.offset) if cache is not None else 0

            keys = self.k_norm(keys)
            keys = keys.transpose(0, 2, 1, 3)
            keys = self.rope(keys, offset=offset)

            values = self.v_norm(values)
            values = values.transpose(0, 2, 1, 3)

            if cache is not None:
                keys, values = cache.update_and_fetch(keys, values)

        queries = queries.transpose(0, 2, 1, 3)
        queries = self.rope(queries, offset=offset)

        # Only addition vs upstream: trim mask if longer than keys.
        if mask is not None and isinstance(mask, mx.array):
            if mask.shape[-1] != keys.shape[-2]:
                mask = mask[..., -keys.shape[-2] :]

        output = scaled_dot_product_attention(
            queries, keys, values, cache=cache, scale=self.scale, mask=mask
        )
        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)

        return self.o_proj(output), (keys, values), offset

    Gemma4Attention.__call__ = _patched_call
    Gemma4Attention._batch_patched = True
    logger.info("[Gemma4 patch] Attention patched (mask trim for BatchedEngine)")
    return True

Complete contract reference¶

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.patches.gemma4_mllm.patch_gemma4_attention_for_batching · function

vllm_mlx.patches.gemma4_mllm.patch_gemma4_attention_for_batching() -> bool

Patch Gemma 4 Attention.call to trim oversized masks.

Parameters

This callable has no explicit inputs.

Returns

Type: bool
Direct return expressions: False; True

Exceptions and behavior

Function patch_gemma4_attention_for_batching calls logger.debug, getattr, logger.info; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L28-L98.

vllm_mlx.patches.gemma4_mllm.patch_gemma4_attention_for_batching._patched_call · nested function

vllm_mlx.patches.gemma4_mllm.patch_gemma4_attention_for_batching._patched_call(x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, shared_kv: Optional[tuple] = None, offset: Optional[Any] = None) -> Any

Nested Function patch_gemma4_attention_for_batching._patched_call calls self.q_proj(x).reshape, self.q_proj, self.q_norm, self.k_proj(x).reshape; returns (self.o_proj(output), (keys, values), offset).

Parameters

Name	Type	Required	Default	Description
`x`	`mx.array`	`yes`	`none`	Required positional or keyword input.
`mask`	`Optional[mx.array]`	`no`	`None`	Optional positional or keyword input; defaults to `None`.
`cache`	`Optional[Any]`	`no`	`None`	Optional positional or keyword input; defaults to `None`.
`shared_kv`	`Optional[tuple]`	`no`	`None`	Optional positional or keyword input; defaults to `None`.
`offset`	`Optional[Any]`	`no`	`None`	Optional positional or keyword input; defaults to `None`.

Returns

Type: Any
Direct return expressions: (self.o_proj(output), (keys, values), offset)

Exceptions and behavior

Nested Function patch_gemma4_attention_for_batching._patched_call calls self.q_proj(x).reshape, self.q_proj, self.q_norm, self.k_proj(x).reshape; returns (self.o_proj(output), (keys, values), offset). No direct raise statement appears in this definition.

View source #L45-L93.

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
`patch_gemma4_attention_for_batching`	function	`patch_gemma4_attention_for_batching() -> bool`	Patch Gemma 4 Attention.call to trim oversized masks.	#L28-L98
`patch_gemma4_attention_for_batching._patched_call`	nested function	`patch_gemma4_attention_for_batching._patched_call(x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, shared_kv: Optional[tuple] = None, offset: Optional[Any] = None) -> Any`	Nested Function `patch_gemma4_attention_for_batching._patched_call` calls `self.q_proj(x).reshape`, `self.q_proj`, `self.q_norm`, `self.k_proj(x).reshape`; returns `(self.o_proj(output), (keys, values), offset)`.	#L45-L93

vllm_mlx.patches.gemma4_mllm¶