Add code long-context attention profiling grid
This commit is contained in:
@@ -37,6 +37,7 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
parser.add_argument("--warmup-iters", type=int, default=3)
|
||||
parser.add_argument("--repeats", type=int, default=5)
|
||||
parser.add_argument("--max-model-len", type=int, default=40960)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--profile-kv-update", action="store_true")
|
||||
parser.add_argument(
|
||||
@@ -115,7 +116,7 @@ def main() -> None:
|
||||
trust_remote_code=False,
|
||||
dtype="bfloat16",
|
||||
seed=0,
|
||||
max_model_len=40960,
|
||||
max_model_len=args.max_model_len,
|
||||
)
|
||||
cache_config = CacheConfig(block_size=config.block_size, cache_dtype="auto")
|
||||
cache_config.num_gpu_blocks = max_num_blocks
|
||||
@@ -124,7 +125,7 @@ def main() -> None:
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=256,
|
||||
max_num_batched_tokens=8192,
|
||||
max_model_len=40960,
|
||||
max_model_len=args.max_model_len,
|
||||
is_encoder_decoder=False,
|
||||
enable_chunked_prefill=True,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user