Execute ANE kernels on already-prepared buffers. No mx ops. Pure ANE (IOSurface + CoreML) — safe to call from a background thread whi

Run one full draft step with fine-grained phase timing. Replicates _start_draft + thread run + post-join sequence exactly. Returns dict[phase

Begin sampling ANE power at 500ms intervals.

Stop sampling and return summary stats.

Advance GatedDeltaNet SSM state for a subset of tokens. Starting from initial_state, applies the recurrent update for each token using its (k

Python fallback: advance SSM state token-by-token.

Find the index of the first KVCache entry for fa_mask creation.

Forward full-attention layer with explicit KV cache arrays.

Forward linear layer with explicit cache inputs, returning all rollback data.

Install split SDPA + small-proj-padding hooks on full-attention layers.

Install split SDPA hooks on attention layers for hybrid (Qwen3.5) models. At long context, the standard SDPA processes all 16 query positions aga

KOD: find block_size that maximizes expected throughput. Kelly-Optimal Drafting selects the block size gamma that maximizes expected tokens p

Load per-head norm + layer norm weights only (no projection weights).

Quantize one layer's projection weights → Q8LayerWeights pyobject.

Replicate the soft-anchor computation from _start_draft.

Patch an Attention module to use TurboQuant quantized_attention. Replaces attn.__call__ with a version that detects TurboQuantKVCache and rou

Pick the exit layer for Mirror-SD early-exit. Chooses the median target_layer_id as the exit point, balancing prefix and suffix compute. This

Mirror-SD speculative decoding with early-exit + parallel draft. Per Mirror-SD Eq. 10: 1. Run target prefix layers (0..exit_layer) → get hi

Parallel ANE||GPU speculative decoding (Mirror-SD Eq. 10). Pipelines draft and verify so that the next iteration's draft on ANE overlaps with

Attention forward using TurboQuant quantized_attention. Replicates the standard Attention.__call__ logic but uses cache.quantized_attention()

Advance GatedDeltaNet SSM state using custom Metal kernel. Replaces the Python per-token loop with a single GPU dispatch. Falls back to the P

Build the ffn_residual kernel with int8-quantized gate/up/down weights baked in.

Fused final RMSNorm + lm_head projection. Inputs: hidden [1, hidden, 1, w_sq], norm_w [1, hidden, 1, w_sq], lm_head_w [1, hidden, 1, vocab_size] Out

Position-weighted loss: exp decay within each block. w_k = exp(-(k-1) / gamma) for position k within the block. Non-masked positions get weig

Build the mega_qkv kernel with int8-quantized projection weights baked in. Inputs: hidden, in_norm_w, context, k_norm_w, q_norm_w, cos_k, sin_k, cos_

Build the o_proj_residual kernel with int8-quantized wo baked in.

Build block-diagonal attention mask for training. Within each block: bidirectional (non-causal). Between blocks: no attention. Context fe

Convert a DFlash model from HuggingFace format to MLX format. Downloads the model from HuggingFace, converts weights to float16, and saves in

Compiled variant of forward_verifier_states for Qwen3.5.

Forward pass for Qwen3.5 that records rollback data for linear attention. Returns same as forward_with_hidden_states plus rollback_records dict

Compiled forward pass for Qwen3.5 verify with rollback recording. Uses per-layer mx.compile for both linear AND full attention layers. Linear

Qwen3.5 compiled whole-model verify. All 64 layers in one mx.compile graph.

Get or create a compiled full-attention verify function for a layer. Takes explicit KV cache arrays (old_keys, old_values, offset) as inputs

Get or create a compiled whole-model verify function for Qwen3.5. Compiles the entire 64-layer forward pass into a single mx.compile graph. A

Create a prompt cache with TurboQuantKVCache for full-attention layers. For Qwen3.5 (mixed attention), replaces KVCache entries with TurboQuantKV

Phase 1: Generate target model responses and extract hidden states. For each prompt: 1. Generate response using target model (greedy)

Return argmax over the channels (vocab) dimension for each of the first `seq_len` spatial positions. Expects buffer shape [1, vocab_size, 1, w_sq] (f

Roll back Qwen3.5 linear attention caches after partial block rejection. Restores conv_state and SSM state to what they would be after processing

Change the active block size without recompiling kernels. Valid for any new_seq_q in [1, w_sq] (w_sq = align_width(original seq_q)).

Begin sampling ANE power in background.

Stop sampling and return stats.

Phase 2: Train the DFlash draft model.