A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
[BugFix] [ROCm]: Bugfix and handle addition case of input for `rocm_aiter_rms_norm` (#17857)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
May 11, 2025 at 19:17 UTC
a810b5b088b898bdfe25606589af69da30e85ae6
2 files changed
+15
-4
tests/models/language/generation/test_common.py
+4
@@ -28,6 +28,7 @@ AITER_MODEL_LIST = [
28
"Qwen/Qwen-7B-Chat",
29
"Qwen/Qwen2.5-0.5B-Instruct",
30
"TitanML/tiny-mixtral",
31
+ "Qwen/Qwen3-8B",
32
]
33
34
@@ -78,6 +79,9 @@ AITER_MODEL_LIST = [
79
"Qwen/Qwen2.5-0.5B-Instruct", # qwen2
80
marks=[pytest.mark.core_model],
81
),
82
+ pytest.param(
83
+ "Qwen/Qwen3-8B", # qwen (text-only)
84
+ ),
85
pytest.param("stabilityai/stablelm-3b-4e1t"), # stablelm
86
pytest.param("bigcode/starcoder2-3b"), # starcoder2
87
pytest.param(
vllm/model_executor/layers/layernorm.py
+11
-4
@@ -46,6 +46,12 @@ def rocm_aiter_rms_norm(x: torch.Tensor, weight: torch.Tensor,
46
variance_epsilon: float) -> torch.Tensor:
47
48
import aiter as rocm_aiter
49
+ if x.dim() > 2:
50
+ x_original_shape = x.shape
51
+ x = x.reshape(-1, x_original_shape[-1])
52
+ x = rocm_aiter.rms_norm(x, weight, variance_epsilon)
53
+ return x.reshape(x_original_shape)
54
+
55
return rocm_aiter.rms_norm(x, weight, variance_epsilon)
56
57
@@ -55,16 +61,17 @@ def rocm_aiter_fused_add_rms_norm(
61
62
import aiter as rocm_aiter
63
58
- # Assuming the correct signature for rmsnorm2d_fwd_with_add
64
+ residual_out = torch.empty_like(residual)
65
+ output = torch.empty_like(x)
66
rocm_aiter.rmsnorm2d_fwd_with_add(
60
- x, # output
67
+ output, # output
68
x, # input
69
residual, # residual input
63
- residual, # residual output
70
+ residual_out, # residual output
71
weight,
72
variance_epsilon,
73
)
67
- return x, residual
74
+ return output, residual_out
75
76
77
def dispatch_cuda_rmsnorm_func(add_residual: bool):