[ROCm][Spec Decode] Fix EAGLE3 layer0 acceptance collapse

[tanpinsiang]

Summary

This PR fixes a ROCm-specific EAGLE3 speculative decoding correctness/robustness issue in the first draft decoder layer input-prep path.

The issue is backend-sensitive. ROCm EAGLE3 does not always collapse: forcing the draft model to use ROCM_AITER_UNIFIED_ATTN is a strong workaround.
This PR fixes the ROCm layer0 input-prep path so the default path is usable, and it also improves acceptance when ROCM_AITER_UNIFIED_ATTN is explicitly used.

The change is gated to ROCm with current_platform.is_rocm(). Non-ROCm platforms keep the existing path.

Problem

EAGLE3 depends on the draft model staying aligned with the target model. On ROCm, the first EAGLE3 draft layer currently builds its input with separate RMSNorm and concat operations:

embeds = self.input_layernorm(embeds)
hidden_states, residual = self._residual_norm(hidden_states=hidden_states)
hidden_states = torch.cat([embeds, hidden_states], dim=-1)

In the default ROCm backend route, the drafter can select ROCM_ATTN. With that path, acceptance can collapse enough that speculative decoding stops helping.

ROCM_AITER_UNIFIED_ATTN avoids much of the collapse and is a good operational workaround, but it is not a substitute for fixing the ROCm layer0 path. In same-batch testing, this PR still improves acceptance and throughput even with ROCM_AITER_UNIFIED_ATTN.

Fix

For ROCm only, this PR uses a fused layer0 RMSNorm + concat path for EAGLE3:

residual = hidden_states
hidden_states = fused_dual_rmsnorm_cat(
    embeds,
    hidden_states,
    self.input_layernorm.weight,
    self.hidden_norm.weight,
    self.input_layernorm.variance_epsilon,
)

Non-ROCm platforms keep the existing fallback path unchanged.

Validation

MI355X, vllm/vllm-openai-rocm:nightly, amd/MiniMax-M3-MXFP4, EAGLE3 draft model Inferact/MiniMax-M3-EAGLE3.

Benchmark: random 8192 input / 1024 output, 640 prompts, max concurrency 64, --num-warmups 2.

Run	Drafter backend path	Acceptance	Accepted length	Output tok/s
baseline d63c8e9, spec3, default backend	default -> ROCM_ATTN	2.83%	1.08	861.12
baseline d63c8e9, spec3, ROCM_AITER_UNIFIED_ATTN	explicit unified	36.69%	2.10	2192.86
this PR, spec3, ROCM_AITER_UNIFIED_ATTN	explicit unified	56.19%	2.69	2472.26
baseline d63c8e9, spec3, TRITON_ATTN	explicit TRITON_ATTN	58.27%	2.75	2504.88
this PR, spec3, TRITON_ATTN	explicit TRITON_ATTN	58.63%	2.76	2457.61

NVIDIA Regression Check

NVIDIA validation was also run against PR.
Output throughput changed by -0.66%, which is within normal run noise for this benchmark.

[tanpinsiang]

@hongxiayang

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[hongxiayang]

I added a comment to the other PR you were preparing: #47882 (comment)

[tanpinsiang]

Update: vllm-project/recipes#615 now pins the MiniMax-M3 ROCm EAGLE3 recipe to TRITON_ATTN.

With explicit TRITON_ATTN, baseline is already healthy:

Run	Acceptance	Accepted length	Output tok/s
baseline d63c8e9, spec3, TRITON_ATTN	58.27%	2.75	2504.88
this PR, spec3, TRITON_ATTN	58.63%	2.76	2457.61

So this PR is not required for the recipe path.

The PR still fixes/improves other ROCm paths:

• default backend -> ROCM_ATTN: 2.83% acceptance baseline
• this PR restores the default path in prior testing
• ROCM_AITER_UNIFIED_ATTN: 36.69% -> 56.19% acceptance

I’m happy to keep this open as a ROCm robustness fix, but if @tjtanaa prefer the recipes#615 as the supported solution, I can close this PR.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tanpinsiang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork