From 7f9a05a7f85e2d3d8f31c385a8d3feda0c4a42b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 22:17:04 +0900 Subject: [PATCH 1/4] test(kvcache): size the fp8 slot round-trip to the rows it indexes `test_store_kv_writes_the_slot_the_attend_kernel_will_read` reaches row 256 -- its comment asks for the 255/256 page boundary -- against a four-page, 256-slot pool, so `codes[out_loc]` gathers one row past the view. Five pages is what that row list needs. The gather raises a device-side assert, and the CUDA context does not recover from it, so in a single-process run everything scheduled afterwards is reported as failing as well -- not because those tests stop working, but because there is no longer a context to run them in. One file per process does not show that. Assisted-by: Claude Opus 5 --- tests/kvcache/test_qsa_pool_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From dabe93ccd18fdf110d02088a93a754d16adfcd00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 21:02:36 +0900 Subject: [PATCH 2/4] test(kernels): put the scale-one encoder's V tensor on the device `test_encoder_inverts_the_grid_through_the_scale_one_path` moves `rows` to the device and leaves the V tensor beside it on the host, so the kernel is handed a CPU pointer. Assisted-by: Claude Opus 5 --- tests/kernels/test_kv_fp8.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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, :]) From 5febeeedcffbdda105493461070683c93afa0075 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 21:08:05 +0900 Subject: [PATCH 3/4] test(kvcache): give the layer-ids remap test a model deep enough for its ids `test_layer_ids_remap_applies_to_scales_too` backs `layer_ids=(1, 3)` on a pool built with `num_layers=LAYERS`, and LAYERS is 3, so id 3 is one past the end and the constructor raises before the assertion it is there to make. The helper now takes the depth. Assisted-by: Claude Opus 5 --- tests/kvcache/test_mha_pool_fp8.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) 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) From 0820ff4e9b5a774940f34b74b2a79c1ca2fa8013 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 21:09:48 +0900 Subject: [PATCH 4/4] test(kernels): give the triton-attention doubles the scale accessors the backend now reads `k_scale` / `v_scale` arrived with the fp8 store, and the backend reads them on every path, including a 16-bit pool -- which answers None. The two hand-rolled `FakeKVCache` classes in this file do not inherit the base pool, so they were left without them and raise `AttributeError` instead. Assisted-by: Claude Opus 5 --- tests/kernels/test_triton_attention.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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)