| """Run under the N4 UUID/flock wrapper; fixtures are synthetic, never captured.""" |
| from __future__ import annotations |
|
|
| import math |
|
|
| import pytest |
| import torch |
|
|
| from offload.dspark.native.attention import attention_into |
| from offload.dspark.native.fidelity_reference import markov_logits, online_attention |
| from offload.dspark.native.ops import markov_logits_fp32 |
|
|
| pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires authorized N4 GPU hold") |
|
|
|
|
| @pytest.mark.parametrize("last", [1, 62, 63, 64, 126, 127, 128, 255, 900000]) |
| def test_online64_ring_and_graph(last: int) -> None: |
| g = torch.Generator().manual_seed(38) |
| b, k, h, d = 2, 5, 32, 512 |
| |
| q = torch.randint(-2, 3, (b, k, h, d), generator=g).bfloat16() |
| ring = torch.randint(-2, 3, (4, 3, 128, d), generator=g).bfloat16()[:, 1] |
| block = torch.randint(-2, 3, (b, k, d), generator=g).bfloat16() |
| positions = torch.full((4, 128), -1, dtype=torch.int64) |
| lo = max(0, last - 127) |
| pp = torch.arange(lo, last + 1) |
| for slot in (1, 3): |
| positions[slot, pp % 128] = pp |
| slots = torch.tensor([3, 1]) |
| lows = torch.full((b,), lo) |
| highs = torch.full((b,), last) |
| sink = torch.linspace(-2., 4., h) |
| nc = min(128, last + 1) |
| kv = torch.cat((ring[slots], block), 1) |
| idx = torch.cat((torch.arange(nc), 128 + torch.arange(k))).view(1, 1, -1).expand(b, k, -1) |
| want = online_attention(q, kv, sink, idx, d ** -.5, fp64=True) |
| args = (q.cuda(), ring.cuda(), positions.cuda(), slots.cuda(), lows.cuda(), highs.cuda(), block.cuda(), sink.cuda()) |
| out = torch.empty_like(args[0]) |
| ptr = out.data_ptr() |
| attention_into(*args, d ** -.5, out) |
| got = out.cpu() |
| |
| assert (got.float() - want.float()).norm() / want.float().norm() < .003 |
| torch.testing.assert_close(got, want, rtol=.008, atol=.004) |
| stream = torch.cuda.Stream() |
| stream.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(stream): |
| attention_into(*args, d ** -.5, out) |
| stream.synchronize() |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph, stream=stream): |
| attention_into(*args, d ** -.5, out) |
| graph.replay() |
| torch.cuda.synchronize() |
| assert out.data_ptr() == ptr |
| assert torch.equal(out.cpu(), got) |
|
|
|
|
| def test_fp32_markov_exact_dyadic_and_near_tie() -> None: |
| g = torch.Generator().manual_seed(42) |
| embed = torch.randint(-2, 3, (32, 256), generator=g).bfloat16() |
| weight = torch.randint(-2, 3, (513, 256), generator=g).bfloat16() |
| base = torch.randn(3, 513, generator=g) |
| prev = torch.tensor([0, 15, 31]) |
| want = markov_logits(base, prev, embed, weight, fp64=True) |
| got = markov_logits_fp32(base.cuda(), prev.cuda(), embed.cuda(), weight.cuda()) |
| assert torch.equal(got.cpu(), want) |
|
|
|
|
| def test_online64_exact_constant_values() -> None: |
| q = torch.zeros(1, 5, 32, 512, dtype=torch.bfloat16, device="cuda") |
| ring = torch.ones(1, 128, 512, dtype=torch.bfloat16, device="cuda") |
| positions = torch.arange(128, device="cuda").view(1, -1) |
| slot = torch.zeros(1, dtype=torch.int64, device="cuda") |
| last = torch.full_like(slot, 127) |
| block = torch.ones(1, 5, 512, dtype=torch.bfloat16, device="cuda") |
| sink = torch.zeros(32, device="cuda") |
| out = torch.empty_like(q) |
| attention_into(q, ring, positions, slot, slot, last, block, sink, 512 ** -.5, out) |
| assert torch.equal(out, torch.full_like(q, 133 / 134)) |
|
|
|
|
| def test_mixed_short_and_wrapped_windows() -> None: |
| g = torch.Generator().manual_seed(91) |
| q = torch.randint(-2, 3, (2, 5, 32, 512), generator=g).bfloat16() |
| ring = torch.randint(-2, 3, (4, 128, 512), generator=g).bfloat16() |
| block = torch.randint(-2, 3, (2, 5, 512), generator=g).bfloat16() |
| sink = torch.linspace(-2., 4., 32) |
| positions = torch.full((4, 128), -1, dtype=torch.int64) |
| slots, last = torch.tensor([3, 0]), torch.tensor([62, 900000]) |
| lo = (last - 127).clamp_min(0) |
| expected = [] |
| for b in range(2): |
| p = torch.arange(int(lo[b]), int(last[b]) + 1) |
| positions[slots[b], p % 128] = p |
| kv = torch.cat((ring[slots[b]], block[b]), 0).unsqueeze(0) |
| idx = torch.cat((torch.arange(min(128, int(last[b]) + 1)), 128 + torch.arange(5))) |
| expected.append(online_attention(q[b:b+1], kv, sink, idx.view(1, 1, -1).expand(1, 5, -1), |
| 512 ** -.5, fp64=True)) |
| args = (q.cuda(), ring.cuda(), positions.cuda(), slots.cuda(), lo.cuda(), last.cuda(), block.cuda(), sink.cuda()) |
| out = torch.empty_like(args[0]) |
| attention_into(*args, 512 ** -.5, out) |
| want = torch.cat(expected) |
| got = out.cpu() |
| assert (got.float() - want.float()).norm() / want.float().norm() < .003 |
| torch.testing.assert_close(got, want, rtol=.008, atol=.004) |
|
|
|
|
| def test_online_blocks_and_sink_rounding_are_required() -> None: |
| |
| |
| |
| |
| q = torch.zeros(1, 5, 32, 512, dtype=torch.bfloat16, device="cuda") |
| q[..., 0] = 1 |
| ring = torch.zeros(1, 128, 512, dtype=torch.bfloat16, device="cuda") |
| ring[:, :64, 1] = 1 |
| positions = torch.full((1, 128), -1, dtype=torch.int64, device="cuda") |
| positions[:, :64] = torch.arange(64, device="cuda") |
| slot = torch.zeros(1, dtype=torch.int64, device="cuda") |
| last = torch.full_like(slot, 63) |
| block = torch.zeros(1, 5, 512, dtype=torch.bfloat16, device="cuda") |
| block[..., 0] = 1 |
| sink = torch.full((32,), 2 * math.log(3), device="cuda") |
| out = torch.empty_like(q) |
| attention_into(q, ring, positions, slot, slot, last, block, sink, math.log(3), out) |
| assert torch.equal(out[..., 1], torch.full_like(out[..., 1], 8 / 11)) |
|
|
