Add dflash_laguna EAGLE-3 speculator for Laguna

[pierre427]

Add dflash_laguna EAGLE-3 speculator for Laguna

Stacked on #1223 (Laguna-XS-2.1). This PR depends on the laguna.py model
added there and should be reviewed/merged after it. It reuses Laguna's
embed_tokens and lm_head — it has no embedding or LM head of its own.

What this adds

mlx_lm/models/dflash_laguna.py — an EAGLE-3-style block speculator for
Laguna-XS-2.1 (poolside/Laguna-XS-2.1-DFlash). Unlike a standalone draft model,
the DFlash speculator is target-coupled: it borrows the target Laguna model's
token embedding and LM head, and consumes fused auxiliary hidden states pulled
from selected target layers.

Mechanism (per block)

1. Build a token block [anchor, MASK*(block_size-1)] and embed it via the
   target Laguna embed_tokens.
2. Fuse the target's auxiliary hidden states:
   fuse = hidden_norm(fc(concat(aux_norm_j(aux_j)))).
3. Each draft layer's attention takes queries from the mask block and injects
   the fused target hidden states as extra K/V context
   (k/v = concat(proj(target_hidden), proj(block))), with per-head softplus
   output gating and causal masking within the block.
4. One parallel forward predicts the whole block; block position k predicts
   the token at anchor + k (position 0 reproduces the anchor, 1: are the
   speculative tokens). Draft logits come from the target lm_head.

Because it is target-coupled, Model.__call__ deliberately raises — the model is
driven by the speculative-decoding loop that supplies the target embedding, aux
hidden states, and LM head, not as a standalone causal LM.

Loading / registration

No utils.py change is required: model_type: "dflash_laguna" auto-discovers
mlx_lm/models/dflash_laguna.py via the existing filename-based loader. A
sanitize hook splits a fused qkv_proj checkpoint tensor into q/k/v_proj.

Compatibility with #1223's laguna.py

Verified against the laguna.py in #1223 (not a fork copy). The speculator only
relies on interfaces that base provides:

• LagunaModel.embed_tokens (nn.Embedding) — used to embed the mask block.
• Model.lm_head (nn.Linear) — used to produce draft logits.

No extra hooks beyond these are needed; the fork's Laguna had no additional
surface the #1223 version lacks.

Test

tests/test_dflash_laguna.py builds a small Laguna target plus a small DFlash
speculator and exercises the real forward path:

• embeds [anchor, MASK*(block-1)] through the target embed_tokens,
• fuses aux hidden states (fuse),
• runs the block-draft forward (draft_block) through the draft layers with the
  fused target hidden states injected as K/V context,
• projects the block hidden states through the target lm_head,
• asserts output shapes and NaN-freedom, and asserts that calling the speculator
  as a standalone causal LM raises.

pytest tests/test_dflash_laguna.py  ->  1 passed

black and isort --profile black clean.

[pierre427]

Rebuilt on current main now that #1334 provides the laguna model in-tree — the PR slims from +1082 to +470:

• Dropped: the bundled laguna.py (+499) and its model tests. The speculator now sits directly on the upstream laguna module (the draft test drives fuse()/draft_block() against an in-tree laguna.Model target).
• Kept: models/dflash_laguna.py (EAGLE-3 block speculator for poolside/Laguna-XS-2.1-DFlash), the Laguna tool parser + chat-template inference (its <tool_call>function-name / <arg_key>–<arg_value> dialect, checked before glm47 which shares <arg_key>), and the compressed-tensors fix — now consolidated with the nvfp4 branch that landed since: int pack-quantized checkpoints read group_size/num_bits from config_groups instead of assuming 4-bit/32, nvfp4 and the legacy default unchanged.
• Re-verified: the real poolside/Laguna-XS-2.1-DFlash weights strict-load through sanitize (fused-qkv split) and fuse() + draft_block() produce finite output on real weights; test_dflash_laguna, test_tool_parsing, and the two new _compressed_tensors_config unit tests are green; black clean.

Related: #1704 makes the in-tree laguna load the public target checkpoints this speculator pairs with (independent, not stacked).

🤖 Generated with Claude Code