Premchan369/Q-TensorFormer

1"""2Tests for Adaptive KV Cache Module.3"""4 5import sys6import os7sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))8 9import torch10import pytest11from src.kv_cache import AdaptiveKVCache, KVPrecision, QuantizedKVTensor12 13 14def test_quantized_kv_tensor():15    tensor = torch.randn(2, 4, 8, 32)16 17    # FP1618    q_fp16 = QuantizedKVTensor(tensor, KVPrecision.FP16)19    rec_fp16 = q_fp16.dequantize()20    assert torch.allclose(tensor, rec_fp16, atol=1e-2)21 22    # INT823    q_int8 = QuantizedKVTensor(tensor, KVPrecision.INT8)24    rec_int8 = q_int8.dequantize()25    cos_sim_int8 = torch.cosine_similarity(tensor.flatten(), rec_int8.flatten(), dim=0)26    assert cos_sim_int8 > 0.99, f"INT8 cosine similarity {cos_sim_int8} too low"27 28    # INT429    q_int4 = QuantizedKVTensor(tensor, KVPrecision.INT4)30    rec_int4 = q_int4.dequantize()31    cos_sim_int4 = torch.cosine_similarity(tensor.flatten(), rec_int4.flatten(), dim=0)32    assert cos_sim_int4 > 0.90, f"INT4 cosine similarity {cos_sim_int4} too low"33    print("✓ test_quantized_kv_tensor passed")34 35 36def test_adaptive_kv_cache_append_and_evict():37    cache = AdaptiveKVCache(max_capacity=16, default_precision=KVPrecision.FP16, window_size=4)38 39    B, H, D = 1, 2, 1640 41    # Append 10 tokens42    k1 = torch.randn(B, H, 10, D)43    v1 = torch.randn(B, H, 10, D)44    out_k1, out_v1 = cache.update(k1, v1)45    assert cache.seq_len == 1046    assert out_k1.shape[-2] == 1047 48    # Append 10 more tokens (exceeds max_capacity 16 -> should trigger eviction)49    k2 = torch.randn(B, H, 10, D)50    v2 = torch.randn(B, H, 10, D)51    out_k2, out_v2 = cache.update(k2, v2)52 53    assert cache.seq_len == 16, f"Expected cache seq_len 16, got {cache.seq_len}"54    assert cache.evicted_tokens_count == 455    assert cache.current_mb > 0.056    print("✓ test_adaptive_kv_cache_append_and_evict passed")57 58 59def test_adaptive_kv_cache_precision_switch():60    cache = AdaptiveKVCache(max_capacity=32, default_precision=KVPrecision.FP16)61    k = torch.randn(1, 2, 8, 16)62    v = torch.randn(1, 2, 8, 16)63    cache.update(k, v)64 65    bytes_fp16 = cache.current_bytes66 67    # Switch to INT868    cache.set_precision(KVPrecision.INT8)69    bytes_int8 = cache.current_bytes70    assert bytes_int8 < bytes_fp16, f"INT8 ({bytes_int8}) should be smaller than FP16 ({bytes_fp16})"71 72    # Switch to INT473    cache.set_precision(KVPrecision.INT4)74    bytes_int4 = cache.current_bytes75    assert bytes_int4 < bytes_int8, f"INT4 ({bytes_int4}) should be smaller than INT8 ({bytes_int8})"76    print("✓ test_adaptive_kv_cache_precision_switch passed")77 78 79if __name__ == "__main__":80    test_quantized_kv_tensor()81    test_adaptive_kv_cache_append_and_evict()82    test_adaptive_kv_cache_precision_switch()83    print("All KV Cache tests passed!")84