up

Author: metascroy

backends/mlx/README.md

@@ -85,7 +85,7 @@ logging from the partitioner and preprocessor — including ops-to-not-decompose
 lists, graph dumps, per-node support decisions, and serialization details:
 
 ```bash
-ET_MLX_DEBUG=1 python -m executorch.backends.mlx.examples.llm.export_llama ...
+ET_MLX_DEBUG=1 python -m executorch.backends.mlx.examples.llm.export_llm_hf ...
 ```
 
 ---
@@ -107,19 +107,29 @@ backends/mlx/
 │   ├── MLXExecutor.h           # ExecutionState, constant loading, helpers
 │   ├── MLXInterpreter.h        # Op dispatch loop + per-op exec_* functions
 │   └── schema_generated.h      # [GENERATED] FlatBuffer C++ bindings (flatc)
+├── llm/                        # LLM infrastructure (KV cache, attention, etc.)
+│   ├── cache.py                # KV cache implementations (ET + HF static cache)
+│   ├── et_attention.py         # ExecuTorch custom SDPA attention
+│   ├── hf_attention.py         # HuggingFace custom SDPA attention
+│   ├── quantization.py         # TorchAO quantization helpers
+│   └── source_transformation.py # Source transforms for MLX export
 ├── ops.py                      # Op handlers  (ATen target → MLX IR node)
 ├── patterns.py                 # Pattern handlers (multi-node fusions)
+├── passes.py                   # Graph passes (RMSNorm fusion, CSE, etc.)
+├── pattern_utils.py            # Pattern matching utilities for passes
 ├── program_builder.py          # MLXProgramBuilder + REGISTRY
 ├── partitioner.py              # Decides which ops to delegate to MLX
 ├── preprocess.py               # BackendDetails.preprocess() entry point
-├── custom_ops.py               # Custom torch ops (rope, etc.)
+├── custom_ops.py               # Custom torch ops (kv_cache_update, custom_sdpa, rope)
+├── _logging.py                 # Debug logging utilities (ET_MLX_DEBUG)
+├── pte_inspector.py            # .pte file inspection/debugging tool
 ├── test/
 │   ├── test_ops.py             # Op test definitions (models + configs)
 │   ├── test_utils.py           # OpTestCase base class + helpers
 │   ├── op_test_runner.cpp      # C++ test runner (loads .pte, runs, compares)
 │   └── run_all_tests.py        # End-to-end: export → C++ run → compare
 └── examples/
-    ├── llm/                    # LLM export + run (Llama, etc.)
+    ├── llm/                    # LLM export + run via HuggingFace
     └── whisper/                # Whisper export + run
 ```
 

backends/mlx/custom_ops.py

@@ -20,11 +20,6 @@
 from torch import Tensor
 
 
-# =============================================================================
-# kv_cache_update: Functional KV cache update for BHSD layout
-# =============================================================================
-
-
 @torch.library.custom_op("mlx::kv_cache_update", mutates_args=("cache",))
 def kv_cache_update(
     cache: Tensor,  # [B, H, S_max, D] - mutated in place
@@ -73,7 +68,6 @@ def kv_cache_update(
         )
         cache[:, :, start_pos:end_pos, :] = new_values
 
-    # Return dummy tensor like llama.update_cache does
     return torch.empty((1,), dtype=new_values.dtype, device=new_values.device)
 
 
@@ -88,11 +82,6 @@ def kv_cache_update_fake(
     return torch.empty((1,), dtype=new_values.dtype, device="meta")
 
 
-# =============================================================================
-# custom_sdpa: Scaled Dot-Product Attention with KV cache slicing
-# =============================================================================
-
-
 @torch.library.custom_op("mlx::custom_sdpa", mutates_args=())
 def mlx_custom_sdpa(
     query: Tensor,  # [B, num_heads, seq_len, head_dim] - BHSD
@@ -194,11 +183,6 @@ def mlx_custom_sdpa_fake(
     return query.new_empty(query.shape)
 
 
-# =============================================================================
-# rope: Rotary Position Embedding (single tensor)
-# =============================================================================
-
-
 @torch.library.custom_op("mlx::rope", mutates_args=())
 def rope(
     x: Tensor,  # (B, H, T, D)

backends/mlx/examples/whisper/export_whisper.py

@@ -52,11 +52,6 @@
 logger = logging.getLogger(__name__)
 
 
-# =============================================================================
-# Whisper Encoder Wrapper
-# =============================================================================
-
-
 class WhisperEncoderExportable(nn.Module):
     """
     Wrapper around Whisper's encoder for export.
@@ -72,11 +67,6 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor:
         return self.encoder(input_features=input_features).last_hidden_state
 
 
-# =============================================================================
-# Whisper Decoder Self-Attention with KV Cache
-# =============================================================================
-
-
 class WhisperSelfAttentionWithCache(nn.Module):
     """
     Whisper self-attention layer with static KV cache.
@@ -147,11 +137,6 @@ def forward(
         return self.out_proj(attn_out)
 
 
-# =============================================================================
-# Whisper Cross-Attention (no cache update - K/V pre-computed)
-# =============================================================================
-
-
 class WhisperCrossAttention(nn.Module):
     """
     Whisper cross-attention layer.
@@ -192,11 +177,6 @@ def forward(
         return self.out_proj(attn_out)
 
 
-# =============================================================================
-# Whisper Decoder Layer Wrapper
-# =============================================================================
-
-
 class WhisperDecoderLayerWithCache(nn.Module):
     """
     Wrapper for a single Whisper decoder layer with KV cache.
@@ -254,11 +234,6 @@ def forward(
         return hidden_states
 
 
-# =============================================================================
-# Whisper Decoder Wrapper
-# =============================================================================
-
-
 class WhisperDecoderWithCache(nn.Module):
     """
     Whisper decoder wrapper with static KV cache.
@@ -335,11 +310,6 @@ def forward(
         return logits
 
 
-# =============================================================================
-# Cross-KV Projection Module
-# =============================================================================
-
-
 class WhisperCrossKVProjection(nn.Module):
     """
     Compute cross-attention K/V projections from encoder hidden states.
@@ -393,11 +363,6 @@ def forward(
         return tuple(k_list), tuple(v_list)
 
 
-# =============================================================================
-# Export Functions
-# =============================================================================
-
-
 def export_whisper_to_mlx(
     model_id: str,
     output_dir: str,
@@ -528,9 +493,6 @@ def export_whisper_to_mlx(
             logger.error("TorchAO not installed. Run: pip install torchao")
             raise
 
-    # =========================================================================
-    # Export Encoder
-    # =========================================================================
     logger.info("Exporting encoder...")
 
     with torch.no_grad():
@@ -541,9 +503,6 @@ def export_whisper_to_mlx(
 
     _save_to_pte(encoder_ep, os.path.join(output_dir, "encoder.pte"), "encoder")
 
-    # =========================================================================
-    # Export Cross-KV Projection
-    # =========================================================================
     logger.info("Exporting cross-KV projection...")
 
     with torch.no_grad():
@@ -561,9 +520,6 @@ def export_whisper_to_mlx(
 
     _save_to_pte(cross_kv_ep, os.path.join(output_dir, "cross_kv.pte"), "cross_kv")
 
-    # =========================================================================
-    # Export Decoder
-    # =========================================================================
     logger.info("Exporting decoder...")
 
     # Example inputs for decoder

backends/mlx/examples/whisper/run_whisper.py

@@ -115,9 +115,6 @@ def run_whisper_inference(  # noqa: C901
     )
     decoder_forward = decoder_program.load_method("forward")
 
-    # =========================================================================
-    # Step 1: Run encoder
-    # =========================================================================
     logger.info("Running encoder...")
     overall_start = time.time()
     start_time = time.time()
@@ -129,9 +126,6 @@ def run_whisper_inference(  # noqa: C901
     logger.info(f"Encoder time: {encoder_time:.3f}s")
     logger.info(f"Encoder output shape: {encoder_hidden_states.shape}")
 
-    # =========================================================================
-    # Step 2: Compute cross-attention K/V
-    # =========================================================================
     logger.info("Computing cross-attention K/V...")
     start_time = time.time()
 
@@ -145,9 +139,6 @@ def run_whisper_inference(  # noqa: C901
     logger.info(f"Cross-KV time: {cross_kv_time:.3f}s")
     logger.info(f"Cross K/V: {num_layers} layers, each shape {cross_k_tuple[0].shape}")
 
-    # =========================================================================
-    # Step 3: Setup decoder generation
-    # =========================================================================
     # Get forced decoder IDs for language/task
     forced_decoder_ids = processor.get_decoder_prompt_ids(
         language=language,
@@ -177,9 +168,6 @@ def run_whisper_inference(  # noqa: C901
 
     generated_tokens: List[int] = [sot_id]
 
-    # =========================================================================
-    # Step 4: Token-by-token decoder generation
-    # =========================================================================
     logger.info(f"Generating up to {max_new_tokens} tokens...")
     decode_start = time.time()
 

backends/mlx/llm/cache.py

@@ -126,7 +126,7 @@ def update(
         Returns:
             Tuple of (k_cache, v_cache) - slices of the FULL cache buffers
         """
-        # Extract start position as int (SymInt during tracing)
+
         if isinstance(input_pos, torch.Tensor):
             start_pos = input_pos[0].item()
             seq_len = k_val.size(2)
@@ -136,7 +136,6 @@ def update(
         else:
             start_pos = input_pos
 
-        # Use MLX custom op for cache update (mutates in place)
         torch.ops.mlx.kv_cache_update(self.k_cache, k_val, start_pos)
         torch.ops.mlx.kv_cache_update(self.v_cache, v_val, start_pos)
 
@@ -145,11 +144,6 @@ def update(
         return self.k_cache[:, :, :, :], self.v_cache[:, :, :, :]
 
 
-# =============================================================================
-# RingBufferKVCache - Sliding Window KV Cache
-# =============================================================================
-
-
 class RingBufferKVCache(nn.Module):
     """
     Ring buffer KV cache for sliding window attention.
@@ -235,8 +229,6 @@ def update(
             start_pos = input_pos
             seq_len = k_val.size(2)
 
-        # Use MLX custom op for ring buffer cache update (mutates in place)
-        # ring_size enables wrapping: write_pos = start_pos % ring_size
         torch.ops.mlx.kv_cache_update(
             self.k_cache, k_val, start_pos, ring_size=self.buffer_size
         )
@@ -287,10 +279,6 @@ def create_sliding_window_mask(self, start_pos: int, seq_len: int) -> torch.Tens
         return torch.where(attn_mask, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0)
 
 
-# =============================================================================
-# HFStaticCache - Standalone HuggingFace-compatible Static Cache
-# =============================================================================
-
 from transformers.cache_utils import StaticCache
 
 
@@ -424,28 +412,15 @@ def update(
         return self.kv_cache[layer_idx].update(cache_position, key_states, value_states)
 
     def get_seq_length(self, layer_idx: int = 0) -> int:
-        """
-        Get the current sequence length in the cache.
-
-        Note: This is approximate - returns the number of non-zero positions.
-
-        Args:
-            layer_idx: Layer index to check (default: 0)
-
-        Returns:
-            Approximate sequence length
-        """
-        # Check how many positions have been filled by looking for non-zero values
+        """Approximate sequence length (counts non-zero cache positions)."""
         k_cache = self.kv_cache[layer_idx].k_cache
         # Check if any value in the head_dim is non-zero for each position
         return (k_cache[0, 0, :, 0] != 0).sum().item()
 
     def get_max_cache_shape(self, layer_idx: int = 0) -> int:
-        """Get the maximum cache length."""
         return self.max_cache_len
 
     def reset(self):
-        """Reset all cache buffers to zero."""
         for layer_cache in self.kv_cache:
             layer_cache.k_cache.zero_()
             layer_cache.v_cache.zero_()

backends/mlx/llm/hf_attention.py

@@ -99,12 +99,7 @@ def sdpa_mask_passthrough(
     allow_torch_fix: bool = True,
     **kwargs,
 ) -> Optional[torch.Tensor]:
-    """
-    Returns None — the custom SDPA op handles causal masking internally.
-
-    Returning None avoids materializing a mask tensor during export, which
-    would create a bounded tensor that fails at runtime with longer sequences.
-    """
+    """Returns None — custom SDPA handles causal masking, avoiding bounded mask tensors."""
     return None
 
 
@@ -128,11 +123,6 @@ def register_mlx_attention(name: str = "mlx") -> None:
         )
 
 
-# =============================================================================
-# Sliding Window Attention (Ring Buffer)
-# =============================================================================
-
-
 def get_mlx_sliding_window_sdpa(exportable_module) -> Callable:
     """
     Create a closure-based SDPA function for sliding window attention.
@@ -219,16 +209,7 @@ def _sliding_window_sdpa_forward(
 def register_mlx_sliding_window_attention(
     exportable_module, name: str = "mlx_sliding_window"
 ) -> None:
-    """
-    Register MLX sliding window attention with HuggingFace's attention interfaces.
-
-    Creates a closure that captures the model reference for lazy mask creation,
-    following optimum-executorch's get_custom_sdpa_for_ring_kv_cache pattern.
-
-    Args:
-        exportable_module: The model module containing ring buffer caches.
-        name: Name to register the attention implementation under.
-    """
+    """Register MLX sliding window attention with HuggingFace's attention interfaces."""
     try:
         from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
         from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

backends/mlx/llm/quantization.py

@@ -19,7 +19,6 @@
 
 
 def add_quantization_args(parser: argparse.ArgumentParser) -> None:
-    """Add common quantization arguments to an argparse parser."""
     parser.add_argument(
         "--quantize-linear",
         type=str,
@@ -58,7 +57,6 @@ def add_quantization_args(parser: argparse.ArgumentParser) -> None:
 
 
 def _default_group_size(dtype_str: str) -> int:
-    """Return the default group size for a given quantization dtype."""
     return 32 if dtype_str == "int4" else 128
 
 

backends/mlx/llm/source_transformation.py

@@ -30,18 +30,7 @@ def _replace_modules(
     factory: Callable[[nn.Module], nn.Module],
     label: str,
 ) -> nn.Module:
-    """
-    Recursively replace all instances of target_type using factory.
-
-    Args:
-        module: Root module to modify (in place)
-        target_type: Type to match against children
-        factory: Callable that takes the original child and returns a replacement
-        label: Human-readable label for logging (e.g. "RMSNorm → FunctionalRMSNorm")
-
-    Returns:
-        The modified module (same reference, mutated in place)
-    """
+    """Recursively replace all instances of target_type using factory."""
 
     def _recurse(parent: nn.Module) -> int:
         count = 0