[RFC] DFlash

.gitignore

@@ -69,6 +69,7 @@ xcuserdata/
 /src/executorch/include/
 /src/executorch/share/
 /src/executorch/version.py
+/dflash_benchmarks.md
 *_etdump
 
 # Android
@@ -85,3 +86,6 @@ zephyr_dev_root.backup.*/
 # Agents
 .claude/*.local.*
 extension/pybindings/mlx.metallib
+
+# Scratch/WIP work not ready for review -- never committed
+/wip/

Makefile

@@ -484,3 +484,8 @@ clean:
 	rm -rf cmake-out \
 	       extension/llm/tokenizers/build \
 	       extension/llm/tokenizers/pytorch_tokenizers.egg-info
+
+# qwen3_dflash-mlx target removed: it depended on the C++ engine sources
+# (CMakeLists.txt, CMakePresets.json, qwen3_dflash_engine.*), which are
+# gitignored/not yet landed. Restore this target in the follow-up PR that
+# actually lands the C++ engine.

backends/mlx/.gitignore

@@ -0,0 +1,9 @@
+# Auto-generated by backends/mlx/serialization/generate.py — do not commit.
+# See backends/mlx/serialization/README.md for regeneration instructions.
+runtime/MLXLoader.cpp
+runtime/MLXLoader.h
+runtime/schema_generated.h
+serialization/_generated/
+serialization/_generated_serializers.py
+serialization/mlx_graph_schema.py
+_generated_inspector.py

backends/mlx/examples/llm/dflash_draft_cache.py

@@ -0,0 +1,70 @@
+# 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.
+
+"""Persistent per-layer KV cache for the DFlash draft model.
+
+Follows the TorchExportableModuleWithStaticCache pattern (per review): cache
+tensors are registered as mutable buffers and cache_position is passed
+through the call rather than tracked internally. Built on the existing
+KVCache (backends/mlx/llm/cache.py), so writes go through
+torch.ops.mlx.kv_cache_update instead of a Python slice (avoids the
+GuardOnDataDependentSymNode failure hit by an earlier attempt).
+"""
+
+from typing import Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from executorch.backends.mlx.llm.cache import KVCache
+
+
+class DFlashDraftKVCache(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        num_heads: int,
+        head_dim: int,
+        max_seq_len: int,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        self.max_seq_len = max_seq_len
+        self.layers = nn.ModuleList(
+            [
+                KVCache(
+                    max_batch_size=1,
+                    max_context_length=max_seq_len,
+                    n_heads=num_heads,
+                    head_dim=head_dim,
+                    enable_dynamic_shape=True,
+                    dtype=dtype,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def write(
+        self,
+        layer_idx: int,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        cache_position: Union[torch.Tensor, int],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Write K/V at cache_position and return the FULL buffer for this
+        # layer. Caller masks the unwritten tail via valid_mask(). Extract
+        # cache_position once and reuse across layers, not once per layer.
+        return self.layers[layer_idx].update(cache_position, key_states, value_states)
+
+    def valid_mask(self, valid_len: Union[torch.Tensor, int], device=None) -> torch.Tensor:
+        # True for positions [0, valid_len), False for the unwritten tail.
+        positions = torch.arange(self.max_seq_len, device=device)
+        return positions < valid_len
+
+    def reset(self) -> None:
+        for layer in self.layers:
+            layer.k_cache.zero_()
+            layer.v_cache.zero_()