feat(speculative): add the Laguna DFlash drafter (sliding-window gated attention over captured target hidden states) #1351

Summary

Poolside publishes DFlash speculative drafters for every Laguna release (poolside/Laguna-XS-2.1-DFlash, 0.9 GB bf16; poolside/Laguna-S-2.1-DFlash, 2.2 GB; plus NVFP4/FP8/INT4 variants). mlxcel's DFlash machinery (mlxcel_core::drafter::dflash) is hard-wired to the Qwen 3.5 drafter shape: unfused q/k/v_proj, no QK norm, no output gate, a growing KVCache for the context, and one hidden_norm after fc. The Laguna drafter is a 5- or 6-layer sliding-window decoder with a fused qkv_proj, per-head q_norm/k_norm, a per-head softplus gate, a rotating 511-entry context cache, and one RMSNorm per captured target layer. This change adds a laguna_dflash drafter, teaches the loader to route model_type: "laguna" + dflash_config to it, and implements SpeculativeTarget on the Laguna target so --draft-kind dflash works end to end.

Current behavior

• src/lib/mlxcel-core/src/drafter/dflash/config.rs:274 (is_dflash_drafter_config) already recognizes any config with a dflash_config object, so a Laguna drafter directory is correctly refused as a standalone model by src/models/detection.rs:139. But load_drafter (src/lib/mlxcel-core/src/drafter/mod.rs:1097, DrafterKind::Dflash arm) builds DFlashDraftModel (dflash/model.rs:117), whose from_weights expects layers.N.self_attn.{q,k,v}_proj and has no g_proj, q_norm, k_norm, or aux_hidden_norms, so loading fails on missing keys.
• SpeculativeTarget (src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs:136) is implemented only by the Qwen 3.5 targets; no Laguna target exists (it is introduced by feat(models): add the Laguna family (Poolside Laguna XS/S, hybrid sliding/full MoE with NVFP4 experts) #1347).

Expected behavior

Drafter config.json (XS 2.1 DFlash; S 2.1 in parentheses):

model_type: "laguna"            architectures: ["DFlashLagunaForCausalLM"]
hidden_size: 2048 (3072)        intermediate_size: 8192 (12288)
num_hidden_layers: 5 (6)        num_attention_heads: 64 (72)   num_key_value_heads: 8   head_dim: 128
rms_norm_eps: 1e-6              vocab_size: 100352   draft_vocab_size: 100352 (must equal vocab_size)
max_position_embeddings: 262144 (1048576)   rope_theta: 500000.0   (no rope_parameters: plain RoPE on all 128 dims)
sliding_window: 512             layer_types: ["sliding_attention"] * num_hidden_layers
sliding_windows: [512] * L      (absent on XS 2.1; default to [sliding_window] * L when missing)
gating: "per-head"              num_experts: 0   attention_bias: false
eagle_aux_hidden_state_layer_ids: [2, 14, 26, 34, 40] ([2, 11, 20, 30, 39, 48])
dflash_config: {block_size: 16, mask_token_id: 12, num_target_layers: 40 (48),
                target_layer_ids: [1, 13, 25, 33, 39] ([1, 10, 19, 29, 38, 47]), causal: true}

Validation at load: len(target_layer_ids) == num_hidden_layers, strictly increasing, each < num_target_layers; num_target_layers == target.num_layers(); draft_vocab_size == vocab_size == target vocab; 0 <= mask_token_id < vocab_size (12 is 〈|MASK|〉); causal must be true; gating == "per-head"; every layer_types[i] == "sliding_attention".

Weights (XS 2.1; bf16, no model. prefix, no embedding, no lm_head):

fc.weight                               [2048, 5 * 2048]
hidden_norm.weight                      [2048]
norm.weight                             [2048]
aux_hidden_norms.{i}.weight             [2048]            i in 0..5
layers.{i}.input_layernorm.weight       [2048]
layers.{i}.post_attention_layernorm.weight [2048]
layers.{i}.self_attn.qkv_proj.weight    [(64 + 2*8) * 128 = 10240, 2048]   rows: q (8192), k (1024), v (1024)
layers.{i}.self_attn.o_proj.weight      [2048, 8192]
layers.{i}.self_attn.g_proj.weight      [64, 2048]
layers.{i}.self_attn.q_norm.weight      [128]
layers.{i}.self_attn.k_norm.weight      [128]
layers.{i}.mlp.{gate_proj,up_proj}.weight [8192, 2048]   layers.{i}.mlp.down_proj.weight [2048, 8192]

Sanitize: strip a leading model.; if a checkpoint ships split self_attn.{q,k,v}_proj.weight, concatenate them along axis 0 in q, k, v order into qkv_proj.weight; reject a layer that has both or an incomplete split set; after sanitize the key set must equal the list above exactly (quantized variants add .scales/.biases next to each .weight).

Forward for one draft block (B = 1, L = block_size = 16):

target capture: after target layer t_i (0-based, t_i in target_layer_ids) record the residual stream h_i [B, T, H]
ctx = hidden_norm( fc( concat_i aux_hidden_norms[i](h_i) ) )          # [B, T, H]; T = new target positions since the last accept
block = [bonus, mask_id, mask_id, ..., mask_id]                         # [B, L]
x = embed_tokens_target(block) * embed_scale (1.0)                      # [B, L, H]
for each drafter layer i with rotating cache c_i (max_size = sliding_window - 1 = 511, keep = 0):
    ctx_i = ctx;  if ctx_len >= 512: drop the oldest so ctx_len = 511 and c_i.offset += dropped
    h = RMSNorm(x; input_layernorm)
    qkv = qkv_proj(h);  ctx_qkv = qkv_proj(ctx_i)
    q = qkv[..., :8192] -> [B, 64, L, 128];  k_p = qkv[..., 8192:9216] -> [B, 8, L, 128];  v_p = qkv[..., 9216:] -> [B, 8, L, 128]
    k_c = ctx_qkv[..., 8192:9216] -> [B, 8, T, 128];  v_c = ctx_qkv[..., 9216:] -> [B, 8, T, 128]
    q = RMSNorm_128(q; q_norm);  k_p = RMSNorm_128(k_p; k_norm);  k_c = RMSNorm_128(k_c; k_norm)      # per head, before RoPE
    rope (base 500000, all 128 dims, rotate-half): k_c at positions c_i.offset .. c_i.offset + T
                                                   q and k_p at positions c_i.offset + T .. + L
    K_c, V_c = c_i.update_and_fetch(k_c, v_c)               # only the context is cached
    K = concat(K_c, k_p, axis=2);  V = concat(V_c, v_p, axis=2)
    mask = [ ones(L, len(K_c)) | tril(ones(L, L)) ]           # causal inside the proposal block, full view of the context
    o = SDPA(q, K, V, scale = 128^-0.5, mask) -> [B, L, 8192]
    gate = softplus(g_proj(h).astype(f32)).astype(o.dtype)    # [B, L, 64]
    o = (o.reshape(B, L, 64, 128) * gate[..., None]).reshape(B, L, 8192)
    x = x + o_proj(o)
    x = x + down_proj(silu(gate_proj(n)) * up_proj(n)),  n = RMSNorm(x; post_attention_layernorm)
logits = lm_head_target( RMSNorm(x; norm) )[:, 1:]             # L - 1 draft tokens, sampled per position

Acceptance and rollback follow the existing DFlash round loop: the target verifies [bonus, d_0, ..., d_{14}] in one forward with hidden capture, the accepted prefix length trims the target caches (KV on full layers; RotatingKVCache on sliding layers, which must still be unwrapped or support trim from the tail), and the drafter's rotating context caches are advanced only by the accepted target positions.

Implementation plan

1. src/lib/mlxcel-core/src/drafter/laguna_dflash/{config.rs, model.rs, drafter.rs, mod.rs} (new). LagunaDFlashConfig::from_json implements the field table and validation above (sliding_windows defaults to [sliding_window; L]). Keep DFLASH_DRAFT_ARCHITECTURE detection; add "DFlashLagunaForCausalLM" to the recognized architectures in dflash/config.rs:274 and route inside load_drafter (mod.rs:1097): when the drafter config's model_type == "laguna" build LagunaDFlashDrafter, else the Qwen 3.5 DFlashDraftModel.
2. LagunaDFlashAttention (fused qkv_proj: UnifiedLinear, o_proj, g_proj, q_norm, k_norm: RMSNorm) with the forward above; reuse mlxcel_core::fast_rope (base 500000, dims 128, offset per the cache), mlxcel_core::layers::RotatingKVCache::new(511) per layer, mlxcel_core::softplus. The causal block mask is the same shape the Qwen drafter builds in dflash/attention.rs; reuse its helper.
3. LagunaDFlashDraftModel: aux_hidden_norms: Vec<RMSNorm>, fc: UnifiedLinear, hidden_norm, norm, layers. bind_target_embedding / bind_target_lm_head as in dflash/model.rs:216 and :229; validate_target_compat enforces the layer-count and vocabulary checks above. draft_block(last_bonus, target_hidden_concat, caches, block_size, sampler) mirrors dflash/model.rs:352 but splits target_hidden_concat into num_layers slices of hidden_size and applies aux_hidden_norms[i] before fc.
4. Target side (src/models/laguna.rs): implement SpeculativeTarget (capture_layer_ids from the drafter config via verify_forward_with_capture_layers, verify_forward capturing the residual stream after each listed layer, rollback_partial trimming KVCache and RotatingKVCache by block_size - (accepted + 1)). Register the Laguna target in src/server/speculative_dispatch.rs so --draft-kind dflash --draft-model models/Laguna-XS-2.1-DFlash resolves; the CLI flags already exist (src/cli/speculative_args.rs).
5. docs/supported-models.md: note the Laguna DFlash drafter next to the Qwen 3.5 one.

Validation

(a) Unit tests, src/lib/mlxcel-core/src/drafter/laguna_dflash/tests.rs:

• config_rejects_bad_contract: each of causal: false, draft_vocab_size != vocab_size, non-increasing target_layer_ids, gating: "per-element", a full_attention entry in layer_types, and num_target_layers mismatch produces a typed load error; the XS 2.1 config (without sliding_windows) loads.
• sanitize_fuses_split_qkv_and_rejects_mixed: split q/k/v keys fuse into qkv_proj in q, k, v row order; a layer with both fused and split keys is rejected; the resulting key set equals the expected list.
• context_trim_keeps_window_minus_one: with a 600-position context the drafter feeds 511 context positions and advances the cache offset by 89.
• draft_block_is_causal_inside_the_proposal: changing mask position j changes logits only at positions >= j.
• greedy_invariant_on_synthetic_target: a tiny random Laguna target plus a tiny random drafter produce the same greedy tokens with and without speculative decoding (the existing round_loop.rs SyntheticTarget pattern).

(b) Real checkpoints: models/Laguna-XS-2.1-NVFP4 (target, 21.5 GB) and models/Laguna-XS-2.1-DFlash (drafter, 0.9 GB bf16):

./target/release/mlxcel generate -m models/Laguna-XS-2.1-NVFP4 --draft-model models/Laguna-XS-2.1-DFlash --draft-kind dflash -p "Write a Python retry wrapper with exponential backoff." -n 128
./target/release/mlxcel-server -m models/Laguna-XS-2.1-NVFP4 --draft-model models/Laguna-XS-2.1-DFlash --draft-kind dflash

Acceptance: greedy output with the drafter is token-identical to greedy output without it for 128 tokens on three code prompts; the reported mean accepted length is above 3 tokens per round on the code prompt above; decode throughput is higher than the non-speculative run on the same machine (report the numbers in the PR).

Acceptance criteria

• A Laguna DFlash drafter directory loads through load_drafter and is rejected as a standalone model
• Greedy speculative output matches non-speculative output token for token on Laguna XS 2.1 (PR feat(speculative): add the Laguna DFlash drafter #1771: identical on 1 of 6 prompts on a GB10 with the NVFP4 target; the others differ at one or two bf16 logit-tie positions per 128 tokens, where the multi-row verify kernels and single-token decode round differently; the tiny f32 model is token-exact with forced rejections)
• Rollback after a partial accept trims both full and sliding target caches and the drafter context caches correctly (covered by the greedy-invariant test with forced rejections)
• docs/supported-models.md updated
• detection table in src/models/detection.rs updated with a test
• cargo test --workspace --profile test-fast --features metal,accelerate passes (not runnable on the Linux/CUDA host that implemented PR feat(speculative): add the Laguna DFlash drafter #1771; the CUDA-scoped equivalents pass)
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass (run as --lib --bins --tests --features cuda per crate on the CUDA host)

Out of scope

• Batched (B > 1) Laguna DFlash; the B = 1 round loop only. Batched follows the round_loop_batched.rs pattern in a later issue.
• Quantized drafter variants beyond what UnifiedLinear loads from .scales/.biases already.

Dependencies

• Blocked by feat(models): add the Laguna family (Poolside Laguna XS/S, hybrid sliding/full MoE with NVFP4 experts) #1347.