test: keep the fp8 tests runnable in a single pytest process

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, :])

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)

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)

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