| import pytest |
| import torch |
|
|
|
|
| class TestBitonicTopk: |
|
|
| @pytest.fixture |
| def device(self): |
| yield 'cuda' |
|
|
| @pytest.fixture |
| def k(self): |
| yield 2048 |
|
|
| @pytest.fixture |
| def q_seqlens(self, device): |
| ret = [4, 16, 1, 32] |
| ret = torch.tensor(ret, dtype=torch.int32, device=device) |
| yield ret |
|
|
| @pytest.fixture |
| def kv_seqlens(self, device): |
| ret = [1024, 2048, 4096, 4096 + 133] |
| ret = torch.tensor(ret, dtype=torch.int32, device=device) |
| yield ret |
|
|
| @pytest.fixture |
| def batch_size(self, kv_seqlens): |
| return kv_seqlens.numel() |
|
|
| @pytest.fixture |
| def max_kv_len(self, kv_seqlens): |
| return kv_seqlens.max().item() |
|
|
| @pytest.fixture |
| def scores(self, q_seqlens, max_kv_len, device): |
| num_tokens = q_seqlens.sum().item() |
| yield torch.randn((num_tokens, max_kv_len), device=device) |
|
|
| @pytest.fixture |
| def gt(self, scores, q_seqlens, kv_seqlens, k): |
| batch_size = kv_seqlens.numel() |
| num_tokens, _ = scores.shape |
| topk_indices = torch.empty((num_tokens, k), dtype=torch.int32, device=scores.device) |
| topk_indices.fill_(-1) |
|
|
| start = 0 |
| for i in range(batch_size): |
| q_seqlen = q_seqlens[i].item() |
| seqlen = kv_seqlens[i].item() |
| tmp_k = min(seqlen, k) |
| end = start + q_seqlen |
| _, topk_indices[start:end, :seqlen] = torch.topk(scores[start:end, :seqlen], |
| tmp_k, |
| largest=True, |
| sorted=True) |
| start = end |
| return topk_indices |
|
|
| def test_bitonic_topk(self, scores, q_seqlens, kv_seqlens, k, gt): |
| from lmdeploy.pytorch.kernels.cuda.bitonic_topk import bitonic_topk |
| out = bitonic_topk(scores, q_seqlens=q_seqlens, kv_seqlens=kv_seqlens, k=k, fill=-1, sorted=True) |
| gt[gt < 0] = 0 |
| out[out < 0] = 0 |
| gt_score = torch.gather(scores, 1, gt.to(torch.int64)) |
| out_score = torch.gather(scores, 1, out.to(torch.int64)) |
| torch.testing.assert_close(gt_score, out_score) |
|
|
