DFlash: add Laguna XS 2.1 support

[joelfarthing]

Adds DFlashLagunaForCausalLM support to the DFlash speculative path. New behavior is gated on the dflash-draft.dflash.laguna GGUF key. The diff is behavior-preserving rather than purely additive: it modifies a few shared lines (chiefly the capture path below), and existing DFlash drafts are verified byte-identical vs the base build on gpt-oss (spec and non-spec, including under a small -ub that splits prompt processing into multiple microbatches).

One shared change is notable. DFlash seeds the draft's cross-attention from target hidden-state rows captured during prompt processing. When the prompt is longer than -ub, that processing is split into internal microbatches and the capture callback fires per microbatch; the previous code re-sized the row buffer to the current microbatch and wrote it at offset 0, so each microbatch overwrote the last and only the final one survived. That's enough for existing drafts but not for Laguna, whose target graph reduces to the selected output row at the final layer, so capture only saw one row and multi-token prompts silently disabled speculation (and smaller -ub shifted acceptance depth). The fix sizes the buffer once to the full logical batch (n_tokens_all) and writes each microbatch at its running offset, accumulating across microbatches, with guards for width mismatch, row overflow, and incomplete capture, plus a sync of the async device readback before the buffer is read. A paired change keeps the Laguna target's full final-layer prompt rows when capture is active. Capture is now complete and -ub-invariant, and existing drafts are unaffected (byte-identical, per above).

Draft: official 462M -> Q8_0. Target: official Q4_K_M. Greedy, RTX 4070 + Core i7-11700K.

Throughput by content (server, CPU target, n_max=4, TG tok/s, same prompt):

Fixture	Kind	no-spec	DFlash	Ratio
double-linked-list	code	19.7	35.9	1.82x
quicksort	code	20.0	29.9	1.50x
sayap	table	19.8	28.2	1.43x
bulgaria	narrative	20.8	16.6	0.80x

Proposal-depth economics (doubly-linked list, TG tok/s, ratio vs that config's no-spec):

n_max	CUDA, draft GPU	CUDA, draft CPU	CPU target
no-spec	46.7 (1.00)	46.7 (1.00)	20.3 (1.00)
1	51.1 (1.10)	37.7 (0.81)	23.2 (1.14)
2	51.3 (1.10)	43.3 (0.93)	26.5 (1.31)
4	49.2 (1.05)	44.3 (0.95)	30.5 (1.50)
8	36.7 (0.79)	36.7 (0.79)	27.9 (1.37)
16	28.9 (0.62)	29.4 (0.63)	18.5 (0.91)
autotune (#2112)	44.1 (0.94)	n/a	26.8 (1.32)

Optimum depth tracks target speed: a CPU target favors n_max=4, a CUDA target favors n_max=1-2 and loses past n_max=8. The draft must sit on GPU when the target does; a CPU draft loses at every depth. --spec-autotune converges near the best fixed value but does not beat it, and on the fast CUDA target its exploration overhead makes it a net loss (as expected; autotune is intended for convenience, worked well on these tests, and is not expected to beat manual tuning).

Validation:

Output is coherent and correct but not guaranteed bit-identical to non-speculative decoding on longer greedy runs: batched block verification flips occasional near-tie argmaxes in the target (a target-model property, reproducible with no draft by varying -ub; generic DFlash targets are unaffected).

Ready-to-test Q8_0 draft: https://huggingface.co/ji-farthing/Laguna-XS-2.1-DFlash-ik-llama-GGUF. Pair with official https://huggingface.co/poolside/Laguna-XS-2.1-GGUF.

• I have read the contributing guidelines
• Self-reported review complexity:
  • Low
  • Medium
  • High

[SamuelOliveirads]

LGTM.

I didn't get good performance with DFlash, but that's not related to the implementation, although I did notice a limitation for 2+ GPUs that I'll figure out how to resolve in a future PR.

Now speaking specifically about Laguna: The context-KV precompute builds each draft-layer projection independently. The vLLM implementation groups the five Laguna projections and RMSNorm operations into batched operations, reducing graph nodes and kernel launches. This is executed for every DFlash context update, so batching this path would be a follow-up optimization.

[ikawrakow]

Will merge when back in a few days.

[jeverson-bmi]

Note that I also didn't get good performance with Laguna S 2.1 and Dflash on llama-cpp, to confirm your experience @SamuelOliveirads

[SamuelOliveirads]

@ikawrakow can we merge it?