| import pytest |
| import torch |
|
|
| import kernels |
|
|
| lp = kernels.get_kernel("phanerozoic/logits-processor", version=1, trust_remote_code=True) |
|
|
| requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") |
|
|
|
|
| def _pack(allowed_sets, V, device): |
| nwords = (V + 31) // 32 |
| mask = torch.zeros(len(allowed_sets), nwords, dtype=torch.int32, device=device) |
| for r, s in enumerate(allowed_sets): |
| for tok in s: |
| mask[r, tok >> 5] |= (1 << (tok & 31)) |
| return mask |
|
|
|
|
| @requires_cuda |
| @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) |
| def test_bitmask_sets_disallowed_to_neg_inf(dtype): |
| V = 130 |
| logits = torch.randn(2, V, device="cuda", dtype=dtype) |
| allowed = [{3, 7, 50, 129}, {0, 42, 63}] |
| ref = logits.clone() |
| lp.apply_token_bitmask(logits, _pack(allowed, V, "cuda")) |
| for r, s in enumerate(allowed): |
| for tok in range(V): |
| if tok in s: |
| assert logits[r, tok] == ref[r, tok] |
| else: |
| assert logits[r, tok].float() == float("-inf") |
|
|
|
|
| @requires_cuda |
| def test_greedy_and_topk1_are_argmax(): |
| logits = torch.randn(8, 2000, device="cuda") |
| amax = logits.argmax(-1) |
| assert torch.equal(lp.sample(logits, temperature=0.0), amax) |
| assert torch.equal(lp.sample(logits, top_k=1, temperature=1.0), amax) |
|
|
|
|
| @requires_cuda |
| def test_min_p_and_topp_exclude_low_prob(): |
| logits = torch.full((4, 50), -10.0, device="cuda") |
| logits[:, 7] = 10.0 |
| assert torch.equal(lp.sample(logits, temperature=1.0, min_p=0.5, seed=1), |
| torch.full((4,), 7, device="cuda")) |
| assert torch.equal(lp.sample(logits, temperature=1.0, top_p=0.9, seed=2), |
| torch.full((4,), 7, device="cuda")) |
|
|
|
|
| @requires_cuda |
| def test_deterministic(): |
| logits = torch.randn(16, 500, device="cuda") |
| a = lp.sample(logits, temperature=1.0, seed=99, offset=3) |
| b = lp.sample(logits, temperature=1.0, seed=99, offset=3) |
| assert torch.equal(a, b) |
|
|
|
|
| @requires_cuda |
| def test_distribution_matches_softmax(): |
| torch.manual_seed(0) |
| V, N = 32, 40000 |
| base = torch.randn(V, device="cuda") |
| logits = base.unsqueeze(0).repeat(N, 1) |
| toks = lp.sample(logits, temperature=1.0, seed=1234) |
| emp = torch.bincount(toks, minlength=V).float() / N |
| ref = torch.softmax(base, -1) |
| assert (emp - ref).abs().max().item() < 0.03 |
|
|
