| from __future__ import annotations |
|
|
| import torch |
|
|
| from offload.dspark.native.fidelity_reference import head_logits, markov_logits, online_attention |
|
|
|
|
| def test_head_and_markov_preserve_fp32_near_tie() -> None: |
| hidden = torch.tensor([[1., 1.]], dtype=torch.bfloat16) |
| weight = torch.tensor([[1., 2 ** -9], [1., 2 ** -8]], dtype=torch.bfloat16) |
| base = head_logits(hidden, weight) |
| assert torch.equal(base, head_logits(hidden, weight, fp64=True)) |
| emb = torch.tensor([[1., 0.]], dtype=torch.bfloat16) |
| mw = torch.tensor([[0., 0.], [2 ** -10, 0.]], dtype=torch.bfloat16) |
| got = markov_logits(base, torch.tensor([0]), emb, mw) |
| assert torch.equal(got, markov_logits(base, torch.tensor([0]), emb, mw, fp64=True)) |
| assert got.argmax(-1).item() == 1 |
| assert (base.bfloat16() + head_logits(emb, mw).bfloat16()).argmax(-1).item() == 0 |
|
|
|
|
| def test_online_exact_dots_and_block_rescaling() -> None: |
| g = torch.Generator().manual_seed(37) |
| q = torch.randint(-2, 3, (1, 2, 3, 16), generator=g).bfloat16() |
| kv = torch.randint(-2, 3, (1, 133, 16), generator=g).bfloat16() |
| idx = torch.arange(133).view(1, 1, -1).expand(1, 2, -1) |
| sink = torch.tensor([0., 2., -3.]) |
| |
| a = online_attention(q, kv, sink, idx, .25) |
| b = online_attention(q, kv, sink, idx, .25, fp64=True) |
| assert torch.equal(a, b) |
|
|
|
|
| def test_sink_last_and_empty_blocks() -> None: |
| q = torch.zeros(1, 1, 1, 16, dtype=torch.bfloat16) |
| kv = torch.ones(1, 1, 16, dtype=torch.bfloat16) |
| idx = torch.full((1, 1, 129), -1) |
| assert torch.equal(online_attention(q, kv, torch.tensor([0.]), idx, 1.), q) |
| idx[..., -1] = 0 |
| assert torch.equal(online_attention(q, kv, torch.tensor([0.]), idx, 1.), q + .5) |
|
|
|
|
| def test_online_differs_from_single_pass_bf16_p() -> None: |
| g = torch.Generator().manual_seed(45) |
| q = torch.randn(1, 2, 4, 32, generator=g).bfloat16() |
| kv = torch.randn(1, 133, 32, generator=g).bfloat16() |
| idx = torch.arange(133).view(1, 1, -1).expand(1, 2, -1) |
| sink = torch.full((4,), 4.) |
| got = online_attention(q, kv, sink, idx, .2) |
| s = torch.einsum("bmhd,btd->bmht", q.float(), kv.float()) * .2 |
| mx = torch.maximum(s.amax(-1), sink) |
| p = torch.exp(s - mx[..., None]) |
| den = p.sum(-1) + torch.exp(sink - mx) |
| old = (torch.einsum("bmht,btd->bmhd", p.bfloat16().float(), kv.float()) / den[..., None]).bfloat16() |
| assert not torch.equal(got, old) |
|
|
