File size: 2,401 Bytes
6aac1a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | import pytest
import torch
import kernels
sdo = kernels.get_kernel("phanerozoic/spec-decode-ops", version=1, trust_remote_code=True)
def hf_filter_reference(logits, temperature, top_k, top_p, min_p):
"""Reference filter chain with transformers semantics."""
l = logits.float() / temperature
if top_k > 0:
kth = torch.topk(l, top_k, dim=-1).values[..., -1, None]
l = l.masked_fill(l < kth, float("-inf"))
if top_p < 1.0:
sl, si = torch.sort(l, descending=True, dim=-1)
p = torch.softmax(sl, dim=-1)
cum = p.cumsum(dim=-1)
keep_sorted = (cum - p) < top_p # keep tokens whose exclusive cumsum < top_p
keep_sorted[..., 0] = True
keep = torch.zeros_like(keep_sorted).scatter(-1, si, keep_sorted)
l = l.masked_fill(~keep, float("-inf"))
if min_p > 0.0:
pmax = l.max(dim=-1, keepdim=True).values
l = l.masked_fill(l < pmax + torch.log(torch.tensor(min_p)), float("-inf"))
return l
@pytest.mark.kernels_ci
@pytest.mark.parametrize("top_k,top_p,min_p", [(50, 1.0, 0.0), (0, 0.9, 0.0), (0, 1.0, 0.05), (40, 0.92, 0.02)])
def test_kept_set_matches_reference(top_k, top_p, min_p):
torch.manual_seed(0)
logits = torch.randn(8, 4096, device="cuda") * 3.0
ours = sdo.filter_logits(logits, temperature=0.7, top_k=top_k, top_p=top_p, min_p=min_p)
ref = hf_filter_reference(logits, 0.7, top_k, top_p, min_p)
assert torch.equal(torch.isinf(ours), torch.isinf(ref))
@pytest.mark.kernels_ci
def test_greedy_matches_argmax():
torch.manual_seed(1)
logits = torch.randn(16, 32000, device="cuda", dtype=torch.bfloat16)
tok = sdo.sample(logits, temperature=0.0)
assert torch.equal(tok, logits.float().argmax(dim=-1))
@pytest.mark.kernels_ci
def test_verify_greedy_contract():
torch.manual_seed(2)
B, k, V = 4, 5, 1024
target = torch.randn(B, k + 1, V, device="cuda")
# draft tokens equal to target argmax for the first j positions
am = target.argmax(dim=-1) # [B, k+1]
draft_tokens = am[:, :k].clone()
draft_tokens[0, 2] = (draft_tokens[0, 2] + 1) % V # force mismatch at position 2
draft_logits = torch.randn(B, k, V, device="cuda")
alen, out = sdo.verify(target, draft_logits, draft_tokens, temperature=0.0)
assert alen[0].item() == 2
assert (alen[1:] == k).all()
assert torch.equal(out[1, : k + 1], am[1])
|