| |
| |
| |
|
|
| 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 |
|
|
