feat(dflash): add reusable hybrid KDA attention

[maocheng23]

Summary

• add Kimi Delta Attention (KDA) as a reusable DFlash-family attention implementation;
• support per-layer hybrid gqa/mha/mla/kda layouts through dflash_config.attention_modes;
• reuse the same draft backbone and FSDP training path from DFlash, DSpark, and Domino;
• add an optional acceptance-aligned DSpark objective without coupling it to KDA;
• add a Qwen3-4B hybrid KDA/GQA example and offline training recipe.

This PR is intentionally a draft. The implementation, correctness, FSDP, and performance gates pass, but a freshly trained KDA recipe has not yet improved the frozen end-to-end acceptance baseline. The experimental recipe is therefore not promoted as a new default.

Design

KDA and KV/latent attention share a small decoder-facing base class. GQA/MHA/MLA keep their shared cache, mask, and backend machinery, while KDA owns its recurrent projections, causal depthwise convolutions, gated RMSNorm, and optimized/reference dispatch.

Each block-parallel proposal is an independent KDA sequence, so recurrent and convolution state reset at proposal boundaries. Hybrid layouts retain at least one consistent GQA, MHA, or MLA layer for target-context injection.

The optimized KDA backend is optional and imported lazily. Non-KDA installations and existing uniform attention_mode configurations remain backward-compatible.

Correctness

Coverage includes:

• scalar recurrence-oracle parity;
• proposal state isolation and batched/individual equivalence;
• forward/backward and gradient parity;
• hybrid composition through DFlash, DSpark, and Domino;
• malformed-layout and KDA-configuration validation;
• cache non-mutation, checkpoint naming, save/reload, generation, and train/eval equivalence;
• optimized-backend launch splitting and dynamic-shape bucketing;
• real two-rank FSDP forward/backward, optimizer update, auto-wrap, checkpoint materialization, and dtype checks;
• export, resize, packaging, and loss-contract tests.

Exact-stack result: 1,103 passed, 13 skipped, 1 xfailed, 685 subtests.

Optimized/reference maximum absolute differences on H200:

• forward: 0.000244140625;
• gradients: 7.34e-8.

Performance

For a five-layer Qwen3-4B-sized hybrid draft, the selected 2KDA/3GQA layout was 1.031x faster, used 6.10% fewer attention parameters, and used 5.04% less incremental peak memory than the matched GQA draft in the isolated forward/backward benchmark.

The frozen Qwen3-4B speculative-generation harness also verified identical deterministic target outputs. KDA candidates were consistently faster, but the clean training runs did not clear the acceptance gate:

Clean run	GQA acceptance	KDA acceptance	Relative delta	GQA tok/s	KDA tok/s
Off-policy	1.158750	1.149500	-0.80%	30.49	40.97
On-policy LK	1.494447	1.434528	-4.01%	36.72	49.29

A diagnostic checkpoint composition did produce a +7.83% acceptance result with matching target outputs, which motivates further training work, but it is not presented as a clean trained result.

Serving boundary

SpecForge forward, training, save/reload, export, and spec_generate support KDA. Production speculative serving still requires a paired serving-runtime change for recurrent KDA state, so the export normalizer rejects KDA/MLA serving layouts today.

Stack

Built on the MLA attention scaffold merged in #758.