[Bugfix][ROCm] Fix Unsupported attention metadata type for speculative decoding in eagle.py

[vllmellm]

Purpose

this PR fixes the issue mentioned here #30811 (comment)

Test Plan

Only testing VLLM_ATTENTION_BACKEND=ROCM_ATTN

1. run vllm serve:
   VLLM_USE_V1=1 VLLM_ROCM_USE_AITER=0 VLLM_ATTENTION_BACKEND=ROCM_ATTN vllm serve meta-llama/Llama-3.3-70B-Instruct -tp 4 --speculative-config '{"model": "yuhuili/EAGLE3-LLaMA3.3-Instruct-70B", "num_speculative_tokens": 3, "method":"eagle3", "draft_tensor_parallel_size":1}'
2. run lm_eval
   lm_eval --model local-completions --tasks gsm8k --model_args model=meta-llama/Llama-3.3-70B-Instruct,base_url=http://localhost:8000/v1/completions --trust_remote_code --num_fewshot 5 --batch_size 128

Test Result

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9333	±	0.0144
		strict-match	5	exact_match	↑	0.9033	±	0.0171

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses a ValueError that occurs during speculative decoding on ROCm when using the ROCM_ATTN backend. The fix correctly includes RocmAttentionMetadata in the list of allowed attention metadata types. My review includes a suggestion to make this fix more comprehensive by also adding support for the ROCM_AITER_UNIFIED_ATTN backend, which appears to have been overlooked and would cause a similar issue.

[gemini-code-assist]

While this change correctly adds support for RocmAttentionMetadata, it appears that another ROCm attention backend, ROCM_AITER_UNIFIED_ATTN, might have been missed. This backend likely uses RocmAiterUnifiedAttentionMetadata, which is not included in rocm_types and could lead to a similar ValueError during speculative decoding. To ensure a more comprehensive fix, I recommend adding RocmAiterUnifiedAttentionMetadata to the list of allowed types.

            from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata
            from vllm.v1.attention.backends.rocm_aiter_unified_attn import RocmAiterUnifiedAttentionMetadata

            rocm_types = [
                TritonAttentionMetadata,
                FlashAttentionMetadata,
                RocmAttentionMetadata,
                RocmAiterUnifiedAttentionMetadata,
            ]

[vllmellm]

RocmAiterUnifiedAttentionMetadata does not exist. RocmAiterUnifiedAttentionBackend shares the same metadata as RocmAttentionMetadata and the lm_eval scores are as follow:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9333	±	0.0144
		strict-match	5	exact_match	↑	0.9133	±	0.0163

run vllm serve command:

VLLM_USE_V1=1 \ VLLM_ROCM_USE_AITER=1 VLLM_ATTENTION_BACKEND=ROCM_AITER_UNIFIED_ATTN vllm serve meta-llama/Llama-3.3-70B-Instruct -tp 4 --speculative-config '{"model": "yuhuili/EAGLE3-LLaMA3.3-Instruct-70B", "num_speculative_tokens": 3, "method":"eagle3", "draft_tensor_parallel_size":1}'

[tjtanaa]

@vllmellm is FlashAttentionMetadata still relevant for ROCm platform?

[vllmellm]

removed

[tjtanaa]

LGTM. Thanks for cleaning up the conditions