184 lines
8.7 KiB
Python
184 lines
8.7 KiB
Python
"""CPU interface/provenance checks; GPU tests are explicitly marked by skips."""
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from safetensors.torch import save_file, load_file
|
|
|
|
from yue2 import fast
|
|
from yue2 import quantization as quant
|
|
from yue2.protocol import Sampling
|
|
from yue2.modeling_yue2 import YuE2Config, YuE2ForCausalLM
|
|
|
|
|
|
def make_checkpoint(path):
|
|
path.mkdir()
|
|
config = YuE2Config(hidden_size=32, intermediate_size=64, num_hidden_layers=1,
|
|
num_attention_heads=4, num_key_value_heads=2, head_dim=8,
|
|
vocab_size=71, max_position_embeddings=64, max_latent_frames=64).to_dict()
|
|
(path / "config.json").write_text(json.dumps(config))
|
|
weights = {name: torch.randn(2, 2, dtype=torch.bfloat16) for name in fast.ar_keys(config)}
|
|
weights["model.layers.0.nar_mlp.up_proj.weight"] = torch.ones(3, 2, dtype=torch.bfloat16)
|
|
save_file(weights, path / "model.safetensors")
|
|
return weights, config
|
|
|
|
|
|
def test_optional_imports_do_not_load_cuda_extensions():
|
|
code = """
|
|
import sys, importlib.abc
|
|
class Deny(importlib.abc.MetaPathFinder):
|
|
def find_spec(self, fullname, path=None, target=None):
|
|
if fullname.split('.')[0] in {'vllm', 'triton'}:
|
|
raise RuntimeError('Unexpected optional import: ' + fullname)
|
|
sys.meta_path.insert(0, Deny())
|
|
import yue2.fast
|
|
import yue2.quantization
|
|
"""
|
|
result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True)
|
|
assert result.returncode == 0, result.stderr
|
|
|
|
|
|
def test_derived_ar_subset_and_hash_validation(tmp_path):
|
|
original, config = make_checkpoint(tmp_path / "full")
|
|
derived = fast.derive_ar_checkpoint(tmp_path / "full", tmp_path / "cache")
|
|
actual = load_file(derived / "model.safetensors")
|
|
assert set(actual) == fast.ar_keys(config)
|
|
assert all(torch.equal(actual[key], original[key]) for key in actual)
|
|
assert json.loads((derived / "config.json").read_text())["architectures"] == ["Qwen3ForCausalLM"]
|
|
assert fast.derive_ar_checkpoint(tmp_path / "full", tmp_path / "cache") == derived
|
|
(derived / "config.json").write_text('{}')
|
|
with pytest.raises(ValueError, match="Corrupt"):
|
|
fast.derive_ar_checkpoint(tmp_path / "full", tmp_path / "cache")
|
|
|
|
|
|
def test_missing_ar_tensor_fails(tmp_path):
|
|
weights, _ = make_checkpoint(tmp_path / "full")
|
|
del weights['lm_head.weight']
|
|
save_file(weights, tmp_path / "full/model.safetensors")
|
|
with pytest.raises(ValueError, match="Missing AR"):
|
|
fast.derive_ar_checkpoint(tmp_path / "full", tmp_path / "cache")
|
|
|
|
|
|
def test_full_context_kv_budget():
|
|
config = YuE2Config().to_dict()
|
|
assert fast.kv_cache_bytes(config) == int(2.625 * 2**30)
|
|
assert fast.kv_cache_bytes(config, max_sequences=2) == int(5.25 * 2**30)
|
|
|
|
|
|
@pytest.mark.parametrize("cfg,legacy,reason", [(1.01, True, "paired_cfg_uses_torch"),
|
|
(1., True, "historical_off_sampling_uses_torch"),
|
|
(1., False, "vllm_requires_cuda")])
|
|
def test_cfg_and_device_fallback_preserve_protocol(monkeypatch, cfg, legacy, reason):
|
|
from yue2 import sampling
|
|
calls, model = [], object()
|
|
pipe = SimpleNamespace(device=torch.device("cpu"), quantization="none", _load_model=lambda: model)
|
|
prefix, negative = [1, 2], [3]
|
|
sample = Sampling(max_tokens=4, min_tokens=0)
|
|
def generate(*args, **kwargs):
|
|
calls.append((args, kwargs))
|
|
return [9], {"output_tokens": 1}, True
|
|
monkeypatch.setattr(sampling, "generate_tokens", generate)
|
|
result = fast.generate_vllm(pipe, prefix, sample, 71, "semantic", negative=negative,
|
|
cfg_scale=cfg, legacy_off=legacy)
|
|
assert calls[0][0] == (model, prefix, sample, 71, "semantic")
|
|
assert calls[0][1]['negative'] is negative
|
|
assert calls[0][1]['cfg_scale'] == cfg
|
|
assert calls[0][1]['legacy_off'] == legacy
|
|
assert result[1]['backend_actual'] == 'torch' and result[1]['fallback_reason'] == reason
|
|
|
|
|
|
def test_worker_transport_handles_many_buffered_lines(monkeypatch):
|
|
original_popen = subprocess.Popen
|
|
worker_code = """
|
|
import sys,json
|
|
json.loads(sys.stdin.readline())
|
|
for line in sys.stdin:
|
|
q=json.loads(line)
|
|
lines=['YUE2_FAST\\t'+json.dumps({'event':'token','token':i}) for i in range(100)]
|
|
lines.append('YUE2_FAST\\t'+json.dumps({'event':'result','ids':[12], 'timing':{},'truncated':False}))
|
|
sys.stdout.write('\\n'.join(lines)+'\\n'); sys.stdout.flush()
|
|
"""
|
|
monkeypatch.setattr(fast.subprocess, "Popen", lambda command, **kwargs: original_popen([sys.executable, '-c', worker_code], **kwargs))
|
|
pipe = SimpleNamespace(model_dir=Path('.'), device=torch.device('cpu'), memory_budget_gib=24)
|
|
worker = fast._Worker(pipe)
|
|
tokens = []
|
|
try:
|
|
for _ in range(2):
|
|
assert worker.request({'phase': 'abc'}, on_token=lambda phase, token: tokens.append((phase, token)))[0] == [12]
|
|
assert len(tokens) == 200
|
|
finally:
|
|
worker.close()
|
|
assert worker.process.poll() is not None
|
|
|
|
|
|
def test_fp8_tensor_scaling_and_cpu_execution_guard():
|
|
value = torch.tensor([-448., -2., 0., 1., 448.])
|
|
encoded, scale = quant.quantize_tensor(value)
|
|
torch.testing.assert_close(encoded.float() * scale, value)
|
|
layer = quant.FP8Linear(torch.nn.Linear(32, 16, bias=False), "cpu")
|
|
assert layer.weight.dtype == torch.float8_e4m3fn
|
|
with pytest.raises(RuntimeError, match="requires CUDA"):
|
|
layer(torch.randn(1, 32))
|
|
|
|
|
|
def test_fp8_cpu_rejected():
|
|
with pytest.raises(RuntimeError, match="requires CUDA"):
|
|
quant.prepare_fp8_ar(torch.nn.Linear(32, 32), "cpu")
|
|
|
|
|
|
@pytest.mark.skipif(importlib.util.find_spec('vllm') is None, reason='Optional vLLM package is not installed')
|
|
def test_window_processor_counting_and_request_moves():
|
|
cls = fast.WindowedPenalty
|
|
processor = cls(SimpleNamespace(scheduler_config=SimpleNamespace(max_num_seqs=2)), 'cpu', False)
|
|
a, b = [2, 2, 3], [1]
|
|
added = [(0, SimpleNamespace(extra_args={fast.PENALTY: 2., fast.WINDOW: 3}), [], a),
|
|
(1, SimpleNamespace(extra_args={fast.PENALTY: 3., fast.WINDOW: 2}), [], b)]
|
|
processor.update_state(SimpleNamespace(removed=[], added=added, moved=[]))
|
|
logits = torch.tensor([[0., 9., 12., -2.], [0., 9., 12., -2.]])
|
|
actual = processor.apply(logits.clone())
|
|
torch.testing.assert_close(actual, torch.tensor([[0., 9., 3., -4.], [0., 3., 12., -2.]]))
|
|
processor.update_state(SimpleNamespace(removed=[], added=[], moved=[(0, 1, SimpleNamespace(name='SWAP'))]))
|
|
torch.testing.assert_close(processor.apply(logits.clone()), actual.flip(0))
|
|
processor.update_state(SimpleNamespace(removed=[0], added=[], moved=[(1, 0, SimpleNamespace(name='UNIDIRECTIONAL'))]))
|
|
torch.testing.assert_close(processor.apply(logits[:1].clone()), actual[:1])
|
|
|
|
|
|
def test_fp8_ar_selection_idempotence_and_exact_restore(monkeypatch):
|
|
config = YuE2Config(hidden_size=32, intermediate_size=64, num_hidden_layers=1,
|
|
num_attention_heads=4, num_key_value_heads=2, head_dim=8,
|
|
vocab_size=71, max_latent_frames=64)
|
|
model = YuE2ForCausalLM(config).to(torch.bfloat16)
|
|
original = {key: value.clone() for key, value in model.state_dict().items()}
|
|
constructor = quant.FP8Linear
|
|
monkeypatch.setattr(quant, 'FP8Linear', lambda linear, device: constructor(linear, 'cpu'))
|
|
monkeypatch.setattr(torch.cuda, 'is_available', lambda: True)
|
|
monkeypatch.setattr(torch.cuda, 'get_device_capability', lambda device: (9, 0))
|
|
status = quant.prepare_fp8_ar(model, 'cuda')
|
|
assert status['active_ar_linears'] == 7
|
|
assert isinstance(model.model.layers[0].nar_mlp.up_proj, torch.nn.Linear)
|
|
refs = model._yue2_fp8_originals
|
|
assert quant.prepare_fp8_ar(model, 'cuda') == status and model._yue2_fp8_originals is refs
|
|
quant.restore_ar(model)
|
|
assert all(torch.equal(value, model.state_dict()[key]) for key, value in original.items())
|
|
assert quant.quantization_status(model)['active_ar_linears'] == 0
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='Needs an allocated NVIDIA GPU')
|
|
def test_fp8_scaled_mm_cuda():
|
|
if torch.cuda.get_device_capability() < (8, 9):
|
|
pytest.skip('FP8 kernel requires SM89 or newer')
|
|
torch.manual_seed(17)
|
|
linear = torch.nn.Linear(64, 32, bias=False, dtype=torch.bfloat16, device='cuda')
|
|
layer = quant.FP8Linear(linear, 'cuda')
|
|
for rows in [1, 17]:
|
|
x = torch.randn(rows, 64, dtype=torch.bfloat16, device='cuda')
|
|
actual, expected = layer(x), linear(x)
|
|
assert actual.shape == expected.shape and torch.isfinite(actual).all()
|
|
assert (actual.float() - expected.float()).norm() / expected.float().norm() < .15
|