| 29 | |
| 30 | |
| 31 | def parse_args(): |
| 32 | parser = argparse.ArgumentParser(description="KVCOMM Experiments on TTFT Benchmark") |
| 33 | parser.add_argument( |
| 34 | "--mode", |
| 35 | type=str, |
| 36 | default="FullConnected", |
| 37 | choices=["DirectAnswer", "FullConnected", "Random", "Chain", "Debate", "Layered", "Star", "Mesh"], |
| 38 | help="The communication topology among agents.", |
| 39 | ) |
| 40 | parser.add_argument("--batch_size", type=int, default=1) |
| 41 | parser.add_argument( |
| 42 | "--agent_names", |
| 43 | nargs="+", |
| 44 | type=str, |
| 45 | default=["CopyMachine"], |
| 46 | help="List of agent names to include in the graph." |
| 47 | ) |
| 48 | parser.add_argument( |
| 49 | "--agent_nums", |
| 50 | nargs="+", |
| 51 | type=int, |
| 52 | default=[5], |
| 53 | help="List of counts corresponding to each agent name." |
| 54 | ) |
| 55 | parser.add_argument("--llm_name", type=str, default="meta-llama/Llama-3.1-8B-Instruct") |
| 56 | parser.add_argument("--domain", type=str, default="COPY") |
| 57 | parser.add_argument("--decision_method", type=str, default=None) |
| 58 | parser.add_argument( |
| 59 | "--execution_mode", |
| 60 | type=str, |
| 61 | default="allow_kv_reuse", |
| 62 | choices=["default", "allow_kv_reuse"], |
| 63 | help="Execution strategy for the graph.", |
| 64 | ) |
| 65 | parser.add_argument("--output_dir", type=str, default=str(PROJECT_ROOT / "result" / "TTFT_Benchmark"), help="Directory to save the output results.") |
| 66 | parser.add_argument("--prefix", type=str, default="The task is:\n\n", help="The prefix text for the input query, kept the same as the default dense prefill mode.") |
| 67 | parser.add_argument("--samples", type=int, default=100, help="Number of 1K-token samples") |
| 68 | parser.add_argument("--kv-threshold", type=float, default=1.0, help="Threshold for key-value memory usage.") |
| 69 | parser.add_argument("--kv-max-anchor-num", type=int, default=20, help="Maximum number of anchors for key-value memory.") |
| 70 | parser.add_argument("--kv-window-size", type=int, default=5, help="Window size for key-value memory update.") |
| 71 | parser.add_argument("--kv-thread-workers", type=int, default=None, help="Number of thread workers for key-value memory processing.") |
| 72 | parser.add_argument("--kv-worker-timeout", type=float, default=None, help="Timeout for key-value memory workers processing.") |
| 73 | |
| 74 | args = parser.parse_args() |
| 75 | result_path = Path(args.output_dir) |
| 76 | result_path.mkdir(parents=True, exist_ok=True) |
| 77 | if len(args.agent_names) != len(args.agent_nums): |
| 78 | parser.error("The number of agent names must match the number of agent counts.") |
| 79 | return args |
| 80 | |
| 81 | def _make_random_token_sequence(length: int) -> str: |
| 82 | symbols = random.choices(["Δ", "Ω"], k=length) |