42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
|
|
from vllm_ascend.worker.utils import copy_snapshot_to_gpu
|
|
|
|
|
|
class TestQueryStartLocCopy(unittest.TestCase):
|
|
def test_copy_uses_stable_cpu_snapshot(self):
|
|
class DeferredCopy:
|
|
def copy_(self, source, non_blocking=False):
|
|
self.source = source
|
|
self.non_blocking = non_blocking
|
|
return self
|
|
|
|
cpu = torch.tensor([0, 2, 5], dtype=torch.int32)
|
|
gpu = DeferredCopy()
|
|
query_start_loc = SimpleNamespace(cpu=cpu, gpu=gpu)
|
|
|
|
with patch.object(torch.Tensor, "pin_memory", lambda tensor: tensor):
|
|
copy_snapshot_to_gpu(query_start_loc)
|
|
cpu.fill_(99)
|
|
|
|
self.assertEqual(gpu.source.tolist(), [0, 2, 5])
|
|
self.assertNotEqual(gpu.source.data_ptr(), cpu.data_ptr())
|
|
self.assertTrue(gpu.non_blocking)
|
|
|
|
def test_copy_pins_snapshot(self):
|
|
cpu = MagicMock()
|
|
snapshot = MagicMock()
|
|
pinned_snapshot = MagicMock()
|
|
cpu.clone.return_value = snapshot
|
|
snapshot.pin_memory.return_value = pinned_snapshot
|
|
gpu = MagicMock()
|
|
|
|
copy_snapshot_to_gpu(SimpleNamespace(cpu=cpu, gpu=gpu))
|
|
|
|
snapshot.pin_memory.assert_called_once_with()
|
|
gpu.copy_.assert_called_once_with(pinned_snapshot, non_blocking=True)
|