| 400 | |
| 401 | |
| 402 | def parse_args() -> argparse.Namespace: |
| 403 | parser = argparse.ArgumentParser(description="Stateless Reference/Reward executor benchmark") |
| 404 | parser.add_argument( |
| 405 | "--mode", |
| 406 | choices=["reference", "reward", "both", "paged-kv-compare"], |
| 407 | default="reference", |
| 408 | ) |
| 409 | parser.add_argument("--smoke", action="store_true") |
| 410 | parser.add_argument( |
| 411 | "--compare-paged-kv", |
| 412 | action="store_true", |
| 413 | help="Run the local generation-style paged-KV reservation baseline too.", |
| 414 | ) |
| 415 | parser.add_argument("--batch-size", type=int, default=8) |
| 416 | parser.add_argument("--seq-len", type=int, default=1024) |
| 417 | parser.add_argument("--completion-len", type=int, default=256) |
| 418 | parser.add_argument("--vocab-size", type=int, default=4096) |
| 419 | parser.add_argument("--hidden-dim", type=int, default=256) |
| 420 | parser.add_argument("--paged-kv-layers", type=int, default=4) |
| 421 | parser.add_argument("--paged-kv-heads", type=int, default=8) |
| 422 | parser.add_argument("--paged-kv-head-dim", type=int, default=32) |
| 423 | parser.add_argument("--paged-kv-block-size", type=int, default=16) |
| 424 | parser.add_argument("--paged-kv-blocks", type=int, default=None) |
| 425 | parser.add_argument("--paged-kv-dtype", choices=sorted(DTYPE_CHOICES), default="float16") |
| 426 | parser.add_argument("--valid-density", type=float, default=1.0) |
| 427 | parser.add_argument("--seed", type=int, default=47) |
| 428 | parser.add_argument("--device", default="cpu") |
| 429 | parser.add_argument( |
| 430 | "--output", |
| 431 | type=Path, |
| 432 | default=REPO_ROOT / "task-workspace/issues/issue_47/benchmark.csv", |
| 433 | ) |
| 434 | parser.add_argument("--json-output", type=Path, default=None) |
| 435 | args = parser.parse_args() |
| 436 | |
| 437 | if args.smoke: |
| 438 | args.batch_size = min(args.batch_size, 2) |
| 439 | args.seq_len = min(args.seq_len, 16) |
| 440 | args.completion_len = min(args.completion_len, 8) |
| 441 | args.vocab_size = min(args.vocab_size, 128) |
| 442 | args.hidden_dim = min(args.hidden_dim, 32) |
| 443 | args.paged_kv_layers = min(args.paged_kv_layers, 2) |
| 444 | args.paged_kv_heads = min(args.paged_kv_heads, 2) |
| 445 | args.paged_kv_head_dim = min(args.paged_kv_head_dim, 8) |
| 446 | args.paged_kv_block_size = min(args.paged_kv_block_size, 16) |
| 447 | if args.batch_size <= 0: |
| 448 | raise ValueError("batch-size must be greater than zero") |
| 449 | if args.seq_len < 2: |
| 450 | raise ValueError("seq-len must be at least 2") |
| 451 | if args.completion_len <= 0 or args.completion_len >= args.seq_len: |
| 452 | raise ValueError("completion-len must be greater than zero and less than seq-len") |
| 453 | if args.vocab_size <= 1: |
| 454 | raise ValueError("vocab-size must be greater than one") |
| 455 | if args.hidden_dim <= 0: |
| 456 | raise ValueError("hidden-dim must be greater than zero") |
| 457 | if args.paged_kv_layers <= 0: |
| 458 | raise ValueError("paged-kv-layers must be greater than zero") |
| 459 | if args.paged_kv_heads <= 0: |