[BugFix] [ROCm]: Bugfix and handle addition case of input for rocm_aiter_rms_norm

[tjtanaa]

After this PR (#17735) , a new case to be handled in rocm_aiter_rms_norm.

x could be a 3-dimensional tensor. x is reshaped to 2-dimensional tensors before calling aiter.rms_norm

lm_eval

Dataset: GSM8K
Engine: V0

Qwen/Qwen3-235B-A22B-FP8 without AITER

vllm (pretrained=Qwen/Qwen3-235B-A22B-FP8,tensor_parallel_size=4,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8726	±	0.0092
		strict-match	5	exact_match	↑	0.8567	±	0.0097

Qwen/Qwen3-235B-A22B-FP8 with AITER RMS Norm

vllm (pretrained=Qwen/Qwen3-235B-A22B-FP8,tensor_parallel_size=4,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8734	±	0.0092
		strict-match	5	exact_match	↑	0.8506	±	0.0098

Qwen/Qwen3-32B without AITER

vllm (pretrained=Qwen/Qwen3-32B,tensor_parallel_size=1,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.6255	±	0.0133
		strict-match	5	exact_match	↑	0.7369	±	0.0121

Qwen/Qwen3-32B with AITER RMS Norm

vllm (pretrained=Qwen/Qwen3-32B,tensor_parallel_size=4,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.6209	±	0.0134
		strict-match	5	exact_match	↑	0.7453	±	0.0120

Add Qwen3 dense to the list of test_common.py

Passed: tests/models/language/generation/test_common.py

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[DarkLight1337]

Why wasn't this copy necessary before? Would making the copy conditional improve performance in cases where a copy isn't required?

[tjtanaa]

@DarkLight1337 It is a bugfix of using the AITER commit: 5a77249 . The usage of the ops has changed at some point. However, it hasn't been throwing error in any of the unit tests. To ensure the correctness, we have followed the usage of the kernel based on this commit 5a77249 specifically.