diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py index fe6b8094f6b..93f6d354a0f 100644 --- a/backends/mlx/examples/llm/export_llm_hf.py +++ b/backends/mlx/examples/llm/export_llm_hf.py @@ -312,6 +312,148 @@ def _export_with_custom_components( _save_program(executorch_program, output_path) +def _export_with_offgraph_cache( + model_id: str, + revision: Optional[str], + output_path: str, + max_seq_len: int, + dtype: str, + qlinear: Optional[str], + qembedding: Optional[str], + no_tie_word_embeddings: bool = False, + qlinear_group_size: Optional[int] = None, + qembedding_group_size: Optional[int] = None, +) -> None: + """ + Export using the off-graph KV cache op (kvcache::update_and_attend). + + Unlike the custom-components path, the cache is not a graph tensor: the model + is run with use_cache=False (no StaticCache wrapper, no HFStaticCache swap), + so each attention layer emits an update_and_attend node fed this step's k/v. + The cache is created and bound at runtime via a cache_key. + """ + import executorch.exir as exir + from executorch.backends.mlx import MLXPartitioner + from executorch.backends.mlx.llm.hf_attention import ( + OffGraphExportWrapper, + register_mlx_offgraph_attention, + ) + from executorch.backends.mlx.passes import get_default_passes + from executorch.exir import EdgeCompileConfig + from executorch.exir.capture._config import ExecutorchBackendConfig + from executorch.exir.passes import MemoryPlanningPass + from transformers import AutoModelForCausalLM + + torch_dtype_map = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, + } + torch_dtype = torch_dtype_map.get(dtype, torch.bfloat16) + + register_mlx_offgraph_attention() + logger.info("Registered MLX off-graph attention (update_and_attend)") + + logger.info(f"Loading HuggingFace model: {model_id}") + load_kwargs = { + "torch_dtype": torch_dtype, + "low_cpu_mem_usage": True, + "attn_implementation": "mlx_offgraph", + } + if revision is not None: + load_kwargs["revision"] = revision + model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) + model.eval() + + from executorch.backends.mlx.llm.quantization import quantize_model_ + + quantize_model_( + model, + qlinear_config=qlinear, + qlinear_group_size=qlinear_group_size, + qembedding_config=qembedding, + qembedding_group_size=qembedding_group_size, + tie_word_embeddings=getattr(model.config, "tie_word_embeddings", False) + and not no_tie_word_embeddings, + ) + + exportable = OffGraphExportWrapper(model) + + # The cache is built before the graph runs, so the runtime cannot learn its + # shape from the model. Publish it, one entry per cache: + # n_caches how many caches to allocate + # kv_heads KV heads per cache + # head_dims head dim per cache -- gemma-4 mixes 256 and 512 + # windows sliding window per cache, 0 = flat + from executorch.backends.mlx.llm.cache import resolve_hf_cache_layout + + # resolve_hf_cache_layout drops the KV-shared tail, so these are indexed by + # cache -- the same indexing _cache_id maps layers onto. + layer_types, cache_kv_heads, cache_head_dims = resolve_hf_cache_layout(model.config) + text_config = model.config.get_text_config() + sliding_window = getattr(text_config, "sliding_window", None) or 0 + cache_windows = [ + sliding_window if t == "sliding_attention" else 0 for t in layer_types + ] + kv_metadata = { + # The cache count, not num_hidden_layers: gemma-4 E2B shares KV across + # its tail, so 15 caches for 35 layers. + "get_n_caches": len(layer_types), + "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32), + "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32), + "get_windows": torch.tensor(cache_windows, dtype=torch.int32), + } + logger.info( + f"KV cache layout: {len(layer_types)} caches, " + f"{sum(1 for w in cache_windows if w)} sliding (window {sliding_window})" + ) + + logger.info("Exporting model with torch.export...") + seq_length = 3 + example_input_ids = torch.zeros((1, seq_length), dtype=torch.long) + example_cache_position = torch.arange(seq_length, dtype=torch.long) + + seq_len_dim = torch.export.Dim("seq_length_dim", max=max_seq_len - 1) + dynamic_shapes = { + "input_ids": {1: seq_len_dim}, + "cache_position": {0: seq_len_dim}, + } + + with torch.no_grad(): + exported_program = torch.export.export( + exportable, + args=(), + kwargs={ + "input_ids": example_input_ids, + "cache_position": example_cache_position, + }, + dynamic_shapes=dynamic_shapes, + strict=True, + ) + + logger.info("Delegating to MLX backend...") + edge_program = exir.to_edge_transform_and_lower( + {"forward": exported_program}, + transform_passes=get_default_passes(), + constant_methods=kv_metadata, + partitioner=[MLXPartitioner()], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + + logger.info("Exporting to ExecuTorch...") + executorch_program = edge_program.to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=True), + ) + ) + + _save_program(executorch_program, output_path) + + def _save_program(executorch_program, output_path: str) -> None: """Save the ExecuTorch program to disk.""" os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) @@ -332,6 +474,7 @@ def export_llama_hf( qembedding: Optional[str] = None, use_custom_sdpa: bool = False, use_custom_kv_cache: bool = False, + use_offgraph_cache: bool = False, no_tie_word_embeddings: bool = False, qlinear_group_size: Optional[int] = None, qembedding_group_size: Optional[int] = None, @@ -349,7 +492,26 @@ def export_llama_hf( use_custom_sdpa: Use MLX custom SDPA (mlx::custom_sdpa) use_custom_kv_cache: Use MLX custom KV cache (mlx::kv_cache_update) """ - if use_custom_sdpa or use_custom_kv_cache: + if use_offgraph_cache: + if use_custom_sdpa or use_custom_kv_cache: + raise ValueError( + "--use-offgraph-cache is exclusive with --use-custom-sdpa / " + "--use-custom-kv-cache (it replaces both)" + ) + logger.info("Using off-graph KV cache (update_and_attend)") + _export_with_offgraph_cache( + model_id=model_id, + revision=revision, + output_path=output_path, + max_seq_len=max_seq_len, + dtype=dtype, + qlinear=qlinear, + qembedding=qembedding, + no_tie_word_embeddings=no_tie_word_embeddings, + qlinear_group_size=qlinear_group_size, + qembedding_group_size=qembedding_group_size, + ) + elif use_custom_sdpa or use_custom_kv_cache: logger.info( f"Using custom components: sdpa={use_custom_sdpa}, " f"kv_cache={use_custom_kv_cache}" @@ -434,6 +596,13 @@ def main(): default=False, help="Use MLX custom KV cache (mlx::kv_cache_update)", ) + parser.add_argument( + "--use-offgraph-cache", + action="store_true", + default=False, + help="Use the off-graph KV cache op (kvcache::update_and_attend); " + "replaces --use-custom-sdpa/--use-custom-kv-cache", + ) args = parser.parse_args() @@ -447,6 +616,7 @@ def main(): qembedding=args.qembedding, use_custom_sdpa=args.use_custom_sdpa, use_custom_kv_cache=args.use_custom_kv_cache, + use_offgraph_cache=args.use_offgraph_cache, no_tie_word_embeddings=args.no_tie_word_embeddings, qlinear_group_size=args.qlinear_group_size, qembedding_group_size=args.qembedding_group_size, diff --git a/backends/mlx/llm/cache.py b/backends/mlx/llm/cache.py index 013a37c6916..3d1ef536637 100644 --- a/backends/mlx/llm/cache.py +++ b/backends/mlx/llm/cache.py @@ -525,7 +525,9 @@ def update( # Current HF ExecuTorch wrappers copy the requested cache position # into each StaticCache layer's cumulative_length before forward(). if hasattr(self.layers[layer_idx], "cumulative_length"): - cache_position = self.layers[layer_idx].cumulative_length + # cumulative_length is a scalar; KVCache.update indexes [0], so + # give it the 1-D shape a cache_kwargs caller would have passed. + cache_position = self.layers[layer_idx].cumulative_length.reshape(1) else: raise RuntimeError( "cache_position was not provided and the pinned " diff --git a/backends/mlx/llm/hf_attention.py b/backends/mlx/llm/hf_attention.py index f2a01c9e653..1977ef3ce7b 100644 --- a/backends/mlx/llm/hf_attention.py +++ b/backends/mlx/llm/hf_attention.py @@ -9,14 +9,23 @@ """ MLX-optimized attention for HuggingFace models. -Registers a custom attention implementation ("mlx") with HuggingFace's -attention interface, following the same pattern as optimum-executorch's -custom_sdpa: - -1. Mask function returns None (custom op handles causal masking internally) -2. Attention function extracts start_pos from position_ids[0][0] -3. mlx::custom_sdpa receives full K/V cache + start_pos, slices K/V internally -4. MLX pattern handler serializes custom_sdpa as SliceNode(K), SliceNode(V), SdpaNode +Two implementations are registered with HuggingFace's attention interface, +differing in where the KV cache lives. Both register a mask function returning +None, since the op masks internally and a mask tensor would be traced at a +fixed size. + +"mlx" (register_mlx_attention) keeps the cache in the graph: the model runs +with a StaticCache, so key/value are the full history, and the attention +function extracts start_pos from position_ids and hands both to +mlx::custom_sdpa, which slices K/V and applies causal masking. The MLX pattern +handler serializes it as SliceNode(K), SliceNode(V), SdpaNode. + +"mlx_offgraph" (register_mlx_offgraph_attention) keeps the cache out of the +graph: the model must run with use_cache=False, so key/value are only this +step's projections, and each layer emits one kvcache::update_and_attend fed the +token positions. The cache is created and bound at run time via a cache_key. +OffGraphExportWrapper supplies the (input_ids, cache_position) signature that +path needs. Usage: from executorch.backends.mlx.llm.hf_attention import register_mlx_attention @@ -25,7 +34,7 @@ model = AutoModelForCausalLM.from_pretrained( model_id, - attn_implementation="mlx", + attn_implementation="mlx", # or "mlx_offgraph" ) """ @@ -56,7 +65,21 @@ def mlx_sdpa_with_start_pos_forward( Returns (output, None) where output is [B, seq_len, num_heads, head_dim] (BSHD). """ + # HuggingFace calls every registered implementation with the same argument + # list. Drop the ones handled elsewhere, and refuse the rest rather than + # ignoring them silently. softcap and head_mask stay named because models + # pass them unconditionally, often as None -- the value is what matters. kwargs.pop("is_causal", None) + kwargs.pop("use_cache", None) + kwargs.pop("sliding_window", None) # the export routes sliding layers away + if kwargs.pop("dropout", 0.0): + raise ValueError("mlx attention does not support dropout") + if kwargs: + raise ValueError(f"mlx attention got unsupported args: {sorted(kwargs)}") + if softcap is not None: + raise ValueError("mlx attention does not support softcap") + if head_mask is not None: + raise ValueError("mlx attention does not support head_mask") is_causal = getattr(module, "is_causal", True) if is_causal: @@ -87,6 +110,82 @@ 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). + """ + kwargs.pop("is_causal", None) + kwargs.pop("use_cache", None) + kwargs.pop("sliding_window", None) # the cache applies the window + if kwargs.pop("dropout", 0.0): + raise ValueError("mlx_offgraph attention does not support dropout") + if kwargs: + raise ValueError( + f"mlx_offgraph attention got unsupported args: {sorted(kwargs)}" + ) + 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, @@ -125,6 +224,62 @@ def register_mlx_attention(name: str = "mlx") -> None: ) +class OffGraphExportWrapper(torch.nn.Module): + """forward(input_ids, cache_position) -> logits, with no in-graph cache. + + The analog of TorchExportableModuleWithStaticCache for the off-graph op: + runs the model with use_cache=False so each attention layer sees only this + step's k/v (the op owns history), and exposes the (input_ids, cache_position) + signature the runner drives. + """ + + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + + def forward( + self, input_ids: torch.Tensor, cache_position: torch.Tensor + ) -> torch.Tensor: + # Single sequence: the op takes [q_len, n_dims] positions and the + # attention function reads position_ids[0], so a batch would be placed + # at row 0's positions. + assert input_ids.shape[0] == 1, "off-graph export supports batch size 1" + return self.model( + input_ids=input_ids, + cache_position=cache_position, + # Pass positions rather than letting the model infer them: it derives + # them from past_key_values.get_seq_length(), which is 0 here because + # the cache is out of the graph, so every decode step would place its + # token at position 0. + position_ids=cache_position.unsqueeze(0), + use_cache=False, + past_key_values=None, + ).logits + + +def register_mlx_offgraph_attention(name: str = "mlx_offgraph") -> None: + """ + Register off-graph KV-cache attention with HuggingFace's attention interface. + + 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. + """ + 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" + ) + + def get_mlx_sliding_window_sdpa(exportable_module) -> Callable: """ Create a closure-based SDPA function for sliding window attention.