From 1b08c37dabcad1cadd341332dec976668f85d96c Mon Sep 17 00:00:00 2001 From: Po-Han Huang Date: Fri, 2 Oct 2026 00:42:36 -0700 Subject: [PATCH] [Test] Add get_swa_key_page_size to the Q8KV8 sparse-prefill fake KV pool #42128 made DeepseekV4AttnBackend read the SWA page size through token_to_kv_pool.get_swa_key_page_size(), but the fake _TokenToKVPool in test_q8kv8_sparse_prefill_backend.py does not define it. Four tests in the jit-kernel 1-gpu-large shard now fail on main with AttributeError. Add the accessor to the fake pool, mirroring the real pool's default branch. Co-Authored-By: Claude Opus 5.5 --- .../kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py index 57d061339eb2..ef87a256979a 100644 --- a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py +++ b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py @@ -95,6 +95,9 @@ def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor: def get_swa_key_layout(self) -> KVLayout: return KVLayout.V4 + def get_swa_key_page_size(self) -> int: + return self.swa_kv_pool.page_size + def get_extra_key_layout(self, layer_id: int) -> KVLayout: _ = layer_id return KVLayout.V4