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
This commit is contained in:
@@ -111,7 +111,7 @@ struct PagedKV {
|
|||||||
return p.max_q_len;
|
return p.max_q_len;
|
||||||
}
|
}
|
||||||
HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
|
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)
|
// prefill: Q rows start at qo_indptr[batch] (ragged batch base)
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import torch
|
|||||||
|
|
||||||
from astrai.config.model_config import AutoRegressiveLMConfig
|
from astrai.config.model_config import AutoRegressiveLMConfig
|
||||||
from astrai.model.transformer import AutoRegressiveLM
|
from astrai.model.transformer import AutoRegressiveLM
|
||||||
|
from tests.conftest import skip_no_kernel # noqa: F401 re-export for test modules
|
||||||
|
|
||||||
D = 64
|
D = 64
|
||||||
CFG = dict(
|
CFG = dict(
|
||||||
|
|||||||
Reference in New Issue
Block a user