Load a DFlash draft model from a HuggingFace model directory. Args: model_path: Path to local model directory or HuggingFace repo ID

Create proper fa_mask and ssm_mask for Qwen3.5 models.

De-interleave ANE Q/K output back to standard format. Input: [B, n_heads, seq, head_dim] in interleaved format Output: [B, n_heads, seq, head_

Quantize [oc, ic] weight to (int8_bytes, fp16_scale_bytes, oc, ic). Symmetric per-output-channel int8: scale = max(|w|, axis=1) / 127.

1×1 projection with an int8-quantized constant weight baked into the kernel. `w`: weight in row-major `[oc, ic]` order, pre-quantized to int8 with pe

Forward a full-attention layer using compiled function with explicit cache. Only used during verify (hidden_states.shape[1] > 1) when cache is

Forward through a Qwen3.5 linear attention layer recording rollback data. Replicates GatedDeltaNet.__call__ + DecoderLayer structure but saves

Detect if model has linear attention layers (Qwen3.5 architecture).

Remove the last n positions from the cache. After each draft step, trim(n) removes the noise positions (the speculative draft tokens)

Neox (split-half) RoPE: matches mx.fast.rope(traditional=False) used by Qwen3. x_rot = cat(-x[hd/2:], x[:hd/2]) along the W dimension. cos/sin are til

Apply chat template. By default thinking is disabled for DFlash compatibility. Set enable_thinking=True to enable Qwen3/Qwen3.5 thinking mode

Get stop token IDs including EOS and chat-format end tokens.

Run one transformer layer using baked int8 projection weights.

Prepare all ANE buffers on the main thread, then fork a pure-ANE thread. All mx.eval() / Metal calls happen HERE (main thread) before the thr

Get or create a compiled linear attention verify function for a layer. Compiled functions have fixed input shapes (determined by block_size),

Normalize attention head dimensions across Qwen3 / Qwen3.5.

Load DFlashConfig from a HuggingFace config.json.

Load weights from safetensors files using MLX native loader. Handles both single file and sharded (model-00001-of-000NN.safetensors) formats.

Run a single full forward pass, timing every kernel dispatch. Returns dict[kernel_name -> ms].

Compute position-weighted CE loss on masked positions. Args: input_ids: [B, S] token ids (with MASK tokens at positions to predic

Forward pass for training. Returns raw logits (before lm_head).

Forward pass that returns norm'd hidden states before lm_head. Returns (norm_hidden, embed, hidden_states, rollback_records). norm_hidden is

Load weights as W8A16 int8 and compile per-layer quantized kernels. This replaces load_weights() for the quantized inference path. Af

Write all ANE input buffers on the calling (main) thread. Runs all mx.eval() / Metal operations here so that run_prepared() can be ca

Apply per-head RMSNorm. q_4d: [B, n_heads, seq, head_dim] norm_weight: [head_dim] Returns: [B, n_heads, seq, head_dim]

Apply RoPE (neox/half-rotation style, traditional=False) to [B, n_heads, seq, head_dim].

Argmax over vocab from b_logits → token ids [1, seq_q]. Only valid when vocab_size was set at construction. Must be called on the mai

Read b_output into an mx.array. Must be called on the main thread.