[vLLM-ATOM] Support Eagle 3.1 spec decoding

[kliuae]

Motivation

This PR adds EAGLE 3/3.1 speculative decoding support to vLLM-ATOM and its support for Kimi-K2.6.
Supported combinations are MLA target with MLA/MHA draft models. In the case of MLA target with MHA draft, separate KV cache pools are built so that vLLM doesn't attempt to unify the page sizes.

Technical Details

Test Plan

Target model: amd/Kimi-K2.6-MXFP4
Draft model: lightseekorg/kimi-k2.6-eagle3.1-mla (MLA), lightseekorg/kimi-k2.6-eagle3 (MHA)

Server command

vllm serve amd/Kimi-K2.6-MXFP4 \
    --host localhost \
    --port 8000 \
    --async-scheduling \
    --trust-remote-code \
    --gpu_memory_utilization 0.8 \
    --kv-cache-dtype fp8 \
    --no-enable-prefix-caching \
    --enable-auto-tool-choice \
    --tool-call-parser kimi_k2 \
    --reasoning-parser kimi_k2  \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
    --tensor-parallel-size 8 \
    --speculative-config.method eagle3 \
    --speculative-config.model "$DRAFT_MODEL" \
    --speculative-config.num_speculative_tokens 3

lm_eval

lm_eval --model local-completions --model_args model=amd/Kimi-K2.6-MXFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 5

Test Result

lm_eval with gsm8k

lightseekorg/kimi-k2.6-eagle3.1-mla

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	_	0.9416	_	0.0065
		strict-match	5	exact_match	_	0.9409	_	0.0065

Per-position acceptance rate: 0.800, 0.602, 0.429, Avg Draft acceptance rate: 61.0%

lightseekorg/kimi-k2.6-eagle3

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	_	0.9348	_	0.0068
		strict-match	5	exact_match	_	0.9356	_	0.0068

Per-position acceptance rate: 0.819, 0.605, 0.388, Avg Draft acceptance rate: 60.4%

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[zejunchen-zejun]

maybe we need to raise some warning here if import path doesn't exist when upgrading vllm to avoid the silent return

[kliuae]

Good point, added a warning message.