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 ca2adb70e..7347bada1 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
@@ -203,6 +203,7 @@ def _make_sparse_prefill_case(
         ),
         page_size=page_size,
     )
+    backend.token_to_kv_pool = token_to_kv_pool
 
     generator = torch.Generator(device=device)
     generator.manual_seed(11)