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 7, 17:45 UTC
[CI][AMD][BugFix] Add torch.cuda.set_device to test_punica_ops so punica kernels execute on same device as tensor (#34985)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Feb 25, 2026 at 13:18 UTC
b188bab4417f7f94df0c4e84399d163e6c0db316
1 file changed
+2
tests/lora/test_punica_ops.py
+2
@@ -395,6 +395,7 @@ def test_kernels(
395
Tests LoRA kernels.
396
"""
397
torch.set_default_device(device)
398
+ torch.cuda.set_device(device)
399
set_random_seed(seed)
400
401
if op_type == "shrink":
@@ -447,6 +448,7 @@ def test_kernels_hidden_size(
448
Tests SGMV and LoRA kernels.
449
"""
450
torch.set_default_device(device)
451
+ torch.cuda.set_device(device)
452
set_random_seed(seed)
453
454
if op_type == "shrink":