[Spec Decode][Perf] Optimize padded EAGLE input prep and EAGLE3 layer0 RMSNorm concat

[tanpinsiang]

This PR improves the V1 EAGLE/EAGLE3 speculative decode path.

It replaces a small padded EAGLE input-prep Triton helper with torch tensor ops, and adds a guarded fused EAGLE3 layer-0 helper for:

cat([input_layernorm(embeds), hidden_norm(hidden_states)], dim=-1)

What changed:

• Use torch ops for padded EAGLE input preparation in SpecDecodeBaseProposer.
• Remove the unused Python Triton padded-input helper.
• Add a fused EAGLE3 dual-RMSNorm + concat Triton helper.
• Enable the fused path only for compatible llama_eagle3 layer-0 inputs.
• Preserve the original unfused fallback path.

How to enable:

• No new user-facing flag.
• The padded input-prep change applies automatically to the existing V1 padded EAGLE path.
• The EAGLE3 RMSNorm concat fusion is automatically used only when tensors are GPU, contiguous, same shape/device/dtype, and RMSNorm eps matches.

Scope and limitations:

• The padded input-prep change is general V1 EAGLE path code.
• The fused RMSNorm concat path is specific to the llama_eagle3 layer-0 draft model path.
• The fused helper uses Triton and falls back to the original torch/RMSNorm path when guards do not match.

Verification

FP4 EAGLE3 MTP

Benchmark random input 8192 / output 1024 / random-range-ratio 0.8 / concurrent 64 / prompts 640, TP4.

Model: amd/MiniMax-M3-MXFP4
Draft model: Inferact/MiniMax-M3-EAGLE3

Note: These benchmark numbers were collected before recipes#621 clarified the nested text_config override. The command below is updated to the working form; the A/B delta remains comparable because both arms used the same config, but absolute tok/s should be rerun with nested index-topk reuse.

Baseline checkout

Run from the baseline checkout. For the isolated EAGLE3 layer0 fusion comparison, this was the parent of the fusion commit, i.e. padded EAGLE input prep included but RMSNorm concat fusion removed.

export PORT=8010
export HIP_VISIBLE_DEVICES=0,1,2,3

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4
export VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16=0
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB=256
export VLLM_USE_BREAKABLE_CUDAGRAPH=0

vllm serve amd/MiniMax-M3-MXFP4 \
  --port "$PORT" \
  --tensor-parallel-size 4 \
  --trust-remote-code \
  --block-size 128 \
  --no-enable-prefix-caching \
  --language-model-only \
  --max-model-len 32768 \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 256 \
  --attention-backend TRITON_ATTN \
  --moe-backend aiter \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --kv-cache-dtype fp8 \
  --hf-overrides '{"text_config": {"use_index_cache": true, "index_topk_freq": 4}}' \
  --speculative-config '{"method":"eagle3","model":"Inferact/MiniMax-M3-EAGLE3","num_speculative_tokens":3}' \
  --no-async-scheduling

This PR: Run the same command from this PR checkout.
same environment and vllm serve command as baseline
Common benchmark command:

vllm bench serve \
  --backend openai-chat \
  --base-url "http://127.0.0.1:${PORT}" \
  --model amd/MiniMax-M3-MXFP4 \
  --endpoint /v1/chat/completions \
  --dataset-name random \
  --random-input-len 8192 \
  --random-output-len 1024 \
  --random-range-ratio 0.8 \
  --request-rate inf \
  --max-concurrency 64 \
  --num-prompts 640 \
  --ignore-eos \
  --temperature 0 \
  --save-result

Metric	Baseline	This PR	Improvement
Completed / failed	640 / 0	640 / 0	same
Output throughput / GPU	212.514 tok/s	348.736 tok/s	+64.10%
Total throughput / GPU	2001.188 tok/s	3283.959 tok/s	+64.10%
Mean TPOT	69.584 ms	43.819 ms	37.03% lower
Mean TTFT	2173.242 ms	2228.905 ms	2.56% higher (TTFT regression)
Acceptance	2.895%, length 1.087	56.467%, length 2.694	improved
Duration	773.193 s	471.170 s	39.06% faster

GSM8K

uvx --from 'lm-eval[api]' lm_eval \
  --model local-chat-completions \
  --model_args model='amd/MiniMax-M3-MXFP4',base_url='http://127.0.0.1:8010/v1/chat/completions',tokenizer='amd/MiniMax-M3-MXFP4',num_concurrent=200 \
  --tasks gsm8k \
  --num_fewshot 25 \
  --apply_chat_template \
  --limit 200 \
  --batch_size 1 \
  --gen_kwargs max_tokens=4096,temperature=0 \
  --output_path /tmp/feature18_candidate_lm_eval_gsm8k_25shot_c200_n200_chattemplate

Result:

• Strict exact match: 0.945 +/- 0.0162
• Flexible exact match: 0.945 +/- 0.0162

[tanpinsiang]

superseded by #47854