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, 09:45 UTC
[Bugfix] fix rotary embedding test for _get_padded_tensor_shape (#18229)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Lucas Wilkinson committed
May 15, 2025 at 21:32 UTC
4e1c6a02641e427a6140d33262f1467906817781
1 file changed
+4
tests/kernels/core/test_pos_encoding.py
+4
@@ -152,6 +152,10 @@ def test_batched_rotary_embedding(
152
query = torch.randn(query_shape, dtype=dtype)
153
key = torch.randn_like(query) if use_key else None
154
155
+ # slice tensor if required, noop otherwise
156
+ query = query[..., :head_size]
157
+ key = key[..., :head_size] if use_key else None
158
+
159
# NOTE(woosuk): The reference implementation should be executed first
160
# because the custom kernel is in-place.
161
ref_query, ref_key = rope.forward_native(positions, query, key)