| """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) |
|
|
