29 lines
1.0 KiB
Python
29 lines
1.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""CUDA-lane checks for GPU-resident intermediate buffer values."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
|
|
|
|
|
|
@pytest.mark.parametrize("key,dtype", [(("embed", "prefill"), torch.float32), (("codes", "ref"), torch.long)])
|
|
def test_gpu_resident_values_stay_on_device_after_update(key, dtype):
|
|
# Declared gpu-resident keys must keep tensors on device (CLA snapshot real).
|
|
buffer = OmniIntermediateBuffer(1)
|
|
value = torch.ones(2, 3, device="cuda", dtype=dtype)
|
|
buffer.add_request(0, SimpleNamespace(req_id="r1", mm_features=[]))
|
|
|
|
buffer.update(0, {key[0]: {key[1]: value}}, {key})
|
|
|
|
stored = buffer.buffers[0][key[0]][key[1]]
|
|
assert stored.device.type == "cuda"
|
|
assert stored.dtype == dtype
|
|
assert torch.equal(stored.cpu(), value.cpu())
|