[tests] Added test_niah_standalone.py.

This commit is contained in:
Zijie Tian
2026-01-12 00:16:37 +08:00
parent 5895de0c97
commit a6cc703d73
6 changed files with 686 additions and 9 deletions

View File

@@ -61,6 +61,15 @@ class Config:
self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
assert self.max_num_batched_tokens >= self.max_model_len
# CPU offload mode only supports single sequence (layer-wise processing)
if self.enable_cpu_offload and self.max_num_seqs != 1:
import logging
logging.warning(
f"CPU offload mode only supports single sequence. "
f"Overriding max_num_seqs from {self.max_num_seqs} to 1."
)
self.max_num_seqs = 1
# Override torch_dtype if user specified
if self.dtype is not None:
dtype_map = {