fix: use max_context_len for stable num_splits in paged decode

- PagedKV::host_kv_len now returns max_context_len instead of max_seq_len - Eliminates grid-z instability for CUDA graph capture/replay - Restore skip_no_kernel re-export accidentally removed by ruff --fix
2026-08-07 14:42:53 +08:00
parent 55ee258e95
commit 0e7fe57d96
2 changed files with 2 additions and 1 deletions
@@ -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(