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