File size: 6,470 Bytes
e1acc15 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """CPU checks of the actual K2 branch boundary; GPU numerics have a solo probe."""
from __future__ import annotations
from contextlib import nullcontext
from types import SimpleNamespace
from typing import Any, cast
import pytest
import torch
from offload.dspark.worker import fused as fused_ops
from offload.dspark.worker import sampled_fused, tp2
def test_sampled_replaces_original_selector(monkeypatch: pytest.MonkeyPatch) -> None:
calls: list[int] = []
def forbidden(*args: Any, **kwargs: Any) -> None:
raise AssertionError("duplicate original selector")
def select(ranks: Any, logits: Any, prev: Any, outputs: Any, position: int) -> None:
calls.append(position)
for out in outputs.values():
out[:, position] = 9 + position
def gemm(rows: torch.Tensor, out: torch.Tensor) -> None:
out.zero_()
rank = SimpleNamespace(r=0, dev=torch.device("cpu"), v0=0, v1=64,
head=SimpleNamespace(gemm=gemm), norm=torch.ones(2),
markov=SimpleNamespace(embed=lambda ids: torch.zeros(*ids.shape, 1)),
conf_w=torch.zeros(3))
model = SimpleNamespace(s=SimpleNamespace(hidden=2), eps=1e-6, HEAD_CHUNK=16,
ranks=[rank], comm=SimpleNamespace(real=False),
_pp4_context=SimpleNamespace(select=select))
model._each = lambda ranks, fn: {r: fn(rank) for r in ranks}
monkeypatch.setattr(tp2, "markov_local_argmax", forbidden)
monkeypatch.setattr(tp2, "markov_partials", forbidden)
monkeypatch.setattr(fused_ops, "rms_norm", lambda hh, norm, eps: hh)
monkeypatch.setattr(torch.cuda, "device", lambda dev: nullcontext())
result, conf, margins = tp2.TPDrafter.heads(cast(tp2.TPDrafter, model), (0,), {0: torch.ones(5, 2)},
{0: (None, None, None, torch.zeros(1, dtype=torch.int64))}, 1, 5)
assert calls == list(range(5))
assert result[0].tolist() == [[9, 10, 11, 12, 13]]
assert torch.equal(conf, torch.full((1, 5), .5))
assert margins is None
def test_all_t0_reuses_original_graph_objects(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DSPARK_SAMPLED_TOP32", "1")
ids = torch.tensor([[7, 8, 9, 10, 0]])
conf = torch.tensor([[.1, .2, .3, .4, 0.]])
called: list[int] = []
def graph(inp: torch.Tensor, k: int, n_out: int) -> tuple[torch.Tensor, torch.Tensor, None]:
called.append(n_out)
return ids, conf, None
model = SimpleNamespace(dev=torch.device("cpu"), has_graph=lambda *args: True, draft_graphed=graph)
actual, actual_conf, support, q = tp2.TPDrafter.draft_sampled(
cast(tp2.TPDrafter, model), torch.zeros(4, 1, dtype=torch.int64), 5, torch.zeros(1, 3, dtype=torch.int64),
torch.zeros(1, dtype=torch.bool), n_out=4)
assert actual is ids and actual_conf is conf
assert called == [4]
assert support.shape == q.shape == (1, 5, 32)
assert torch.count_nonzero(q).item() == 0
def test_missing_sampled_graph_refuses_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DSPARK_SAMPLED_TOP32", "1")
model = SimpleNamespace(dev=torch.device("cpu"), graphs={(1, 5, 5): {}})
with pytest.raises(RuntimeError, match="sampled graph missing"):
tp2.TPDrafter.draft_sampled(cast(tp2.TPDrafter, model), torch.zeros(4, 1, dtype=torch.int64), 5,
torch.zeros(1, 3, dtype=torch.int64), torch.ones(1, dtype=torch.bool))
def test_off_refuses_sampled_entry(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("DSPARK_SAMPLED_TOP32", raising=False)
with pytest.raises(RuntimeError, match="explicit switch"):
tp2.TPDrafter.draft_sampled(cast(tp2.TPDrafter, SimpleNamespace()), torch.zeros(4, 1, dtype=torch.int64), 5,
torch.zeros(1, 3, dtype=torch.int64), torch.ones(1, dtype=torch.bool))
def paired_fixture() -> Any:
ranks = []
for r in (0, 1):
markov = SimpleNamespace(slot_of_token=torch.zeros(128, dtype=torch.int64),
cache=torch.zeros(1, 64, dtype=torch.bfloat16),
w1=torch.zeros(128, 256, dtype=torch.bfloat16),
w2=torch.zeros(64, 256, dtype=torch.bfloat16))
ranks.append(SimpleNamespace(v0=r * 64, v1=(r + 1) * 64, markov=markov, dev=torch.device("cpu")))
return SimpleNamespace(ALL=(0, 1), devs=[torch.device("cpu")] * 2, ranks=ranks, emulated=False)
def test_pair_primes_both_before_any_launch(monkeypatch: pytest.MonkeyPatch) -> None:
events: list[str] = []
class Compiler:
def warmup(self, *args: Any, **kwargs: Any) -> SimpleNamespace:
events.append("compile")
return SimpleNamespace(run=lambda *a, **kw: None)
def __getitem__(self, grid: Any) -> Any:
def launch(*args: Any, **kwargs: Any) -> None:
events.append("launch")
return launch
monkeypatch.setattr(sampled_fused, "_kernel", Compiler())
monkeypatch.setattr(torch.cuda, "device", lambda dev: nullcontext())
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
model = paired_fixture()
ctx = sampled_fused.FusedContext(model, 1, 5)
logits = {r: torch.zeros(1, 5, 64, dtype=torch.bfloat16) for r in (0, 1)}
previous = {r: torch.zeros(1, dtype=torch.int64) for r in (0, 1)}
output = {r: torch.zeros(1, 5, dtype=torch.int64) for r in (0, 1)}
ctx.select((0, 1), logits, previous, output, 0)
assert events == ["compile", "compile", "launch", "launch"]
events.clear()
ctx.select((0, 1), logits, previous, output, 0)
assert events == ["launch", "launch"]
assert not torch.count_nonzero(ctx.state[0]["ids"][:, 4]).item()
assert not torch.count_nonzero(ctx.state[0]["q"][:, 4]).item()
def test_unprimed_pair_capture_refuses(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(torch.cuda, "device", lambda dev: nullcontext())
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
ctx = sampled_fused.FusedContext(paired_fixture(), 1, 5)
with pytest.raises(RuntimeError, match="primed together"):
ctx.select((0,), {0: torch.zeros(1, 5, 64, dtype=torch.bfloat16)},
{0: torch.zeros(1, dtype=torch.int64)}, {0: torch.zeros(1, 5, dtype=torch.int64)}, 0)
|