diff --git a/backends/mlx/llm/hf_attention.py b/backends/mlx/llm/hf_attention.py index ff8fc3ff9a3..d55f370ff68 100644 --- a/backends/mlx/llm/hf_attention.py +++ b/backends/mlx/llm/hf_attention.py @@ -43,6 +43,10 @@ import executorch.backends.mlx.custom_ops as _mlx_custom_ops # noqa: F401 import torch from executorch.backends.mlx.llm.exportable import _LogitsToKeepMixin, LogitsToKeepMode +from executorch.extension.llm.cache.hf_attention import ( + register_offgraph_attention, + sdpa_mask_passthrough, +) def mlx_sdpa_with_start_pos_forward( @@ -111,94 +115,6 @@ def mlx_sdpa_with_start_pos_forward( return output.transpose(1, 2).contiguous(), None -def _cache_id(module: torch.nn.Module) -> int: - """Which cache a layer addresses -- a cache id, not a model layer index. - - A KV-sharing layer (gemma 4's YOCO) computes no k/v of its own; HuggingFace - hands it the k/v of the last non-sharing layer of the same attention type. - Pointing it at that donor's cache makes its write repeat what the donor - already stored and its read the shared history, so only donors own a cache - (15 rather than 35 for gemma-4-E2B). - """ - if not getattr(module, "is_kv_shared_layer", False): - return module.layer_idx - cfg = module.config - first_shared = cfg.num_hidden_layers - getattr(cfg, "num_kv_shared_layers", 0) - donors = list(cfg.layer_types[:first_shared]) - # Last index in `donors` of this layer's attention type. Kept identical to - # how Gemma4TextAttention derives `store_full_length_kv`, since that is the - # layer whose k/v we are handed; the two must agree. - return len(donors) - 1 - donors[::-1].index(cfg.layer_types[module.layer_idx]) - - -def mlx_offgraph_attention_forward( - module: torch.nn.Module, - query: torch.Tensor, # [B, num_heads, q_len, head_dim] - BHSD - key: torch.Tensor, # [B, num_kv_heads, q_len, head_dim] - BHSD (this step) - value: torch.Tensor, # [B, num_kv_heads, q_len, head_dim] - BHSD (this step) - attention_mask: Union[torch.Tensor, "BlockMask"], # noqa: F821 - position_ids: Optional[torch.Tensor] = None, - scaling: Optional[float] = None, - softcap: Optional[float] = None, - head_mask: Optional[torch.Tensor] = None, - **kwargs, -) -> Tuple[torch.Tensor, None]: - """ - Attention over the off-graph KV cache (kvcache::update_and_attend). - - Unlike the mlx::custom_sdpa path, no cache is in the graph: the model must be - run with use_cache=False so key/value are this step's projections, not full - history. The op owns the cache, placing/reading by `position` and masking - itself, so this is causal-only and never materializes a mask tensor. - - Returns (output, None) where output is [B, q_len, num_heads, head_dim] (BSHD). - """ - # As above: reject only what changes the result, ignore the rest. - if kwargs.get("dropout"): - raise ValueError("mlx_offgraph attention does not support dropout") - if softcap is not None: - raise ValueError("mlx_offgraph attention does not support softcap") - if head_mask is not None: - raise ValueError("mlx_offgraph attention does not support head_mask") - assert ( - position_ids is not None - ), "position_ids must be provided to place tokens in the off-graph cache" - # Per-token absolute positions [q_len, 1]; the op places + masks from these. - position = position_ids[0].reshape(-1, 1) - assert scaling is not None, "scaling must be provided by the attention module" - - output = torch.ops.kvcache.update_and_attend( - query, - key, - value, - position, - layer_id=_cache_id(module), - scale=scaling, - out_dtype=query.dtype, - ) - - # Transpose BHSD → BSHD for HF - return output.transpose(1, 2).contiguous(), None - - -def sdpa_mask_passthrough( - batch_size: int, - cache_position: Optional[torch.Tensor] = None, - q_length: Optional[int] = None, - kv_length: Optional[int] = None, - q_offset: Optional[Union[int, torch.Tensor]] = None, - kv_offset: int = 0, - mask_function: Optional[Callable] = None, - attention_mask: Optional[torch.Tensor] = None, - local_size: Optional[int] = None, - allow_is_causal_skip: bool = True, - allow_torch_fix: bool = True, - **kwargs, -) -> Optional[torch.Tensor]: - """Returns None — custom SDPA handles causal masking, avoiding bounded mask tensors.""" - return None - - def register_mlx_attention(name: str = "mlx") -> None: """ Register MLX attention with HuggingFace's attention interfaces. @@ -263,21 +179,9 @@ def register_mlx_offgraph_attention(name: str = "mlx_offgraph") -> None: Models using attn_implementation="mlx_offgraph" must be exported with use_cache=False (no StaticCache): the cache lives outside the graph, bound at - runtime via cache_key. Importing the op module registers the custom op. + runtime via cache_key. """ - from executorch.extension.llm.cache import update_and_attend # noqa: F401 - - try: - from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS - from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS - - ALL_ATTENTION_FUNCTIONS.register(name, mlx_offgraph_attention_forward) - ALL_MASK_ATTENTION_FUNCTIONS.register(name, sdpa_mask_passthrough) - - except ImportError: - raise ImportError( - "transformers is not installed. Please install it: pip install transformers" - ) + register_offgraph_attention(name) def get_mlx_sliding_window_sdpa(exportable_module) -> Callable: diff --git a/extension/llm/cache/BUCK b/extension/llm/cache/BUCK index ab6da8818c8..111d1eb11c0 100644 --- a/extension/llm/cache/BUCK +++ b/extension/llm/cache/BUCK @@ -30,6 +30,19 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "hf_attention", + srcs = [ + "hf_attention.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":cache", + "//caffe2:torch", + "fbsource//third-party/pypi/transformers:transformers", + ], +) + fbcode_target(_kind = runtime.python_test, name = "test_update_and_attend", srcs = [ diff --git a/extension/llm/cache/hf_attention.py b/extension/llm/cache/hf_attention.py new file mode 100644 index 00000000000..08011642f58 --- /dev/null +++ b/extension/llm/cache/hf_attention.py @@ -0,0 +1,101 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Hugging Face attention over the off-graph KV cache, for any backend. + +``register_offgraph_attention`` registers an HF attention implementation whose +layers each emit one ``kvcache::update_and_attend``. The model must run with +``use_cache=False`` so key/value are this step's projections; the op places +them by ``position_ids`` and masks causally itself, so the registered mask +function returns None. +""" + +from typing import Callable, Optional, Tuple, Union + +import executorch.extension.llm.cache.update_and_attend # noqa: F401 +import torch + + +def _cache_id(module: torch.nn.Module) -> int: + """Which cache a layer addresses -- a cache id, not a model layer index. + + A KV-sharing layer (gemma 4's YOCO) computes no k/v of its own; HuggingFace + hands it the k/v of the last non-sharing layer of the same attention type. + Pointing it at that donor's cache makes its write repeat what the donor + already stored and its read the shared history, so only donors own a cache. + """ + if not getattr(module, "is_kv_shared_layer", False): + return module.layer_idx + cfg = module.config + first_shared = cfg.num_hidden_layers - getattr(cfg, "num_kv_shared_layers", 0) + donors = list(cfg.layer_types[:first_shared]) + # Must agree with how Gemma4TextAttention derives `store_full_length_kv`. + return len(donors) - 1 - donors[::-1].index(cfg.layer_types[module.layer_idx]) + + +def offgraph_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, # [B, num_heads, q_len, head_dim] + key: torch.Tensor, # [B, num_kv_heads, q_len, head_dim], this step only + value: torch.Tensor, # [B, num_kv_heads, q_len, head_dim], this step only + attention_mask: Union[torch.Tensor, "BlockMask"], # noqa: F821 + position_ids: Optional[torch.Tensor] = None, + scaling: Optional[float] = None, + softcap: Optional[float] = None, + head_mask: Optional[torch.Tensor] = None, + **kwargs, +) -> Tuple[torch.Tensor, None]: + """Returns (output, None) with output [B, q_len, num_heads, head_dim].""" + # Reject only what changes the result; HF forwards open-ended kwargs. + if kwargs.get("dropout"): + raise ValueError("off-graph attention does not support dropout") + if softcap is not None: + raise ValueError("off-graph attention does not support softcap") + if head_mask is not None: + raise ValueError("off-graph attention does not support head_mask") + assert ( + position_ids is not None + ), "position_ids must be provided to place tokens in the off-graph cache" + position = position_ids[0].reshape(-1, 1) + assert scaling is not None, "scaling must be provided by the attention module" + + output = torch.ops.kvcache.update_and_attend( + query, + key, + value, + position, + layer_id=_cache_id(module), + scale=scaling, + out_dtype=query.dtype, + ) + return output.transpose(1, 2).contiguous(), None + + +def sdpa_mask_passthrough( + batch_size: int, + cache_position: Optional[torch.Tensor] = None, + q_length: Optional[int] = None, + kv_length: Optional[int] = None, + q_offset: Optional[Union[int, torch.Tensor]] = None, + kv_offset: int = 0, + mask_function: Optional[Callable] = None, + attention_mask: Optional[torch.Tensor] = None, + local_size: Optional[int] = None, + allow_is_causal_skip: bool = True, + allow_torch_fix: bool = True, + **kwargs, +) -> Optional[torch.Tensor]: + """Returns None: the attention op masks itself, so no bounded mask tensor is traced.""" + return None + + +def register_offgraph_attention(name: str = "offgraph") -> None: + """Register off-graph attention as HF ``attn_implementation=name``.""" + from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + ALL_ATTENTION_FUNCTIONS.register(name, offgraph_attention_forward) + ALL_MASK_ATTENTION_FUNCTIONS.register(name, sdpa_mask_passthrough)