Inesh Reddy ChappidiandGao Han
66d938100b
Signed-off-by: Inesh Reddy <ineshreddy249@gmail.com> Signed-off-by: IneshReddy249 <ineshreddy249@gmail.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Unit tests for the shared runner helpers in vllm_omni.worker.mixins."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.worker.mixins import maybe_unpad_input_ids
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"flag, padded_len, unpadded, expected",
|
|
[
|
|
(True, 4, 3, 3),
|
|
(False, 4, 3, 4),
|
|
(True, 3, 3, 3),
|
|
],
|
|
)
|
|
def test_maybe_unpad_input_ids(flag, padded_len, unpadded, expected):
|
|
"""#6712: opted-in models get input_ids at the real token count; everything
|
|
else keeps the graph-bucket padded buffer."""
|
|
model = SimpleNamespace(requires_exact_input_shape=flag)
|
|
ids = torch.zeros(padded_len, dtype=torch.int32)
|
|
assert maybe_unpad_input_ids(model, ids, unpadded).numel() == expected
|
|
|
|
|
|
def test_maybe_unpad_input_ids_without_capability_attribute():
|
|
ids = torch.zeros(4, dtype=torch.int32)
|
|
assert maybe_unpad_input_ids(SimpleNamespace(), ids, 3).numel() == 4
|
|
|
|
|
|
def test_maybe_unpad_input_ids_handles_none():
|
|
model = SimpleNamespace(requires_exact_input_shape=True)
|
|
assert maybe_unpad_input_ids(model, None, 3) is None
|