File size: 653 Bytes
c92e699 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
#
# SPDX-License-Identifier: Apache-2.0
import pytest
from models.tt_transformers.tt.rope import get_batch_size_per_device_group
@pytest.mark.parametrize(
("batch_size", "use_qk_fused", "num_devices", "mesh_shape", "expected"),
[
(1, True, 32, (8, 4), 2),
(1, False, 32, (8, 4), 1),
(32, True, 32, (8, 4), 16),
(1, True, 1, (), 2),
],
)
def test_get_batch_size_per_device_group(batch_size, use_qk_fused, num_devices, mesh_shape, expected):
assert get_batch_size_per_device_group(batch_size, use_qk_fused, num_devices, mesh_shape, 1) == expected
|