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, 23:30 UTC
[BugFix] Enhance test_pos_encoding to support execution on multi-devices (#13187)
Signed-off-by: wchen61 <wchen61@foxmail.com>
wchen61 committed
Feb 16, 2025 at 16:59 UTC
dc0f7ccf8bc1b2b74113fab8ea69d8420a91274a
1 file changed
+3
-3
tests/kernels/test_pos_encoding.py
+3
-3
@@ -70,7 +70,7 @@ def test_rotary_embedding(
70
if rotary_dim is None:
71
rotary_dim = head_size
72
rope = get_rope(head_size, rotary_dim, max_position, base, is_neox_style)
73
- rope = rope.to(dtype=dtype)
73
+ rope = rope.to(dtype=dtype, device=torch.get_default_device())
74
75
positions = torch.randint(0, max_position, (batch_size, seq_len))
76
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
@@ -125,7 +125,7 @@ def test_batched_rotary_embedding(
125
"rope_type": "linear",
126
"factor": (1, )
127
})
128
- rope = rope.to(dtype=dtype)
128
+ rope = rope.to(dtype=dtype, device=torch.get_default_device())
129
130
positions = torch.randint(0, max_position, (batch_size, seq_len))
131
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
@@ -184,7 +184,7 @@ def test_batched_rotary_embedding_multi_lora(
184
"rope_type": "linear",
185
"factor": tuple(scaling_factors)
186
})
187
- rope = rope.to(dtype=dtype)
187
+ rope = rope.to(dtype=dtype, device=torch.get_default_device())
188
189
positions = torch.randint(0, max_position, (batch_size, seq_len))
190
query = torch.randn(batch_size,