diff --git a/tests/kernels/test_kv_fp8.py b/tests/kernels/test_kv_fp8.py index 26deef2ac..817646886 100644 --- a/tests/kernels/test_kv_fp8.py +++ b/tests/kernels/test_kv_fp8.py @@ -179,7 +179,8 @@ def test_encoder_inverts_the_grid_through_the_scale_one_path(): expected[t, j + 1] = code k_cache, _, k_scale, _ = _store( - rows.to(DEV, dtype=torch.bfloat16), torch.zeros(tokens, 1, dim, dtype=torch.bfloat16) + rows.to(DEV, dtype=torch.bfloat16), + torch.zeros(tokens, 1, dim, dtype=torch.bfloat16, device=DEV), ) assert torch.equal(k_scale[:, 0], torch.ones_like(k_scale[:, 0])) got = _canonical_zero(_as_bytes(k_cache)[:, 0, :]) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index ffff4a09b..61ce5360d 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -67,6 +67,14 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + # A 16-bit pool answers None here. The backend reads it on every path since the + # fp8 store landed, so the double answers it too. + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + kv_cache = FakeKVCache() monkeypatch.setattr( "freetoken.attention.triton.get_global_ctx", @@ -636,6 +644,14 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + # A 16-bit pool answers None here. The backend reads it on every path since the + # fp8 store landed, so the double answers it too. + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 page_table = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32, device=device) diff --git a/tests/kvcache/test_mha_pool_fp8.py b/tests/kvcache/test_mha_pool_fp8.py index 77bd2986e..dd3f0542d 100644 --- a/tests/kvcache/test_mha_pool_fp8.py +++ b/tests/kvcache/test_mha_pool_fp8.py @@ -35,13 +35,13 @@ def _tp(): _init_tp() -def _pool(kv_quant="fp8", num_pages=PAGES, layer_ids=None): +def _pool(kv_quant="fp8", num_pages=PAGES, layer_ids=None, num_layers=LAYERS): from freetoken.kvcache.mha_pool import MHAKVCache _init_tp() return MHAKVCache( num_kv_heads=HEADS, - num_layers=LAYERS, + num_layers=num_layers, head_dim=DIM, num_pages=num_pages, page_size=PAGE_SIZE, @@ -84,8 +84,10 @@ def test_fp8_pool_keeps_geometry_and_adds_scale_views(): def test_layer_ids_remap_applies_to_scales_too(): # Hybrid GDN models back only their full-attention layers; a scale view that # forgot the remap would hand layer 7's rows to layer 2's attention. + # (1, 3) names a subset of a model, so layer id 3 has to be inside it: LAYERS is 3, which + # makes 3 one past the end, so this pool is told the depth the ids imply. layer_ids = (1, 3) - pool = _pool(layer_ids=layer_ids) + pool = _pool(layer_ids=layer_ids, num_layers=4) assert pool._kv_buffer.shape[1] == 2 with pytest.raises(KeyError): pool.k_scale(0) diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py index 13fa7f468..04426b876 100644 --- a/tests/kvcache/test_qsa_pool_fp8.py +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -168,8 +168,8 @@ def test_store_kv_writes_the_slot_the_attend_kernel_will_read(): """out_loc numbering (page * page_size + offset) is the contract between the fused writer and the attend kernel's scale slot arithmetic -- this is that round trip.""" torch.manual_seed(0) - pool = _pool(num_pages=4) - slots = 4 * PAGE_SIZE + pool = _pool(num_pages=5) + slots = 5 * PAGE_SIZE rows = (0, 1, 63, 64, 255, 256) # page boundaries included: 63/64 and 255/256 k = torch.randn(len(rows), HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 3.0 v = torch.randn(len(rows), HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 0.25