Add code long-context attention profiling grid

This commit is contained in:
2026-07-23 23:58:49 +08:00
parent 1d9182f305
commit 288f7b239f
5 changed files with 196 additions and 9 deletions

View File

@@ -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,
)