vllm_mlx.attention

MLX Attention Backend for vLLM.

View the complete module source at #L1-L245.

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

MLX Attention Backend for vLLM.

This module provides an attention backend that uses MLX's native attention implementation, optimized for Apple Silicon.

vllm_mlx.attention.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.attention.MLXAttentionMetadata dataclass

MLXAttentionMetadata(seq_lens: list[int], max_seq_len: int, num_prefill_tokens: int = 0, num_decode_tokens: int = 0, block_tables: Any | None = None, slot_mapping: Any | None = None)

Metadata for MLX attention computation.

vllm_mlx.attention.MLXAttentionMetadata.seq_lens instance-attribute

seq_lens: list[int]

vllm_mlx.attention.MLXAttentionMetadata.max_seq_len instance-attribute

max_seq_len: int

vllm_mlx.attention.MLXAttentionMetadata.num_prefill_tokens class-attribute instance-attribute

num_prefill_tokens: int = 0

vllm_mlx.attention.MLXAttentionMetadata.num_decode_tokens class-attribute instance-attribute

num_decode_tokens: int = 0

vllm_mlx.attention.MLXAttentionMetadata.block_tables class-attribute instance-attribute

block_tables: Any | None = None

vllm_mlx.attention.MLXAttentionMetadata.slot_mapping class-attribute instance-attribute

slot_mapping: Any | None = None

vllm_mlx.attention.MLXAttentionBackend

Attention backend using MLX's native attention.

MLX provides optimized attention implementations that run on Apple Silicon's GPU via Metal. This backend wraps those implementations for use with vLLM.

Note: mlx-lm handles attention internally, so this backend primarily serves as a compatibility layer.

vllm_mlx.attention.MLXAttentionBackend.get_name staticmethod

get_name() -> str

Return backend name.

Source code in vllm_mlx/attention.py

@staticmethod
def get_name() -> str:
    """Return backend name."""
    return "MLX"

vllm_mlx.attention.MLXAttentionBackend.get_impl_cls staticmethod

get_impl_cls() -> type

Return the implementation class.

Source code in vllm_mlx/attention.py

@staticmethod
def get_impl_cls() -> type:
    """Return the implementation class."""
    return MLXAttentionImpl

vllm_mlx.attention.MLXAttentionBackend.get_metadata_cls staticmethod

get_metadata_cls() -> type

Return the metadata class.

Source code in vllm_mlx/attention.py

@staticmethod
def get_metadata_cls() -> type:
    """Return the metadata class."""
    return MLXAttentionMetadata

vllm_mlx.attention.MLXAttentionBackend.get_kv_cache_shape staticmethod

get_kv_cache_shape(num_blocks: int, block_size: int, num_kv_heads: int, head_size: int) -> tuple[int, ...]

Get the shape of KV cache.

Parameters:

• num_blocks (int) –

  Number of cache blocks

• block_size (int) –

  Tokens per block

• num_kv_heads (int) –

  Number of KV attention heads

• head_size (int) –

  Size of each attention head

Returns:

• tuple[int, ...] –

  Shape tuple for KV cache tensor

Source code in vllm_mlx/attention.py

@staticmethod
def get_kv_cache_shape(
    num_blocks: int,
    block_size: int,
    num_kv_heads: int,
    head_size: int,
) -> tuple[int, ...]:
    """
    Get the shape of KV cache.

    Args:
        num_blocks: Number of cache blocks
        block_size: Tokens per block
        num_kv_heads: Number of KV attention heads
        head_size: Size of each attention head

    Returns:
        Shape tuple for KV cache tensor
    """
    # Shape: (num_blocks, block_size, num_kv_heads, head_size)
    return (num_blocks, block_size, num_kv_heads, head_size)

vllm_mlx.attention.MLXAttentionBackend.get_supported_head_sizes staticmethod

get_supported_head_sizes() -> list[int]

Return supported attention head sizes.

Source code in vllm_mlx/attention.py

@staticmethod
def get_supported_head_sizes() -> list[int]:
    """Return supported attention head sizes."""
    return [64, 80, 96, 112, 128, 256]

vllm_mlx.attention.MLXAttentionBackend.validate_configuration staticmethod

validate_configuration(num_heads: int, head_size: int, num_kv_heads: int, dtype: dtype, block_size: int, **kwargs) -> list[str]

Validate attention configuration.

Returns list of error messages (empty if valid).

Source code in vllm_mlx/attention.py

@staticmethod
def validate_configuration(
    num_heads: int,
    head_size: int,
    num_kv_heads: int,
    dtype: "torch.dtype",
    block_size: int,
    **kwargs,
) -> list[str]:
    """
    Validate attention configuration.

    Returns list of error messages (empty if valid).
    """
    errors = []

    if head_size not in MLXAttentionBackend.get_supported_head_sizes():
        errors.append(
            f"Head size {head_size} not in supported sizes: "
            f"{MLXAttentionBackend.get_supported_head_sizes()}"
        )

    return errors

vllm_mlx.attention.MLXAttentionBackend.supports_dtype staticmethod

supports_dtype(dtype: dtype) -> bool

Check if dtype is supported.

Source code in vllm_mlx/attention.py

@staticmethod
def supports_dtype(dtype: "torch.dtype") -> bool:
    """Check if dtype is supported."""
    import torch

    return dtype in [torch.float16, torch.bfloat16, torch.float32]

vllm_mlx.attention.MLXAttentionBackend.supports_block_size staticmethod

supports_block_size(block_size: int) -> bool

Check if block size is supported.

Source code in vllm_mlx/attention.py

@staticmethod
def supports_block_size(block_size: int) -> bool:
    """Check if block size is supported."""
    return block_size in [8, 16, 32]

vllm_mlx.attention.MLXAttentionBackend.supports_attn_type staticmethod

supports_attn_type(attn_type: str) -> bool

Check if attention type is supported.

Source code in vllm_mlx/attention.py

@staticmethod
def supports_attn_type(attn_type: str) -> bool:
    """Check if attention type is supported."""
    return attn_type in ["decoder", "encoder", "encoder_decoder"]

vllm_mlx.attention.MLXAttentionImpl

MLXAttentionImpl(num_heads: int, head_size: int, scale: float, num_kv_heads: int | None = None, alibi_slopes: list[float] | None = None, sliding_window: int | None = None, kv_cache_dtype: str = 'auto', blocksparse_params: dict | None = None, logits_soft_cap: float | None = None, **kwargs)

MLX attention implementation.

This class provides the actual attention computation using MLX. Since mlx-lm handles attention internally during generation, this serves as a compatibility interface.

Initialize MLX attention.

Parameters:

• num_heads (int) –

  Number of attention heads

• head_size (int) –

  Size of each head

• scale (float) –

  Attention scale factor

• num_kv_heads (int | None, default: None ) –

  Number of KV heads (for GQA/MQA)

• alibi_slopes (list[float] | None, default: None ) –

  ALiBi position encoding slopes

• sliding_window (int | None, default: None ) –

  Sliding window attention size

• kv_cache_dtype (str, default: 'auto' ) –

  KV cache data type

• blocksparse_params (dict | None, default: None ) –

  Block-sparse attention params

• logits_soft_cap (float | None, default: None ) –

  Soft cap for logits

Source code in vllm_mlx/attention.py

def __init__(
    self,
    num_heads: int,
    head_size: int,
    scale: float,
    num_kv_heads: int | None = None,
    alibi_slopes: list[float] | None = None,
    sliding_window: int | None = None,
    kv_cache_dtype: str = "auto",
    blocksparse_params: dict | None = None,
    logits_soft_cap: float | None = None,
    **kwargs,
):
    """
    Initialize MLX attention.

    Args:
        num_heads: Number of attention heads
        head_size: Size of each head
        scale: Attention scale factor
        num_kv_heads: Number of KV heads (for GQA/MQA)
        alibi_slopes: ALiBi position encoding slopes
        sliding_window: Sliding window attention size
        kv_cache_dtype: KV cache data type
        blocksparse_params: Block-sparse attention params
        logits_soft_cap: Soft cap for logits
    """
    self.num_heads = num_heads
    self.head_size = head_size
    self.scale = scale
    self.num_kv_heads = num_kv_heads or num_heads
    self.alibi_slopes = alibi_slopes
    self.sliding_window = sliding_window
    self.kv_cache_dtype = kv_cache_dtype
    self.logits_soft_cap = logits_soft_cap

    logger.debug(
        f"MLXAttentionImpl initialized: heads={num_heads}, "
        f"kv_heads={self.num_kv_heads}, head_size={head_size}"
    )

vllm_mlx.attention.MLXAttentionImpl.num_heads instance-attribute

num_heads = num_heads

vllm_mlx.attention.MLXAttentionImpl.head_size instance-attribute

head_size = head_size

vllm_mlx.attention.MLXAttentionImpl.scale instance-attribute

scale = scale

vllm_mlx.attention.MLXAttentionImpl.num_kv_heads instance-attribute

num_kv_heads = num_kv_heads or num_heads

vllm_mlx.attention.MLXAttentionImpl.alibi_slopes instance-attribute

alibi_slopes = alibi_slopes

vllm_mlx.attention.MLXAttentionImpl.sliding_window instance-attribute

sliding_window = sliding_window

vllm_mlx.attention.MLXAttentionImpl.kv_cache_dtype instance-attribute

kv_cache_dtype = kv_cache_dtype

vllm_mlx.attention.MLXAttentionImpl.logits_soft_cap instance-attribute

logits_soft_cap = logits_soft_cap

vllm_mlx.attention.MLXAttentionImpl.forward

forward(query: Any, key: Any, value: Any, kv_cache: Any | None = None, attn_metadata: MLXAttentionMetadata | None = None, output: Any | None = None, **kwargs) -> Any

Compute attention.

Note: In the MLX backend, attention is handled internally by mlx-lm during the generation process. This method is provided for compatibility but may not be called directly.

Parameters:

• query (Any) –

  Query tensor

• key (Any) –

  Key tensor

• value (Any) –

  Value tensor

• kv_cache (Any | None, default: None ) –

  Optional KV cache

• attn_metadata (MLXAttentionMetadata | None, default: None ) –

  Attention metadata

• output (Any | None, default: None ) –

  Optional output buffer

Returns:

• Any –

  Attention output tensor

Source code in vllm_mlx/attention.py

def forward(
    self,
    query: Any,
    key: Any,
    value: Any,
    kv_cache: Any | None = None,
    attn_metadata: MLXAttentionMetadata | None = None,
    output: Any | None = None,
    **kwargs,
) -> Any:
    """
    Compute attention.

    Note: In the MLX backend, attention is handled internally by mlx-lm
    during the generation process. This method is provided for
    compatibility but may not be called directly.

    Args:
        query: Query tensor
        key: Key tensor
        value: Value tensor
        kv_cache: Optional KV cache
        attn_metadata: Attention metadata
        output: Optional output buffer

    Returns:
        Attention output tensor
    """
    try:
        import mlx.core as mx

        # Convert inputs to MLX arrays if needed
        if not isinstance(query, mx.array):
            query = mx.array(query.numpy() if hasattr(query, "numpy") else query)
        if not isinstance(key, mx.array):
            key = mx.array(key.numpy() if hasattr(key, "numpy") else key)
        if not isinstance(value, mx.array):
            value = mx.array(value.numpy() if hasattr(value, "numpy") else value)

        # Use MLX's scaled dot product attention
        # Shape: (batch, seq_len, num_heads, head_size)
        attn_output = mx.fast.scaled_dot_product_attention(
            query,
            key,
            value,
            scale=self.scale,
        )

        return attn_output

    except Exception as e:
        logger.error(f"MLX attention forward failed: {e}")
        raise

vllm_mlx.attention.create_mlx_attention_backend

create_mlx_attention_backend() -> type

Factory function to create MLX attention backend.

Source code in vllm_mlx/attention.py

def create_mlx_attention_backend() -> type:
    """Factory function to create MLX attention backend."""
    return MLXAttentionBackend

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.attention.MLXAttentionMetadata · class

vllm_mlx.attention.MLXAttentionMetadata(seq_lens: list[int], max_seq_len: int, num_prefill_tokens: int = 0, num_decode_tokens: int = 0, block_tables: Any | None = None, slot_mapping: Any | None = None)

Metadata for MLX attention computation.

Parameters

Name	Type	Required	Default	Description
seq_lens	list[int]	yes	none	Required constructor field.
max_seq_len	int	yes	none	Required constructor field.
num_prefill_tokens	int	no	0	Optional constructor field; defaults to 0.
num_decode_tokens	int	no	0	Optional constructor field; defaults to 0.
block_tables	Any \| None	no	None	Optional constructor field; defaults to None.
slot_mapping	Any \| None	no	None	Optional constructor field; defaults to None.

Returns

• Constructs: vllm_mlx.attention.MLXAttentionMetadata

Exceptions and behavior

Class MLXAttentionMetadata declares 0 direct member(s). No direct raise statement appears in this definition.

View source #L20-L39.

vllm_mlx.attention.MLXAttentionBackend · class

vllm_mlx.attention.MLXAttentionBackend()

Attention backend using MLX's native attention.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.attention.MLXAttentionBackend

Exceptions and behavior

Class MLXAttentionBackend declares 9 direct member(s). No direct raise statement appears in this definition.

View source #L42-L135.

vllm_mlx.attention.MLXAttentionBackend.get_name · method

vllm_mlx.attention.MLXAttentionBackend.get_name() -> str

Return backend name.

Parameters

This callable has no explicit inputs.

Returns

• Type: str
• Direct return expressions: 'MLX'

Exceptions and behavior

Method MLXAttentionBackend.get_name returns 'MLX'. No direct raise statement appears in this definition.

View source #L55-L57.

vllm_mlx.attention.MLXAttentionBackend.get_impl_cls · method

vllm_mlx.attention.MLXAttentionBackend.get_impl_cls() -> type

Return the implementation class.

Parameters

This callable has no explicit inputs.

Returns

• Type: type
• Direct return expressions: MLXAttentionImpl

Exceptions and behavior

Method MLXAttentionBackend.get_impl_cls returns MLXAttentionImpl. No direct raise statement appears in this definition.

View source #L60-L62.

vllm_mlx.attention.MLXAttentionBackend.get_metadata_cls · method

vllm_mlx.attention.MLXAttentionBackend.get_metadata_cls() -> type

Return the metadata class.

Parameters

This callable has no explicit inputs.

Returns

• Type: type
• Direct return expressions: MLXAttentionMetadata

Exceptions and behavior

Method MLXAttentionBackend.get_metadata_cls returns MLXAttentionMetadata. No direct raise statement appears in this definition.

View source #L65-L67.

vllm_mlx.attention.MLXAttentionBackend.get_kv_cache_shape · method

vllm_mlx.attention.MLXAttentionBackend.get_kv_cache_shape(num_blocks: int, block_size: int, num_kv_heads: int, head_size: int) -> tuple[int, ...]

Get the shape of KV cache.

Parameters

Name	Type	Required	Default	Description
num_blocks	int	yes	none	Number of cache blocks
block_size	int	yes	none	Tokens per block
num_kv_heads	int	yes	none	Number of KV attention heads
head_size	int	yes	none	Size of each attention head

Returns

• Type: tuple[int, ...]
• Direct return expressions: (num_blocks, block_size, num_kv_heads, head_size)

Exceptions and behavior

Method MLXAttentionBackend.get_kv_cache_shape returns (num_blocks, block_size, num_kv_heads, head_size). No direct raise statement appears in this definition.

View source #L70-L89.

vllm_mlx.attention.MLXAttentionBackend.get_supported_head_sizes · method

vllm_mlx.attention.MLXAttentionBackend.get_supported_head_sizes() -> list[int]

Return supported attention head sizes.

Parameters

This callable has no explicit inputs.

Returns

• Type: list[int]
• Direct return expressions: [64, 80, 96, 112, 128, 256]

Exceptions and behavior

Method MLXAttentionBackend.get_supported_head_sizes returns [64, 80, 96, 112, 128, 256]. No direct raise statement appears in this definition.

View source #L92-L94.

vllm_mlx.attention.MLXAttentionBackend.validate_configuration · method

vllm_mlx.attention.MLXAttentionBackend.validate_configuration(num_heads: int, head_size: int, num_kv_heads: int, dtype: 'torch.dtype', block_size: int, **kwargs) -> list[str]

Validate attention configuration.

Parameters

Name	Type	Required	Default	Description
num_heads	int	yes	none	Required positional or keyword input.
head_size	int	yes	none	Required positional or keyword input.
num_kv_heads	int	yes	none	Required positional or keyword input.
dtype	'torch.dtype'	yes	none	Required positional or keyword input.
block_size	int	yes	none	Required positional or keyword input.
**kwargs	not annotated	no	none	Additional variadic keyword inputs accepted by this callable.

Returns

• Type: list[str]
• Direct return expressions: errors

Exceptions and behavior

Method MLXAttentionBackend.validate_configuration calls MLXAttentionBackend.get_supported_head_sizes, errors.append; returns errors. No direct raise statement appears in this definition.

View source #L97-L118.

vllm_mlx.attention.MLXAttentionBackend.supports_dtype · method

vllm_mlx.attention.MLXAttentionBackend.supports_dtype(dtype: 'torch.dtype') -> bool

Check if dtype is supported.

Parameters

Name	Type	Required	Default	Description
dtype	'torch.dtype'	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: dtype in [torch.float16, torch.bfloat16, torch.float32]

Exceptions and behavior

Method MLXAttentionBackend.supports_dtype returns dtype in [torch.float16, torch.bfloat16, torch.float32]. No direct raise statement appears in this definition.

View source #L121-L125.

vllm_mlx.attention.MLXAttentionBackend.supports_block_size · method

vllm_mlx.attention.MLXAttentionBackend.supports_block_size(block_size: int) -> bool

Check if block size is supported.

Parameters

Name	Type	Required	Default	Description
block_size	int	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: block_size in [8, 16, 32]

Exceptions and behavior

Method MLXAttentionBackend.supports_block_size returns block_size in [8, 16, 32]. No direct raise statement appears in this definition.

View source #L128-L130.

vllm_mlx.attention.MLXAttentionBackend.supports_attn_type · method

vllm_mlx.attention.MLXAttentionBackend.supports_attn_type(attn_type: str) -> bool

Check if attention type is supported.

Parameters

Name	Type	Required	Default	Description
attn_type	str	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: attn_type in ['decoder', 'encoder', 'encoder_decoder']

Exceptions and behavior

Method MLXAttentionBackend.supports_attn_type returns attn_type in ['decoder', 'encoder', 'encoder_decoder']. No direct raise statement appears in this definition.

View source #L133-L135.

vllm_mlx.attention.MLXAttentionImpl · class

vllm_mlx.attention.MLXAttentionImpl(num_heads: int, head_size: int, scale: float, num_kv_heads: int | None = None, alibi_slopes: list[float] | None = None, sliding_window: int | None = None, kv_cache_dtype: str = 'auto', blocksparse_params: dict | None = None, logits_soft_cap: float | None = None, **kwargs)

MLX attention implementation.

Parameters

Name	Type	Required	Default	Description
num_heads	int	yes	none	Number of attention heads
head_size	int	yes	none	Size of each head
scale	float	yes	none	Attention scale factor
num_kv_heads	int \| None	no	None	Number of KV heads (for GQA/MQA)
alibi_slopes	list[float] \| None	no	None	ALiBi position encoding slopes
sliding_window	int \| None	no	None	Sliding window attention size
kv_cache_dtype	str	no	'auto'	KV cache data type
blocksparse_params	dict \| None	no	None	Block-sparse attention params
logits_soft_cap	float \| None	no	None	Soft cap for logits
**kwargs	not annotated	no	none	Additional variadic keyword inputs accepted by this callable.

Returns

• Constructs: vllm_mlx.attention.MLXAttentionImpl

Exceptions and behavior

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

View source #L138-L240.

vllm_mlx.attention.MLXAttentionImpl.__init__ · method

vllm_mlx.attention.MLXAttentionImpl.__init__(num_heads: int, head_size: int, scale: float, num_kv_heads: int | None = None, alibi_slopes: list[float] | None = None, sliding_window: int | None = None, kv_cache_dtype: str = 'auto', blocksparse_params: dict | None = None, logits_soft_cap: float | None = None, **kwargs) -> not annotated

Initialize MLX attention.

Parameters

Name	Type	Required	Default	Description
num_heads	int	yes	none	Number of attention heads
head_size	int	yes	none	Size of each head
scale	float	yes	none	Attention scale factor
num_kv_heads	int \| None	no	None	Number of KV heads (for GQA/MQA)
alibi_slopes	list[float] \| None	no	None	ALiBi position encoding slopes
sliding_window	int \| None	no	None	Sliding window attention size
kv_cache_dtype	str	no	'auto'	KV cache data type
blocksparse_params	dict \| None	no	None	Block-sparse attention params
logits_soft_cap	float \| None	no	None	Soft cap for logits
**kwargs	not annotated	no	none	Additional variadic keyword inputs accepted by this callable.

Returns

• Type: not annotated

Exceptions and behavior

Method MLXAttentionImpl.__init__ updates self.num_heads, self.head_size, self.scale, self.num_kv_heads; calls logger.debug. No direct raise statement appears in this definition.

View source #L147-L186.

vllm_mlx.attention.MLXAttentionImpl.forward · method

vllm_mlx.attention.MLXAttentionImpl.forward(query: Any, key: Any, value: Any, kv_cache: Any | None = None, attn_metadata: MLXAttentionMetadata | None = None, output: Any | None = None, **kwargs) -> Any

Compute attention.

Parameters

Name	Type	Required	Default	Description
query	Any	yes	none	Query tensor
key	Any	yes	none	Key tensor
value	Any	yes	none	Value tensor
kv_cache	Any \| None	no	None	Optional KV cache
attn_metadata	MLXAttentionMetadata \| None	no	None	Attention metadata
output	Any \| None	no	None	Optional output buffer
**kwargs	not annotated	no	none	Additional variadic keyword inputs accepted by this callable.

Returns

• Type: Any
• Direct return expressions: attn_output

Exceptions and behavior

Method MLXAttentionImpl.forward calls isinstance, mx.array, hasattr, query.numpy; returns attn_output. No direct raise statement appears in this definition.

View source #L188-L240.

vllm_mlx.attention.create_mlx_attention_backend · function

vllm_mlx.attention.create_mlx_attention_backend() -> type

Factory function to create MLX attention backend.

Parameters

This callable has no explicit inputs.

Returns

• Type: type
• Direct return expressions: MLXAttentionBackend

Exceptions and behavior

Function create_mlx_attention_backend returns MLXAttentionBackend. No direct raise statement appears in this definition.

View source #L243-L245.

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
MLXAttentionMetadata	class	MLXAttentionMetadata(seq_lens: list[int], max_seq_len: int, num_prefill_tokens: int = 0, num_decode_tokens: int = 0, block_tables: Any \| None = None, slot_mapping: Any \| None = None)	Metadata for MLX attention computation.	#L20-L39
MLXAttentionBackend	class	MLXAttentionBackend()	Attention backend using MLX's native attention.	#L42-L135
MLXAttentionBackend.get_name	method	MLXAttentionBackend.get_name() -> str	Return backend name.	#L55-L57
MLXAttentionBackend.get_impl_cls	method	MLXAttentionBackend.get_impl_cls() -> type	Return the implementation class.	#L60-L62
MLXAttentionBackend.get_metadata_cls	method	MLXAttentionBackend.get_metadata_cls() -> type	Return the metadata class.	#L65-L67
MLXAttentionBackend.get_kv_cache_shape	method	MLXAttentionBackend.get_kv_cache_shape(num_blocks: int, block_size: int, num_kv_heads: int, head_size: int) -> tuple[int, ...]	Get the shape of KV cache.	#L70-L89
MLXAttentionBackend.get_supported_head_sizes	method	MLXAttentionBackend.get_supported_head_sizes() -> list[int]	Return supported attention head sizes.	#L92-L94
MLXAttentionBackend.validate_configuration	method	MLXAttentionBackend.validate_configuration(num_heads: int, head_size: int, num_kv_heads: int, dtype: 'torch.dtype', block_size: int, **kwargs) -> list[str]	Validate attention configuration.	#L97-L118
MLXAttentionBackend.supports_dtype	method	MLXAttentionBackend.supports_dtype(dtype: 'torch.dtype') -> bool	Check if dtype is supported.	#L121-L125
MLXAttentionBackend.supports_block_size	method	MLXAttentionBackend.supports_block_size(block_size: int) -> bool	Check if block size is supported.	#L128-L130
MLXAttentionBackend.supports_attn_type	method	MLXAttentionBackend.supports_attn_type(attn_type: str) -> bool	Check if attention type is supported.	#L133-L135
MLXAttentionImpl	class	MLXAttentionImpl(num_heads: int, head_size: int, scale: float, num_kv_heads: int \| None = None, alibi_slopes: list[float] \| None = None, sliding_window: int \| None = None, kv_cache_dtype: str = 'auto', blocksparse_params: dict \| None = None, logits_soft_cap: float \| None = None, **kwargs)	MLX attention implementation.	#L138-L240
MLXAttentionImpl.__init__	method	MLXAttentionImpl.__init__(num_heads: int, head_size: int, scale: float, num_kv_heads: int \| None = None, alibi_slopes: list[float] \| None = None, sliding_window: int \| None = None, kv_cache_dtype: str = 'auto', blocksparse_params: dict \| None = None, logits_soft_cap: float \| None = None, **kwargs) -> not annotated	Initialize MLX attention.	#L147-L186
MLXAttentionImpl.forward	method	MLXAttentionImpl.forward(query: Any, key: Any, value: Any, kv_cache: Any \| None = None, attn_metadata: MLXAttentionMetadata \| None = None, output: Any \| None = None, **kwargs) -> Any	Compute attention.	#L188-L240
create_mlx_attention_backend	function	create_mlx_attention_backend() -> type	Factory function to create MLX attention backend.	#L243-L245