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 10, 00:15 UTC
[CI][Hardware][AMD] Fix test_rotary_embedding_mla_cache_fused (#32408)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Jan 19, 2026 at 02:25 UTC
11bbf86f6a35d8d30821bed72eb8adf9864b024f
1 file changed
+12
-3
tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
+12
-3
@@ -13,6 +13,7 @@ from tests.kernels.allclose_default import get_default_atol, get_default_rtol
13
from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
14
from vllm import _custom_ops as ops
15
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
16
+from vllm.platforms import current_platform
17
from vllm.utils.torch_utils import set_random_seed
18
19
@@ -68,9 +69,17 @@ def test_concat_and_cache_mla_rope_fused(
69
k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device)
70
kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device)
71
71
- # NOTE(woosuk): The reference implementation should be executed first
72
- # because the custom kernel is in-place.
73
- ref_q_pe, ref_k_pe = rope.forward_native(positions, query, k_pe)
72
+ if current_platform.is_rocm():
73
+ # We use forward_hip for the same numerics as the fused custom kernel on ROCm
74
+ # when dtype is FP16. The torch-native implementation implicitly upcasts
75
+ # FP16 x FP16 multiplications to FP32 before downcasting them, which leads
76
+ # to notable output divergences.
77
+ # Clone the tensors because the implementation modifies them in-place
78
+ ref_q_pe, ref_k_pe = rope.forward_hip(positions, query.clone(), k_pe.clone())
79
+ else:
80
+ # NOTE(woosuk): The reference implementation should be executed first
81
+ # because the custom kernel is in-place.
82
+ ref_q_pe, ref_k_pe = rope.forward_native(positions, query, k_pe)
83
assert ref_k_pe is not None
84
85
ref_k_pe = torch.flatten(ref_k_pe, start_dim=1).to(device=device)