[vLLM-ATOM] Support EAGLE3 spec decoding for MiniMax-M3

[kliuae]

Motivation

This PR adds EAGLE 3/3.1 speculative decoding support for MiniMax-M3 in vllm-atom.
It depends on and currently contains changes from #1201, and will be rebased once the dependency gets merged.

Technical Details

Test Plan

Target model: MiniMaxAI/MiniMax-M3-MXFP8
Draft model: Inferact/MiniMax-M3-EAGLE3

Server command

VLLM_ROCM_USE_AITER=1 VLLM_USE_V1=1 \
vllm serve MiniMaxAI/MiniMax-M3-MXFP8 \
    --host localhost \
    --port 8000 \
    --no-trust-remote-code \
    --dtype auto \
    --kv-cache-dtype auto \
    --max-model-len 32768 \
    --max-num-batched-tokens 32768 \
    --no-enable-prefix-caching \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
    --tensor-parallel-size 4 \
    --block-size 128 \
    --additional-config '{"online_quant_config": {"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "vision_tower", "multi_modal_projector", "patch_merge_mlp", "*block_
sparse_moe"]}}' \
    --language-model-only \
    --speculative-config.method eagle3 \
    --speculative-config.model Inferact/MiniMax-M3-EAGLE3 \
    --speculative-config.num_speculative_tokens 3

lm_eval

lm_eval --model local-completions --model_args model=MiniMaxAI/MiniMax-M3-MXFP8,base_url=http://0.0.0.0:8000/v1/completions,tokenizer_backend=None,tokenized_requests=False,num_concurrent=128,max_retries=10,max_gen_toks=2048,timeout=60000 --batch_size auto --tasks gsm8k --num_fewshot 20

Test Result

lm_eval, EAGLE=3

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	_	0.9431	_	0.0064
		strict-match	20	exact_match	_	0.9431	_	0.0064

gsm8k: Per-position acceptance rate: 0.876, 0.752, 0.626, Avg Draft acceptance rate: 75.1%

Target model: MiniMaxAI/MiniMax-M3-MXFP8
Draft model: Inferact/MiniMax-M3-EAGLE3

Mode	ISL/OSL	Conc	EAGLE3	Req/s	delta	Total tok/s	delta	Output tok/s	delta	median TTFT (ms)	delta	TPOT mean/med/p99 (ms)	delta	ITL mean/med/p99 (ms)	delta
ATOM native	1k/1k	64	no	4.235	-	8672.6	-	4336.3	-	906	-	13.84	-	13.73	-
vLLM plugin	1k/1k	64	no	3.785	-10.6%	7752.2	-10.6%	3876.1	-10.6%	1487	+64%	14.36	+3.8%	14.02	+2.1%
ATOM native	1k/1k	64	spec=3	4.950	-	10144.5	-	5075.6	-	149	-	9.37	-	24.51	-
vLLM plugin	1k/1k	64	spec=3	4.515	-8.8%	9247.7	-8.8%	4623.8	-8.9%	177	+19%	10.24	+9.3%	25.38	+3.6%

Submission Checklist

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

[whx-sjtu]

Why do we need this? Why can't just let vllm side handle the management of kv cache groups?

[kliuae]

This is to match native ATOM's use case of MLA target with MHA draft (e.g. Kimi-K2.6 with lightseekorg/kimi-k2.6-eagle3 drafter). This requires separate kv cache groups for MLA and MHA respectively, but this configuration is currently not supported in vllm which crashes when it tries to unify page sizes of its cache pools. Therefore, we build separate kv cache groups for plugin and patch their peripheries like memory usage calculation to accommodate this use case.

[zejunchen-zejun]

Hi, @kliuae
Please help resolve the format issue
Thank you

[zejunchen-zejun]

Hi, @kliuae
Could you help rebase this PR? We plan to merge it soon.
Thank you

[ZhangLirong-amd]

please fix the ci code format