feat(models): port LLaDA-2 MoE (masked-diffusion LLM, DeepSeek-V3 MoE FFN)

[inureyes]

Summary

Ports LLaDA-2 MoE (model_type: llada2_moe, e.g. inclusionAI/LLaDA2.0-mini): a decoder-only masked-diffusion language model with a DeepSeek-V3-style MoE FFN. It generates by iterative block-wise unmasking of <|mask|> tokens rather than autoregressive decode, so like DiffusionGemma it is a model-owned batch-1 generator (supports_batching() == false) served on the diffusion worker loop.

What changed

New module src/models/llada2_moe/:

• mod.rs — config with the exact serde defaults from the spec, including derived head_dim/rotary_dim and the load-bearing omitted-key fallbacks (eos_token_id -> pad_token_id, mask_token_id -> 156895, use_qk_norm -> true). Bidirectional attention (no causal mask anywhere) with a fused-QKV runtime split (boundaries [num_heads, num_heads + num_kv_heads]), per-head QK-norm, and partial RoPE (rotary_dim of head_dim, non-traditional). Dense SwiGLU MLP for i < first_k_dense_replace, MoE block otherwise. Routed experts reuse switch_layers::SwitchGLU + moe_weighted_sum; the gate follows deepseek_v3::MoEGate with the three LLaDA-2 deltas: bias key expert_bias, + 1e-20 top-k normalization, and a group mask that fills non-kept groups with a large negative constant (not zero, since expert_bias may be negative). Router runs in float32.
• generate.rs — the semi-autoregressive block-unmasking engine: prefill commits full prompt blocks into per-layer KV caches; each generation block is denoised by repeatedly forwarding it against the frozen prefix (read-only KVCache::visible_state) and revealing positions whose chosen-token probability clears a linearly decaying threshold, with a single forced reveal when none clears it (progress guarantee). Schedule and transfer-mask selection are pure host-side functions.
• tests.rs — loop-schedule, config-default, and gate-math unit tests, plus the synthetic parity test and an on-disk tiny checkpoint loaded through the real detection + loading path.

Runtime integration (completion requires the working runtime, not standalone modules): detection.rs, models/mod.rs (variant + registry + display + exhaustiveness), model_metadata.rs, loaded_model.rs (variant + delegation), loading/nonstandard.rs, distributed/tensor_parallel/inference.rs (exclusion), CLI commands/generate.rs + new commands/generate_llada2.rs, commands/chat.rs rejection, and both LoadedModel::DiffusionGemma dispatch points in server/model_worker.rs. The diffusion worker's tokenize/stream/cancel/usage plumbing is factored into one serve_streaming_request helper shared by DiffusionGemma and LLaDA-2 (DiffusionGemma behavior byte-identical); a LLaDA-2 options mapper + worker loop sit beside the gemma ones, with --max-denoising-steps mapping to the per-block step count. Docs row added to docs/supported-models.md.

Also wraps one pre-existing constant assertion in switch_layers tests in a const block so clippy --tests -D warnings stays clean under the pinned 1.93.1 toolchain (byte-identical otherwise; unrelated to LLaDA-2).

⚠️ Real-checkpoint smoke test NOT run (deferred to Phase 3.5)

The issue's stated closure bar includes downloading inclusionAI/LLaDA2.0-mini (~32.5 GB) and generating through the real runtime on both mlxcel generate and mlxcel-server. This was NOT executed on this Linux/GB10 host (it cannot complete within the build/watchdog budget) and remains for orchestrator Phase 3.5 validation. This PR does not imply that bar is met; it delivers everything else (implementation + full integration + synthetic-parity + unit tests, all green locally).

Test plan

• cargo test --lib models::llada2_moe — 20 tests green (loop-schedule, config defaults, gate math, MoE-vs-naive-loop parity, fixed-logits reveal order, on-disk synthetic checkpoint detect + load + generate)
• cargo test --lib server::diffusion_worker — 11 green (6 DiffusionGemma unchanged + 5 new LLaDA-2 option-mapper / media-reject / finish-reason)
• cargo test --lib -- every_variant all_model_types_covers — registry exhaustiveness guards green
• cargo clippy --lib --tests -- -D warnings and cargo clippy --bin mlxcel -- -D warnings clean
• cargo fmt clean
• Real-checkpoint smoke test on inclusionAI/LLaDA2.0-mini — deferred to Phase 3.5 (see caveat above)

Closes #546

[inureyes]

Real-checkpoint validation (orchestrator-run on GB10)

Validated end-to-end against the real inclusionAI/LLaDA2.0-mini checkpoint (30.29 GiB, bf16), through the actual runtime:

• Load: config parsed with the omitted-key defaults, all weight keys mapped (word_embeddings, per-expert stacking, fused QKV), tokenizer and chat template loaded. Model loaded with no key or dtype errors.
• Generate (CUDA backend, sm_121): prompt What is the capital of France? Answer in one word. produced Paris followed by the chat-template end-of-turn token. Coherent and correct.

Notes:

• The CPU backend cannot run this model: MLX's gathered-MoE matmul (gather_mm / gather_qmm) only supports float32 on CPU, while the experts are bf16. This is a general MLX CPU-backend limitation, not specific to this port. Running on an NVIDIA host requires a --features cuda build.
• The synthetic parity test used f32 weights, so it did not exercise the bf16 MoE path. The real-checkpoint run on the CUDA backend covers that gap.

Real-checkpoint closure bar met; moving to status:done.