| |
| """Tests for the node layer: the TQ3 cache wrapper and the attention patch. |
| |
| These paths had no coverage, and they are the ones ComfyUI actually calls. |
| """ |
|
|
| import pytest |
| import torch |
|
|
| import tq3_core |
| import turboquant_nodes |
|
|
|
|
| BLOCK = tq3_core.TQ3_BLOCK |
|
|
|
|
| def _cosine(x, xr): |
| return ( |
| torch.nn.functional.cosine_similarity( |
| x.reshape(-1, x.shape[-1]).float(), |
| xr.reshape(-1, xr.shape[-1]).float(), |
| dim=-1, |
| ) |
| .mean() |
| .item() |
| ) |
|
|
|
|
| @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) |
| def test_wrapper_roundtrip_preserves_dtype_and_shape(dtype): |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| tensor = torch.randn(2, 8, BLOCK, dtype=dtype) |
|
|
| wrapper.store("k", tensor) |
| loaded = wrapper.load("k") |
|
|
| assert loaded.shape == tensor.shape |
| assert loaded.dtype == dtype |
| assert _cosine(tensor, loaded) > 0.97 |
|
|
|
|
| @pytest.mark.parametrize("dim", [320, 640, 1280, BLOCK + 1]) |
| def test_wrapper_pads_dims_that_are_not_a_multiple_of_the_block(dim): |
| """SD/SDXL inner dims are not multiples of 128, so this is the common path.""" |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| tensor = torch.randn(2, 4, dim) |
|
|
| wrapper.store("k", tensor) |
| loaded = wrapper.load("k") |
|
|
| assert loaded.shape == tensor.shape |
| assert _cosine(tensor, loaded) > 0.97 |
|
|
|
|
| def test_wrapper_roundtrip_keeps_structured_tensors(): |
| """A per-channel offset is exactly what real attention tensors carry.""" |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| tensor = torch.randn(2, 8, BLOCK) + torch.linspace(-4.0, 4.0, BLOCK) |
|
|
| wrapper.store("k", tensor) |
|
|
| assert _cosine(tensor, wrapper.load("k")) > 0.97 |
|
|
|
|
| def test_wrapper_stores_raw_tensors_when_disabled(): |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper(enabled=False) |
| tensor = torch.randn(2, BLOCK) |
|
|
| wrapper.store("k", tensor) |
|
|
| assert wrapper.load("k") is tensor |
| assert wrapper.compression_ratio == 0.0 |
|
|
|
|
| def test_wrapper_load_returns_none_for_an_unknown_key(): |
| assert turboquant_nodes.TQ3KVCacheWrapper().load("missing") is None |
|
|
|
|
| def test_wrapper_reports_the_documented_ratio_for_fp16(): |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
|
|
| wrapper.store("k", torch.randn(4, 8, BLOCK, dtype=torch.float16)) |
|
|
| |
| assert wrapper.compression_ratio == pytest.approx(256 / 56, rel=1e-6) |
| assert wrapper.savings_mb > 0 |
| assert "Ratio:" in wrapper.stats_string() |
|
|
|
|
| def test_wrapper_clear_resets_cache_and_stats(): |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| wrapper.store("k", torch.randn(2, BLOCK)) |
|
|
| wrapper.clear() |
|
|
| assert wrapper.load("k") is None |
| assert wrapper.compression_ratio == 0.0 |
|
|
|
|
| def test_attn_patch_passes_through_a_none_context(): |
| """ComfyUI's attn2 path does not coerce a None context (attn1 does).""" |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| patch = turboquant_nodes._make_attn_patch(wrapper) |
| hidden = torch.randn(2, 16, 320) |
|
|
| q, k, v = patch(hidden, None, None, {}) |
|
|
| assert q is hidden |
| assert k is None |
| assert v is None |
|
|
|
|
| def test_attn_patch_roundtrips_kv_and_leaves_no_cache_entries(): |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| patch = turboquant_nodes._make_attn_patch(wrapper) |
| hidden = torch.randn(2, 16, 320) |
|
|
| q, k, v = patch(hidden, hidden, hidden, {}) |
|
|
| assert q is hidden |
| assert k.shape == hidden.shape and v.shape == hidden.shape |
| assert _cosine(hidden, k) > 0.97 |
| assert wrapper._cache == {} |
| assert wrapper._stats["num_stores"] == 2 |
|
|
|
|
| def test_attn_patch_skips_tensors_narrower_than_one_block(): |
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| patch = turboquant_nodes._make_attn_patch(wrapper) |
| narrow = torch.randn(2, 16, 64) |
|
|
| q, k, v = patch(narrow, narrow, narrow, {}) |
|
|
| assert k is narrow and v is narrow |
| assert wrapper._stats["num_stores"] == 0 |
|
|
|
|
| def test_node_registration_exposes_both_documented_nodes(): |
| assert set(turboquant_nodes.NODE_CLASS_MAPPINGS) == { |
| "TurboQuantPatch", |
| "TurboQuantInfo", |
| } |
| assert set(turboquant_nodes.NODE_DISPLAY_NAME_MAPPINGS) == { |
| "TurboQuantPatch", |
| "TurboQuantInfo", |
| } |
| for name, cls in turboquant_nodes.NODE_CLASS_MAPPINGS.items(): |
| assert hasattr(cls, cls.FUNCTION), name |
| assert "required" in cls.INPUT_TYPES() |
|
|
|
|
| def test_info_node_reports_stats_once_a_wrapper_is_active(): |
| turboquant_nodes._active_wrapper = None |
| try: |
| (idle,) = turboquant_nodes.TurboQuantInfo().info(model=None) |
| assert "Not yet active" in idle |
|
|
| wrapper = turboquant_nodes.TQ3KVCacheWrapper() |
| wrapper.store("k", torch.randn(2, BLOCK, dtype=torch.float16)) |
| turboquant_nodes._active_wrapper = wrapper |
|
|
| (active,) = turboquant_nodes.TurboQuantInfo().info(model=None) |
| assert "Stores: 1" in active |
| finally: |
| turboquant_nodes._active_wrapper = None |
|
|
