[RFC] dflash

[metascroy]

DFlash Speculative Decoding for ET

Overview

DFlash is a block diffusion draft model for speculative decoding. A lightweight draft model generates an entire block of tokens in parallel (not autoregressively), then the target model verifies them in a single forward pass. Accepted tokens are kept; rejected ones are discarded by rolling back the KV cache position.

This document describes how to implement DFlash on top of the existing ET MLX delegate for Gemma 4 31B (and similar models).

Architecture Recap

Target Model (Gemma 4 31B)
    │
    ├─ Layer 2 hidden ──────────┐
    ├─ Layer N//2 hidden ───────┤  concat → [B, T, num_tapped * H]
    ├─ Layer N-3 hidden ────────┘
    │                                │
    │                           FC(num_tapped*H → draft_H)
    │                           hidden_norm (RMSNorm)
    │                                │
    │                           Draft Model
    │                           ┌────┴────┐
    │                           │ embed(input) = h_proposal
    │                           │ For each draft layer:
    │                           │   attn(Q=h_proposal, KV=concat(h_ctx, h_proposal))
    │                           │   + MLP
    │                           │ lm_head(norm(h)) → draft_logits
    │                           └─────────┘
    │
    └─ logits

The draft model's attention is a cross-attention variant:

• Q comes from proposal token embeddings (the draft block)
• K/V comes from context hidden states (FC-projected target intermediates) concatenated with proposal K/V
• This lets the draft attend to rich target representations without reprocessing the full context

The draft model borrows embed_tokens and lm_head from the target (no weight duplication for those).

Design

Two .pte Programs

Export two separate .pte files sharing the same weight data where possible:

Program	Inputs	Outputs
target.pte	(tokens, input_pos)	(logits, hidden_states)
draft.pte	(tokens, hidden_states, input_pos)	(draft_logits)

The target program is the existing Gemma 4 export, modified to additionally output concatenated hidden states from tapped layers. The draft program is a new small model export.

Part 1: Target Model — Tapping Hidden States

Modify mlx_source_transformations.py to capture intermediate layer outputs and return them as a second output tensor.

Current model forward (_replace_model_forward in mlx_source_transformations.py):

def _mlx_model_forward(self, tokens, input_pos):
    x = self.embed_tokens(tokens) * self.embed_normalizer
    for layer in self.layers:
        x = layer(x, input_pos)
    x = self.norm(x)
    last = self.lm_head(x[:, -1, :]).float()
    cap = self.logit_softcap.float()
    return torch.tanh(last / cap) * cap

Modified forward — captures hidden states at selected layers:

def _mlx_model_forward_with_hidden(self, tokens, input_pos):
    x = self.embed_tokens(tokens) * self.embed_normalizer
    captured = []
    for i, layer in enumerate(self.layers):
        x = layer(x, input_pos)
        if i in self._dflash_layer_ids:
            captured.append(x)
    x = self.norm(x)
    last = self.lm_head(x[:, -1, :]).float()
    cap = self.logit_softcap.float()
    logits = torch.tanh(last / cap) * cap
    hidden = torch.cat(captured, dim=-1)  # [B, T, num_tapped * H]
    return logits, hidden

This uses the EAGLE3/DFlash pattern of tapping layers at [2, N//2, N-3] (configurable via target_layer_ids from the draft model's config).

Key point: No new ops needed. The torch.cat of hidden states is a standard ATen op the MLX backend already handles. The model just returns a tuple of two tensors instead of one.

Part 2: Draft Model Export

The DFlash draft model is a small standard transformer. All ops it uses are already supported by the MLX backend:

Draft Model Op	ET MLX Support
nn.Linear (FC projection)	Yes — ops.py
RMSNorm (hidden_norm, layernorms)	Yes — ops.py
SDPA (cross-attention)	Yes — mlx::custom_sdpa
RoPE	Yes — mlx::rope
SiLU + MLP	Yes — ops.py
embed_tokens (shared)	Yes — ops.py
lm_head (shared)	Yes — ops.py
Softcap (tanh(x/cap)*cap)	Yes — ops.py

Draft model PyTorch definition: Write a DFlashDraftModel(nn.Module) in PyTorch that mirrors dflash/model_mlx.py lines 132-198. The forward signature:

def forward(self, tokens, hidden_states, input_pos):
    # tokens: [B, block_size] — [last_accepted, mask, mask, ..., mask]
    # hidden_states: [B, T_ctx, num_tapped * H] — from target
    # input_pos: [block_size] — position indices
    h = self.embed_tokens(tokens) * self.embed_scale
    h_ctx = self.hidden_norm(self.fc(hidden_states))
    for layer, cache in zip(self.layers, self.kv_caches):
        h = layer(h, h_ctx, input_pos, cache)
    return self.lm_head(self.norm(h[:, 1:, :]))  # skip first token (logits_start=1)

Draft attention needs to handle the cross-attention pattern: K/V from both context (target hidden) and proposal (draft tokens), with Q from proposal only. This maps to the existing mlx::custom_sdpa after concatenating ctx and proposal K/V into the cache:

# In draft attention forward:
q = self.q_proj(x_proposal)      # Q from proposal only
k_ctx = self.k_proj(x_context)   # K from target hidden
v_ctx = self.v_proj(x_context)   # V from target hidden  
k_prop = self.k_proj(x_proposal) # K from proposal
v_prop = self.v_proj(x_proposal) # V from proposal

# Concat context + proposal K/V → feed to SDPA
k = torch.cat([k_ctx, k_prop], dim=2)
v = torch.cat([v_ctx, v_prop], dim=2)
out = torch.ops.mlx.custom_sdpa(q, k, v, ...)

Note on cross-attention and caching: The draft model's KV cache holds the context (target hidden state) K/V from previous speculative iterations. Each iteration, the context is the same set of hidden states (from the most recent target forward), so the draft cache can be reset between speculation rounds. This simplifies the design — the draft cache does not need rollback, only reset.

Part 3: Weight Sharing

The draft model borrows embed_tokens and lm_head from the target. Two options:

Option A (simple): Duplicate these weights in the draft .pte. They're small relative to the full model (vocab_size × hidden_size ≈ 256K × 3584 × 2 bytes ≈ 1.7 GB for Gemma 4). But with quantization, the lm_head is much smaller.

Option B (shared NamedDataMap): Both .pte files reference the same .ptd weight data. The C++ host loads one NamedDataMap and passes it to both programs. This requires the export to use the same FQN for embed_tokens / lm_head in both models, or a mapping layer at load time.

Recommendation: Start with Option A for simplicity. Optimize to shared weights later if memory is a concern.

Part 4: KV Cache Rollback via Position Reset

As discussed, rollback after rejected draft tokens is free for standard attention:

1. kv_cache_update writes at start_pos — overwriting is the rollback
2. custom_sdpa slices KV to [:stop_pos] — stale data beyond is invisible

The C++ host simply passes the correct input_pos on the next call. No new ops needed.

For the draft model: The draft KV cache is reset between speculation rounds (the draft model re-processes context hidden states each round), so rollback is not needed.

Part 5: C++ Host Loop

The speculative decoding loop lives in C++, extending the existing Gemma4_31BEngine. Pseudocode:

// After prefill:
auto [logits, hidden] = target_module.forward(prompt_tokens, input_pos);
int64_t pos = prompt_len;
auto token = sample(logits);

while (pos < max_seq_len) {
    int bs = std::min(block_size, max_tokens - generated + 1);
    
    // 1. Draft: generate block of tokens
    //    input = [last_token, mask, mask, ..., mask]  (bs tokens)
    auto draft_input = make_draft_input(token, mask_id, bs);
    auto draft_logits = draft_module.forward(draft_input, hidden, draft_pos);
    auto draft_tokens = sample_all(draft_logits);  // [bs-1] tokens
    
    // 2. Verify: run target on [last_token] + draft_tokens
    auto verify_input = concat(token, draft_tokens);  // [bs] tokens
    auto verify_pos = arange(pos, pos + bs);
    auto [target_logits, new_hidden] = target_module.forward(verify_input, verify_pos);
    auto target_tokens = sample_all(target_logits);
    
    // 3. Accept/reject
    int accepted = find_first_mismatch(draft_tokens, target_tokens);
    auto new_tokens = draft_tokens[:accepted] + target_tokens[accepted];
    
    // 4. Position-based rollback (free!)
    pos += accepted + 1;
    // Next iteration uses pos as input_pos — overwrites stale cache entries
    
    // 5. Slice hidden states to match accepted length
    hidden = new_hidden.slice(0, accepted + 1);
    token = new_tokens.back();
    
    emit_tokens(new_tokens);
}

Key simplification vs. the MLX standalone: The MLX standalone (dflash/model_mlx.py) has to deal with _trim_recent_cache and _GDNStateCapture rollback. In ET, position-based rollback eliminates both for standard attention models. GatedDeltaNet models are out of scope for V1.

Part 6: Export Script

A new export script examples/models/gemma4_31b/export_dflash.py:

def export_dflash(
    target_model,
    draft_model_id,     # HuggingFace ID for DFlash draft weights
    config,
    output_dir,
):
    # 1. Load draft weights from HuggingFace
    draft_model = load_dflash_draft(draft_model_id, config)
    
    # 2. Apply MLX source transformations to target (with hidden state output)
    mlx_source_transformations_with_hidden(
        target_model,
        target_layer_ids=draft_model.config.target_layer_ids,
    )
    
    # 3. Export target model → target.pte
    #    Signature: (tokens, input_pos) → (logits, hidden_states)
    export_target(target_model, config, output_dir)
    
    # 4. Apply MLX source transformations to draft
    draft_mlx_source_transformations(draft_model)
    
    # 5. Export draft model → draft.pte
    #    Signature: (tokens, hidden_states, input_pos) → (draft_logits)
    export_draft(draft_model, config, output_dir)

Implementation Plan

Phase 1: Target Model with Hidden State Outputs

• Add _replace_model_forward_with_hidden() to mlx_source_transformations.py
• Add --dflash-layers flag to export.py to enable hidden state tapping
• Verify export produces (logits, hidden_states) tuple correctly

Phase 2: Draft Model Definition + Export

• Create examples/models/gemma4_31b/dflash_draft.py — PyTorch DFlash model
• Create examples/models/gemma4_31b/dflash_mlx_source_transformations.py
• Add draft model loading from HuggingFace (z-lab weights)
• Verify draft model exports to .pte and runs correctly

Phase 3: C++ Speculative Decoding Loop

• Add Gemma4_31BSpecDecEngine (or extend existing engine) with spec decode loop
• Implement draft block construction (last_token + mask tokens)
• Implement verification + acceptance logic
• Implement position-based rollback

Phase 4: Integration + Benchmarking

• End-to-end test: export both models, run spec decode
• Benchmark vs. baseline autoregressive decoding (measure speedup)
• Compare acceptance rates with dflash/benchmark.py MLX standalone results

What's NOT Needed

• No new custom ops
• No new Metal kernels
• No schema changes
• No C++ runtime changes
• No cache trim/rollback ops (position reset is sufficient)

Scope Limitations (V1)

• Standard attention only — GatedDeltaNet/recurrent state rollback is deferred
• Gemma 4 31B — other models follow the same pattern but need model-specific draft weights
• Greedy decoding for draft — temperature=0 simplifies verification (draft and target must match exactly)
• No tree drafting — DFlash uses chain (block) drafting, not tree speculation
• Single-batch — batch size 1 (matches current ET MLX LLM inference)

Available Draft Weights

Pre-trained DFlash draft weights from z-lab/ on HuggingFace:

• z-lab/DFlash-Gemma-4-27B-IT (for Gemma 4 27B)
• Various Qwen3, Llama 3.1, and other model families

Check https://github.com/zlab-pub/DFlash for the latest list.

[cthotti]

Hi, I'd like to work on this if it's still up for grabs. I have some experience with ExecuTorch's export pipeline and LLM inference, but would like to further my knowledge within this by working on this project. If possible, could you assign this to me, or let me know if it's already taken?

[cthotti]

Since Qwen3 draft weights are also available, I'll start development against Qwen/Qwen3-4B (with z-lab/Qwen3-4B-DFlash as the draft) rather than Gemma 4, because it fits my local memory constraints. Happy to switch if there's a reason it needs to be Gemma specifically.