diff --git a/csrc/kernels/attn_kv_source.cuh b/csrc/kernels/attn_kv_source.cuh
index 416cd9c..4c181cd 100644
--- a/csrc/kernels/attn_kv_source.cuh
+++ b/csrc/kernels/attn_kv_source.cuh
@@ -111,7 +111,7 @@ struct PagedKV {
         return p.max_q_len;
     }
     HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
-        return p.max_seq_len;
+        return p.max_context_len;
     }
 
     // prefill: Q rows start at qo_indptr[batch] (ragged batch base)
diff --git a/tests/extension/conftest.py b/tests/extension/conftest.py
index 8e3e373..e0c18ac 100644
--- a/tests/extension/conftest.py
+++ b/tests/extension/conftest.py
@@ -5,6 +5,7 @@ import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.model.transformer import AutoRegressiveLM
+from tests.conftest import skip_no_kernel  # noqa: F401 re-export for test modules
 
 D = 64
 CFG = dict(