[Bugfix][ROCm] Only run FP8 AITER MLA prefill when using FP8 KV

[simondanielsson]

Purpose

Even when KV is not FP8, the FP8 PS ASM prefill kernel's persistent buffers are both (a) initialized/reserved and (b) filled with metadata (through _build_fp8_prefill_ps_metadata) in every forward.

Specifically the latter leads to performance degradation when using BF16 as _build_fp8_prefill_ps_metadata is a host-side function with several sync memcpys in it, causing a unnecessary H->D sync before every forward.

Co-authored with @olka-amd.

Test Plan

1. Runs as expected on gsm8k with BF16 KV

export VLLM_ROCM_USE_AITER=1
vllm bench deepseek-ai/DeepSeek-V3-0324 \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --gpu-memory-utilization 0.8 \
  --max-model-len 16K

lm_eval \
  --model local-completions \
  --model_args model=deepseek-ai/DeepSeek-V3-0324,base_url=http://localhost:8000/v1/completions,tokenized_requests=False,trust_remote_code=True,num_concurrent=256 \
  --tasks gsm8k \
  --num_fewshot 5"

Test Result

This branch:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9466	±	0.0065
		strict-match	5	exact_match	↑	0.9469	±	0.006

Nightly:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9484	±	0.0061
		strict-match	5	exact_match	↑	0.9492	±	0.0060

---

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.

[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.