feat(engine): MTP training loss and speculative decoding for bailing_v3

[Dayuxiaoshui]

What does this PR do?

Wires bailing_v3's native multi-token-prediction (MTP) head into both halves of the train-infer loop. Until now the MTP layer shipped in Ling / Bailing V3 checkpoints was dropped at load time, so it degraded silently during RL fine-tuning and the rollout engine could not use it to draft tokens.

Phase 0, MTP training loss. The checkpoint spec loads and saves the MTP layer through a generic ExtraLayerListSpec, the model runs it as BailingMTPLayer when TrainMeta.mtp_enabled is set, and an auxiliary t+2 loss is added when runtime.mtp_loss_scale is given explicitly. The checkpoint's mtp_loss_scaling_factor is never inherited: the loss trains toward every response token and would push preference objectives such as DPO toward their rejected rows.

Phase 2, rollout speculative decoding, runtime.speculative_draft_tokens = k. One decode forward verifies k + 1 tokens per row via InferMeta.tokens_per_seq (flash backend only, the native backend refuses it). KDA linear-attention layers reuse the recurrent kernel's existing target-verify mode, with their recurrent and conv caches stacked across layers so the commit is four kernels and one shared verify-state buffer serves every CUDA-graph bucket. The MTP layer drafts EAGLE-style with its own paged KV cache; SpeculativeDraftModel in areno/models/base.py is the model-side protocol. areno/engine/runtime/speculative.py implements the sampler-equivalent target distribution and exact chain rejection sampling, so the sampled distribution and the reported logprobs are unchanged. Verify and draft forwards have per-bucket CUDA graphs, and the rollout logs the mean accept length.

Also fixes _cancel_stop_token crashing on ignore_eos=True (the CUDA backend passes eos_token_id=(); that path had never worked) and stops the cancel fallback marker from becoming a truncation stop token. Design and method: docs/issues/mtp-speculative-decoding.md.

Results on inclusionAI/Ling-3.0-tiny-base, 1x H200, flash backend

MTP head wiring (Phase 0), natural text, 24-layer trunk plus one MTP layer:

Predictor	NLL
Trunk, token t+1	0.77
Pretrained MTP head, token t+2	0.78
Uniform baseline	12.0

On-policy accept probability of the MTP head at temperature 1.0 (closed form, alpha = sum min(p, q), 8192 gsm8k positions):

Metric	Value
alpha, depth 1	0.896
greedy top-1 agreement	0.909
alpha, depth 2 given depth 1 accepted	0.417
expected accept length, k=1 / k=2	1.90 / 2.27

Decode phase inside the worker (sum of decode steps, 256 forced tokens per row, temperature 1.0, k=2; the client wall clock also carries ~10 s per session of weight fuse/offload and is not a decode measure):

Concurrent rows	Plain decode	Speculative decode	Accept length	Speedup
1	1.66 s (6.5 ms/step)	1.05 s (9.3 ms/step)	2.3	1.58x
4	1.85 s	1.19 s	2.3	1.55x
8	2.20 s	1.37 s	2.3	1.61x
16	2.47 s (10.1 ms/step)	1.59 s (13.2 ms/step)	2.3	1.56x
32	2.65 s	1.70 s	2.4	1.56x
64	3.24 s	2.35 s	2.4	1.38x

k=1 at 16 rows gives 1.29x (accept length 1.87), so k=2 is the default worth using with this single-layer head. Verify at 3 x rows tokens costs 1.2x to 1.3x a plain step, matching the memory-bound L(B) curve in the RFC.

Per-step optimizations that brought the speculative step from 16.3 ms to 13.2 ms at 16 rows:

Change	Step time
Eager verify, torch.multinomial, chunked logprobs, per-layer commit	16.3 ms
Removed nonzero / boolean-index host syncs, one-pass logprobs, Gumbel-max draws	15.3 ms
KDA caches stacked across layers, commit in 4 kernels	13.2 ms

Related issue

No tracking issue; the RFC at docs/issues/mtp-speculative-decoding.md holds the motivation, method and results.

Type of change

• 🐛 Bug fix
• ✨ New feature
• 💥 Breaking change (public API / CLI behavior changes in a non-backward-compatible way)
• 📝 Documentation update
• ♻️ Refactoring
• ⚡ Performance improvement
• ✅ Test coverage improvement

How was it tested?

CPU (no GPU): pytest tests/test_speculative_cpu.py tests/test_speculative_rollout_cpu.py tests/test_inference_scheduler_cpu.py tests/test_bailing_kv_cache_cpu.py tests/test_bailing_v3_mtp_cpu.py tests/test_cancel_stop_token_cpu.py -k cpu, 161 passed together with the neighbouring LoRA, routing-replay, registry, recompute, config and engine-role suites. They cover sampler equivalence with and without top-k/top-p, distribution preservation of rejection sampling (1.5% over 40k trials, exact under greedy), the real decode loop driven by a deterministic fake model (accepted and rejected drafts, stop token inside a step, length cap, continuous-batching admission, k=1), MTP layer forward and checkpoint round trip, and the cancel-stop regression.

GPU (one H200, Ling-3.0-tiny-base, not in the CPU suite):

Check	Result
Verify forward over 3 fed tokens vs 3 sequential decodes	KL <= 8e-3 per position (prefill-vs-sequential noise floor: 2e-2)
Recurrent, conv and KV state after full and partial commits	Match sequential decode
MTP draft forward (prefill mode) vs Phase 0 training-path MTP logits	100% argmax agreement
Greedy speculative vs greedy plain rollout, 16 gsm8k prompts	12/16 identical, 4 diverge late at bf16 near-ties
Eager vs CUDA-graph speculative rollout	Bit-identical
Load test: EOS on, top-p 0.95, 64 prompts x 2 samples, 512 tokens, 32 running	65k tokens, finite logprobs, consistent lengths, both modes

Hardware limitations: TP=1 on a single GPU with the flash backend; TP>1 is only exercised through the CPU tests' single-rank context; the native backend rejects tokens_per_seq > 1.

Checklist

• The PR title summarizes the contribution.
• Linked the related issue in the description (if any): none exists, the RFC is linked.
• Existing tests pass (pytest tests/ -k cpu): all touched suites pass; the three test_moe_sequence_parallel_cpu and two test_policy_tensor_sync_cpu failures reproduce on the base commit and are unrelated.
• New behavior is covered by tests.
• Described the test commands run and any hardware limitations.
• Public API / CLI changes are additive and backward-compatible (see CONTRIBUTING.md): two new RuntimeConfig fields with inert defaults (mtp_loss_scale=None, speculative_draft_tokens=0), one optional CausalLMOutput field, no CLI changes.

[xsuler]

@Dayuxiaoshui Here are some points need to be resolved.

1. HBM leak during role switching: model-level _kda_xxx and _spec_xxx tensors are plain attributes. clear_kv_caches() does not release them, model.to("cpu") does not move them, and offload_kv_caches() ignores _spec_xxx. Large rollout buffers will remain on GPU.
2. Unused CUDA graphs: speculative mode still captures ordinary decode graphs for every bucket, although the speculative loop never uses them. This wastes initialization time and HBM.
3. Configuration is not exposed: mtp_loss_scale and speculative_draft_tokens are absent from TrainerConfig and areno train.
4. Disabled MTP still has overhead: MTP layers are always built, loaded, and policy-synchronized, including reference workers and rollouts with speculative decoding disabled.