| |
| |
|
|
| import numpy as np |
| import pytest |
| import torch |
|
|
| from megablocks import ops |
|
|
| PADDED_GATHER_TESTS = ( |
| (4, 2, 2, 1), |
| (4, 2, 2, 2), |
| (1024, 1, 4, 1), |
| (1024, 1, 4, 2), |
| (1024, 1, 4, 4), |
| (1024, 1, 64, 1), |
| (1024, 1, 64, 2), |
| (1024, 1, 64, 4), |
| (1024, 1, 128, 1), |
| (1024, 1, 128, 2), |
| (1024, 1, 128, 4), |
| (1024, 1536, 4, 1), |
| (1024, 1536, 4, 2), |
| (1024, 1536, 4, 4), |
| (1024, 1536, 64, 1), |
| (1024, 1536, 64, 2), |
| (1024, 1536, 64, 4), |
| (1024, 1536, 128, 1), |
| (1024, 1536, 128, 2), |
| (1024, 1536, 128, 4), |
| (16384, 768, 4, 1), |
| (16384, 768, 4, 2), |
| (16384, 768, 4, 4), |
| (16384, 768, 64, 1), |
| (16384, 768, 64, 2), |
| (16384, 768, 64, 4), |
| (16384, 768, 128, 1), |
| (16384, 768, 128, 2), |
| (16384, 768, 128, 4), |
| (16384, 1, 4, 1), |
| (16384, 1, 4, 2), |
| (16384, 1, 4, 4), |
| (16384, 1, 64, 1), |
| (16384, 1, 64, 2), |
| (16384, 1, 64, 4), |
| (16384, 1, 128, 1), |
| (16384, 1, 128, 2), |
| (16384, 1, 128, 4), |
| ) |
|
|
|
|
| @pytest.mark.gpu |
| @pytest.mark.parametrize(('sl', 'hs', 'ne', 'top_k'), PADDED_GATHER_TESTS) |
| def testPaddedGather(sl: int, hs: int, ne: int, top_k: int): |
| |
| x = torch.randn((sl, hs)).cuda().half() |
|
|
| |
| top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int() |
| bin_ids, indices = ops.sort(top_expert) |
| tokens_per_expert = ops.histogram(top_expert, ne) |
| padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128) |
| padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0) |
| bins = ops.inclusive_cumsum(tokens_per_expert, 0) |
|
|
| def padded_gather( |
| x: torch.Tensor, |
| indices: torch.Tensor, |
| bin_ids: torch.Tensor, |
| bins: torch.Tensor, |
| padded_bins: torch.Tensor, |
| top_k: int, |
| ): |
| x = x.cpu().numpy() |
| indices = indices.cpu().numpy() |
| bin_ids = bin_ids.cpu().numpy() |
| bins = bins.cpu().numpy() |
| padded_bins = padded_bins.cpu().numpy() |
|
|
| out = np.zeros((padded_bins[-1], hs)) |
| in_idx = 0 |
| for i, end in enumerate(bins): |
| out_idx = 0 if i == 0 else padded_bins[i - 1] |
| end = bins[i] |
| while in_idx < end: |
| load_idx = indices[in_idx] // top_k |
| out[out_idx, :] = x[load_idx, :] |
| in_idx += 1 |
| out_idx += 1 |
| return torch.from_numpy(out).cuda().half() |
|
|
| out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) |
| expected_out = padded_gather(x, indices, bin_ids, bins, padded_bins, top_k) |
| assert torch.all(torch.eq(out, expected_out)) |
|
|
