from __future__ import annotations import argparse import asyncio from pathlib import Path from agentic_pd_hybrid.benchmark import BenchmarkConfig, run_live_benchmark from agentic_pd_hybrid.launcher import build_launch_plan from agentic_pd_hybrid.microbench import SmallAppendTraceConfig, write_small_append_trace from agentic_pd_hybrid.profile import ProfileConfig, print_profile_summary, write_profile from agentic_pd_hybrid.replay import ReplayConfig, replay_trace from agentic_pd_hybrid.sampling import SessionSampleConfig, sample_trace_sessions from agentic_pd_hybrid.trace_profiles import ( NormalizeTraceLengthsConfig, normalize_trace_lengths, ) from agentic_pd_hybrid.topology import build_single_node_topology def _normalize_mechanism_name(name: str) -> str: normalized = name.strip().lower() aliases = { "pd-disagg": "pd-disaggregation", "pd-disaggregation": "pd-disaggregation", "pd-hybrid": "pd-disaggregation", "baseline-pd-disagg": "pd-disaggregation", "pd-colo": "pd-colo", "direct-d": "pd-colo", "colocation": "pd-colo", "kvcache-centric": "kvcache-centric", "turn2+-direct-to-d": "kvcache-centric", "pd-with-d-append": "kvcache-centric", } if normalized not in aliases: raise ValueError(f"Unsupported mechanism: {name}") return aliases[normalized] def _parse_gpu_id_list(value: str | None) -> tuple[int, ...] | None: if value is None: return None items = [item.strip() for item in value.split(",") if item.strip()] if not items: return tuple() return tuple(int(item) for item in items) def main() -> None: parser = argparse.ArgumentParser(description="Agentic PD hybrid prototype") subparsers = parser.add_subparsers(dest="command", required=True) print_launch = subparsers.add_parser( "print-launch", help="Print one-node SGLang PD launch commands", ) _add_topology_arguments(print_launch) print_launch.add_argument("--prefill-policy", default="round_robin") print_launch.add_argument("--decode-policy", default="manual") replay = subparsers.add_parser( "replay", help="Replay trace and log request-level metrics", ) _add_topology_arguments(replay) replay.add_argument("--trace", type=Path, required=True) replay.add_argument("--output", type=Path, required=True) replay.add_argument( "--policy", choices=["default", "sticky", "kv-aware"], default="sticky", ) replay.add_argument( "--mechanism", choices=[ "pd-disaggregation", "pd-hybrid", "pd-disagg", "pd-colo", "direct-d", "kvcache-centric", "turn2+-direct-to-d", "pd-with-d-append", ], default="pd-disaggregation", ) replay.add_argument("--router-url") replay.add_argument("--model") replay.add_argument( "--header-mode", choices=["auto", "none", "routing-key", "target-worker"], default="auto", ) replay.add_argument( "--request-limit", type=int, default=None, help="Replay at most this many requests", ) replay.add_argument( "--no-pace", action="store_true", help="Disable wall-clock pacing from trace timestamps", ) replay.add_argument( "--time-scale", type=float, default=1.0, help="Scale trace timing by this factor when pacing is enabled", ) replay.add_argument( "--concurrency-limit", type=int, default=32, ) replay.add_argument( "--timeout-s", type=float, default=600.0, ) replay.add_argument( "--no-stream", action="store_true", help="Use non-streaming OpenAI responses for more robust E2E-only replay.", ) replay.add_argument( "--stream-idle-timeout-s", type=float, default=900.0, help="Abort a streaming request if no SSE line arrives within this many seconds.", ) replay.add_argument( "--kvcache-direct-max-uncached-tokens", type=int, default=2048, help="For kvcache-centric routing, bypass P when the uncached suffix is at most this many tokens.", ) replay.add_argument( "--kvcache-admission-mode", choices=["router", "worker"], default="router", help=( "For kvcache-centric routing, use router shadow-state admission " "or query the decode worker on the critical path." ), ) replay.add_argument( "--kvcache-seed-max-resident-tokens", type=int, default=None, help=( "For kvcache-centric routing, do not seed/reseed a decode session " "when input+output tokens exceed this value." ), ) replay.add_argument( "--kvcache-seed-max-output-tokens", type=int, default=None, help=( "For kvcache-centric routing, do not seed/reseed a decode session " "when output tokens exceed this value." ), ) replay.add_argument( "--kvcache-seed-min-turn-id", type=int, default=1, help=( "For kvcache-centric routing, do not seed/reseed a decode session " "before this turn id." ), ) replay.add_argument( "--kvcache-seed-only-multiturn-sessions", action="store_true", help=( "Oracle ablation for kvcache-centric routing: only seed sessions " "that have more than one turn in the replay trace." ), ) replay.add_argument( "--kvcache-prefill-backup-policy", choices=["release-after-transfer", "capacity-backup"], default="release-after-transfer", help=( "For kvcache-centric seed/reseed, release the P-side session after " "P->D transfer or keep a capacity-limited P-side backup." ), ) replay.add_argument( "--kvcache-seed-max-inflight-decode", type=int, default=3, help=( "For kvcache-centric routing, skip seed/reseed when the router " "shadow inflight decode load assigned to the target D exceeds this value. " "Use a negative value to disable this filter." ), ) replay.add_argument( "--kvcache-seed-max-decode-transfer-queue-reqs", type=int, default=None, help=( "For kvcache-centric worker admission, skip seed/reseed when the " "target D reports more transfer-queue requests than this value. " "Use a negative value to disable this filter." ), ) replay.add_argument( "--kvcache-direct-max-decode-transfer-queue-reqs", type=int, default=None, help=( "For kvcache-centric worker admission, skip direct-to-D append when " "the target D reports more transfer-queue requests than this value. " "Use a negative value to disable this filter." ), ) replay.add_argument( "--kvcache-prefill-priority-eviction", action="store_true", help=( "For kvcache-centric routing, mark P-side prefixes predicted to move " "direct-to-D as lower priority. Requires P workers to use " "--radix-eviction-policy priority." ), ) replay.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100) replay.add_argument("--kvcache-prefill-normal-priority", type=int, default=100) replay.add_argument( "--pool-poll-interval-s", type=float, default=0.0, help=( "Poll each P/D worker's /server_info every N seconds and write a " "time-series snapshot to /d-pool-timeseries.jsonl. " "0 disables polling." ), ) replay.add_argument( "--pool-poll-no-sessions", action="store_true", help=( "Disable per-session detail in the pool timeseries (smaller files)." ), ) replay.add_argument( "--enable-backpressure", action="store_true", help=( "Honor recommended_pause_ms hints from D's admission endpoint. " "When set, replay sleeps before issuing requests to a saturated D. " "Default off — preserves baseline behavior." ), ) replay.add_argument( "--backpressure-max-pause-s", type=float, default=2.0, help="Cap on per-request backpressure sleep, regardless of D hint.", ) replay.add_argument( "--kvcache-migration-reject-threshold", type=int, default=3, help=( "Per-(session, D) admission-reject count after which KvAwarePolicy " "skips that D for the session (forces migration). 0 disables. " "See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1." ), ) sample = subparsers.add_parser( "sample-sessions", help="Sample a session-granularity trace shard for live benchmarking", ) sample.add_argument("--trace", type=Path, required=True) sample.add_argument("--output", type=Path, required=True) sample.add_argument("--target-duration-s", type=float, default=600.0) sample.add_argument("--start-time-s", type=float, default=0.0) sample.add_argument("--session-sample-rate", type=float, default=0.01) sample.add_argument("--min-turns", type=int, default=1) sample.add_argument("--max-requests", type=int, default=None) sample.add_argument( "--profile", choices=["default", "small-append"], default="default", help="Optional workload-shape filter for live benchmarks.", ) sample.add_argument("--min-initial-input-tokens", type=int, default=None) sample.add_argument("--max-initial-input-tokens", type=int, default=None) sample.add_argument("--max-append-input-tokens", type=int, default=None) sample.add_argument("--max-output-tokens", type=int, default=None) sample.add_argument("--min-overlap-ratio", type=float, default=None) normalize = subparsers.add_parser( "normalize-trace-lengths", help="Rewrite a trace to a fixed turn1/append/output length profile", ) normalize.add_argument("--trace", type=Path, required=True) normalize.add_argument("--output", type=Path, required=True) normalize.add_argument("--initial-input-length", type=int, default=10_000) normalize.add_argument("--append-input-length", type=int, default=1_000) normalize.add_argument("--output-length", type=int, default=1_000) normalize.add_argument("--max-requests", type=int, default=None) profile = subparsers.add_parser( "profile", help="Profile a trace and optional request metrics for routing analysis", ) profile.add_argument("--trace", type=Path, required=True) profile.add_argument("--output", type=Path, default=None) profile.add_argument("--metrics", type=Path, default=None) profile.add_argument("--baseline-metrics", type=Path, default=None) profile.add_argument("--candidate-metrics", type=Path, default=None) profile.add_argument("--direct-max-uncached-tokens", type=int, default=2048) micro = subparsers.add_parser( "make-small-append-trace", help="Generate a synthetic multi-turn trace with small turn2+ appends", ) micro.add_argument("--output", type=Path, required=True) micro.add_argument("--session-count", type=int, default=8) micro.add_argument("--turns-per-session", type=int, default=3) micro.add_argument("--initial-input-length", type=int, default=10_000) micro.add_argument("--append-input-length", type=int, default=1_000) micro.add_argument("--output-length", type=int, default=1_000) micro.add_argument("--inter-turn-gap-s", type=float, default=1.0) micro.add_argument("--session-stagger-s", type=float, default=0.1) benchmark = subparsers.add_parser( "benchmark-live", help="Launch a real PD stack, sample sessions, and collect live E2E numbers", ) _add_topology_arguments(benchmark) benchmark.add_argument("--trace", type=Path, required=True) benchmark.add_argument( "--policy", choices=["default", "sticky", "kv-aware"], default="sticky", ) benchmark.add_argument( "--mechanism", choices=[ "pd-disaggregation", "pd-hybrid", "pd-disagg", "pd-colo", "direct-d", "kvcache-centric", "turn2+-direct-to-d", "pd-with-d-append", ], default="pd-disaggregation", ) benchmark.add_argument("--output-root", type=Path, default=Path("outputs/live")) benchmark.add_argument("--target-duration-s", type=float, default=600.0) benchmark.add_argument("--start-time-s", type=float, default=0.0) benchmark.add_argument("--session-sample-rate", type=float, default=0.01) benchmark.add_argument("--min-turns", type=int, default=1) benchmark.add_argument("--time-scale", type=float, default=1.0) benchmark.add_argument("--concurrency-limit", type=int, default=32) benchmark.add_argument("--timeout-s", type=float, default=1200.0) benchmark.add_argument( "--request-timeout-s", type=float, default=None, help=( "Per-request replay/router timeout. If unset, --timeout-s is used. " "--timeout-s still controls stack startup." ), ) benchmark.add_argument( "--no-stream", action="store_true", help="Use non-streaming OpenAI responses for E2E-only live benchmarking.", ) benchmark.add_argument( "--stream-idle-timeout-s", type=float, default=900.0, help="Abort a streaming request if no SSE line arrives within this many seconds.", ) benchmark.add_argument( "--kvcache-direct-max-uncached-tokens", type=int, default=2048, help="For kvcache-centric routing, bypass P when the uncached suffix is at most this many tokens.", ) benchmark.add_argument( "--kvcache-admission-mode", choices=["router", "worker"], default="router", help=( "For kvcache-centric routing, use router shadow-state admission " "or query the decode worker on the critical path." ), ) benchmark.add_argument( "--kvcache-seed-max-resident-tokens", type=int, default=None, help=( "For kvcache-centric routing, do not seed/reseed a decode session " "when input+output tokens exceed this value." ), ) benchmark.add_argument( "--kvcache-seed-max-output-tokens", type=int, default=None, help=( "For kvcache-centric routing, do not seed/reseed a decode session " "when output tokens exceed this value." ), ) benchmark.add_argument( "--kvcache-seed-min-turn-id", type=int, default=1, help=( "For kvcache-centric routing, do not seed/reseed a decode session " "before this turn id." ), ) benchmark.add_argument( "--kvcache-seed-only-multiturn-sessions", action="store_true", help=( "Oracle ablation for kvcache-centric routing: only seed sessions " "that have more than one turn in the replay trace." ), ) benchmark.add_argument( "--kvcache-prefill-backup-policy", choices=["release-after-transfer", "capacity-backup"], default="release-after-transfer", help=( "For kvcache-centric seed/reseed, release the P-side session after " "P->D transfer or keep a capacity-limited P-side backup." ), ) benchmark.add_argument( "--kvcache-seed-max-inflight-decode", type=int, default=3, help=( "For kvcache-centric routing, skip seed/reseed when the router " "shadow inflight decode load assigned to the target D exceeds this value. " "Use a negative value to disable this filter." ), ) benchmark.add_argument( "--kvcache-seed-max-decode-transfer-queue-reqs", type=int, default=None, help=( "For kvcache-centric worker admission, skip seed/reseed when the " "target D reports more transfer-queue requests than this value. " "Use a negative value to disable this filter." ), ) benchmark.add_argument( "--kvcache-direct-max-decode-transfer-queue-reqs", type=int, default=None, help=( "For kvcache-centric worker admission, skip direct-to-D append when " "the target D reports more transfer-queue requests than this value. " "Use a negative value to disable this filter." ), ) benchmark.add_argument( "--kvcache-prefill-priority-eviction", action="store_true", help=( "For kvcache-centric benchmark-live, launch P workers with priority " "radix eviction and mark direct-to-D predicted prefixes lower priority." ), ) benchmark.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100) benchmark.add_argument("--kvcache-prefill-normal-priority", type=int, default=100) benchmark.add_argument( "--pool-poll-interval-s", type=float, default=0.0, help=( "Poll each P/D worker's /server_info every N seconds and write a " "time-series snapshot to /d-pool-timeseries.jsonl. " "0 disables polling." ), ) benchmark.add_argument( "--pool-poll-no-sessions", action="store_true", help=( "Disable per-session detail in the pool timeseries (smaller files)." ), ) benchmark.add_argument( "--enable-backpressure", action="store_true", help=( "Honor recommended_pause_ms hints from D's admission endpoint." ), ) benchmark.add_argument( "--backpressure-max-pause-s", type=float, default=2.0, help="Cap on per-request backpressure sleep, regardless of D hint.", ) benchmark.add_argument( "--kvcache-migration-reject-threshold", type=int, default=3, help=( "Per-(session, D) admission-reject count after which KvAwarePolicy " "skips that D for the session (forces migration). 0 disables. " "See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1." ), ) benchmark.add_argument( "--sample-profile", choices=["default", "small-append"], default="default", help="Optional session-shape filter applied before live replay.", ) benchmark.add_argument("--min-initial-input-tokens", type=int, default=None) benchmark.add_argument("--max-initial-input-tokens", type=int, default=None) benchmark.add_argument("--max-append-input-tokens", type=int, default=None) benchmark.add_argument("--max-output-tokens", type=int, default=None) benchmark.add_argument("--min-overlap-ratio", type=float, default=None) args = parser.parse_args() if args.command == "print-launch": topology = _topology_from_args(args) has_pd = bool(topology.prefill_workers and topology.decode_workers) has_direct_only = bool( topology.direct_workers and not topology.prefill_workers and not topology.decode_workers ) plan = build_launch_plan( topology, prefill_policy=args.prefill_policy, decode_policy=args.decode_policy, include_router=has_pd or has_direct_only, naive_dp=has_direct_only, ) print(plan.render()) return if args.command == "replay": topology = _topology_from_args(args) config = ReplayConfig( trace_path=args.trace, output_path=args.output, policy_name=args.policy, mechanism_name=_normalize_mechanism_name(args.mechanism), topology=topology, router_url=args.router_url, model_name=args.model, pace=not args.no_pace, time_scale=args.time_scale, request_limit=args.request_limit, concurrency_limit=args.concurrency_limit, header_mode=args.header_mode, timeout_s=args.timeout_s, stream=not args.no_stream, stream_idle_timeout_s=args.stream_idle_timeout_s, kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens, kvcache_admission_mode=args.kvcache_admission_mode, kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens, kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens, kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id, kvcache_seed_only_multiturn_sessions=( args.kvcache_seed_only_multiturn_sessions ), kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy, kvcache_seed_max_inflight_decode=( None if args.kvcache_seed_max_inflight_decode < 0 else args.kvcache_seed_max_inflight_decode ), kvcache_seed_max_decode_transfer_queue_reqs=( None if args.kvcache_seed_max_decode_transfer_queue_reqs is None or args.kvcache_seed_max_decode_transfer_queue_reqs < 0 else args.kvcache_seed_max_decode_transfer_queue_reqs ), kvcache_direct_max_decode_transfer_queue_reqs=( None if args.kvcache_direct_max_decode_transfer_queue_reqs is None or args.kvcache_direct_max_decode_transfer_queue_reqs < 0 else args.kvcache_direct_max_decode_transfer_queue_reqs ), kvcache_prefill_priority_eviction=( args.kvcache_prefill_priority_eviction ), kvcache_prefill_direct_priority=args.kvcache_prefill_direct_priority, kvcache_prefill_normal_priority=args.kvcache_prefill_normal_priority, pool_poll_interval_s=args.pool_poll_interval_s, pool_poll_include_sessions=not args.pool_poll_no_sessions, enable_backpressure=args.enable_backpressure, backpressure_max_pause_s=args.backpressure_max_pause_s, kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold, ) results = asyncio.run(replay_trace(config)) print( f"wrote {len(results)} request records to {args.output} and " f"{args.output}{'.summary.json'}" ) return if args.command == "profile": if (args.baseline_metrics is None) != (args.candidate_metrics is None): raise ValueError( "--baseline-metrics and --candidate-metrics must be provided together" ) report = write_profile( ProfileConfig( trace_path=args.trace, output_path=args.output, metrics_path=args.metrics, baseline_metrics_path=args.baseline_metrics, candidate_metrics_path=args.candidate_metrics, direct_max_uncached_tokens=args.direct_max_uncached_tokens, ) ) print_profile_summary(report) if args.output is not None: print(f"wrote profile to {args.output}") return if args.command == "sample-sessions": summary = sample_trace_sessions( SessionSampleConfig( trace_path=args.trace, output_path=args.output, target_duration_s=args.target_duration_s, start_time_s=args.start_time_s, session_sample_rate=args.session_sample_rate, min_turns=args.min_turns, max_requests=args.max_requests, profile=args.profile, min_initial_input_tokens=args.min_initial_input_tokens, max_initial_input_tokens=args.max_initial_input_tokens, max_append_input_tokens=args.max_append_input_tokens, max_output_tokens=args.max_output_tokens, min_overlap_ratio=args.min_overlap_ratio, ) ) print( f"wrote {summary.request_count} requests from {summary.session_count} sessions " f"covering {summary.sampled_duration_s:.3f}s to {args.output}" ) return if args.command == "normalize-trace-lengths": summary = normalize_trace_lengths( NormalizeTraceLengthsConfig( trace_path=args.trace, output_path=args.output, initial_input_length=args.initial_input_length, append_input_length=args.append_input_length, output_length=args.output_length, max_requests=args.max_requests, ) ) print( f"wrote {summary.request_count} normalized requests from " f"{summary.session_count} sessions to {args.output}" ) return if args.command == "make-small-append-trace": summary = write_small_append_trace( SmallAppendTraceConfig( output_path=args.output, session_count=args.session_count, turns_per_session=args.turns_per_session, initial_input_length=args.initial_input_length, append_input_length=args.append_input_length, output_length=args.output_length, inter_turn_gap_s=args.inter_turn_gap_s, session_stagger_s=args.session_stagger_s, ) ) print( f"wrote {summary.request_count} requests across {summary.session_count} sessions " f"to {args.output}" ) return if args.command == "benchmark-live": topology = _topology_from_args(args) artifacts = run_live_benchmark( BenchmarkConfig( trace_path=args.trace, output_root=args.output_root, topology=topology, policy_name=args.policy, mechanism_name=_normalize_mechanism_name(args.mechanism), target_duration_s=args.target_duration_s, start_time_s=args.start_time_s, session_sample_rate=args.session_sample_rate, min_turns=args.min_turns, time_scale=args.time_scale, concurrency_limit=args.concurrency_limit, timeout_s=args.timeout_s, request_timeout_s=args.request_timeout_s, stream=not args.no_stream, stream_idle_timeout_s=args.stream_idle_timeout_s, kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens, kvcache_admission_mode=args.kvcache_admission_mode, kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens, kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens, kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id, kvcache_seed_only_multiturn_sessions=( args.kvcache_seed_only_multiturn_sessions ), kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy, kvcache_seed_max_inflight_decode=( None if args.kvcache_seed_max_inflight_decode < 0 else args.kvcache_seed_max_inflight_decode ), kvcache_seed_max_decode_transfer_queue_reqs=( None if args.kvcache_seed_max_decode_transfer_queue_reqs is None or args.kvcache_seed_max_decode_transfer_queue_reqs < 0 else args.kvcache_seed_max_decode_transfer_queue_reqs ), kvcache_direct_max_decode_transfer_queue_reqs=( None if args.kvcache_direct_max_decode_transfer_queue_reqs is None or args.kvcache_direct_max_decode_transfer_queue_reqs < 0 else args.kvcache_direct_max_decode_transfer_queue_reqs ), kvcache_prefill_priority_eviction=( args.kvcache_prefill_priority_eviction ), kvcache_prefill_direct_priority=( args.kvcache_prefill_direct_priority ), kvcache_prefill_normal_priority=( args.kvcache_prefill_normal_priority ), pool_poll_interval_s=args.pool_poll_interval_s, pool_poll_include_sessions=not args.pool_poll_no_sessions, enable_backpressure=args.enable_backpressure, backpressure_max_pause_s=args.backpressure_max_pause_s, kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold, sample_profile=args.sample_profile, min_initial_input_tokens=args.min_initial_input_tokens, max_initial_input_tokens=args.max_initial_input_tokens, max_append_input_tokens=args.max_append_input_tokens, max_output_tokens=args.max_output_tokens, min_overlap_ratio=args.min_overlap_ratio, launch_stack=True, ) ) print( f"benchmark artifacts written under {artifacts.run_dir}; " f"metrics={artifacts.metrics_path} summary={artifacts.summary_path}" ) return raise AssertionError(f"Unhandled command: {args.command}") def _add_topology_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--model-path", default="~/models/Qwen/Qwen3-Coder-30B-A3B-Instruct", ) parser.add_argument("--prefill-workers", type=int, default=1) parser.add_argument("--decode-workers", type=int, default=1) parser.add_argument("--direct-workers", type=int, default=0) parser.add_argument("--prefill-tp-size", type=int, default=1) parser.add_argument("--decode-tp-size", type=int, default=1) parser.add_argument("--direct-tp-size", type=int, default=1) parser.add_argument("--gpu-budget", type=int, default=8) parser.add_argument( "--prefill-gpu-ids", default=None, help="Comma-separated GPU IDs for prefill workers, e.g. 3,4", ) parser.add_argument( "--decode-gpu-ids", default=None, help="Comma-separated GPU IDs for decode workers, e.g. 5", ) parser.add_argument( "--direct-gpu-ids", default=None, help="Comma-separated GPU IDs for direct workers, e.g. 6", ) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--router-port", type=int, default=8000) parser.add_argument("--prefill-port-base", type=int, default=30000) parser.add_argument("--decode-port-base", type=int, default=31000) parser.add_argument("--direct-port-base", type=int, default=32000) parser.add_argument("--bootstrap-port-base", type=int, default=8998) parser.add_argument("--transfer-backend", default="nixl") parser.add_argument( "--force-rdma", action="store_true", help=( "Force real RDMA transport for PD KV transfer. " "Currently this requires Mooncake plus --ib-device." ), ) parser.add_argument("--ib-device", default=None) parser.add_argument( "--no-trust-remote-code", action="store_true", ) def _topology_from_args(args: argparse.Namespace): transfer_backend = args.transfer_backend if args.force_rdma: transfer_backend = "mooncake" return build_single_node_topology( model_path=str(Path(args.model_path).expanduser()), prefill_worker_count=args.prefill_workers, decode_worker_count=args.decode_workers, direct_worker_count=args.direct_workers, prefill_tp_size=args.prefill_tp_size, decode_tp_size=args.decode_tp_size, direct_tp_size=args.direct_tp_size, prefill_gpu_ids=_parse_gpu_id_list(args.prefill_gpu_ids), decode_gpu_ids=_parse_gpu_id_list(args.decode_gpu_ids), direct_gpu_ids=_parse_gpu_id_list(args.direct_gpu_ids), total_gpu_budget=args.gpu_budget, host=args.host, router_port=args.router_port, prefill_port_base=args.prefill_port_base, decode_port_base=args.decode_port_base, direct_port_base=args.direct_port_base, bootstrap_port_base=args.bootstrap_port_base, transfer_backend=transfer_backend, force_rdma=args.force_rdma, trust_remote_code=not args.no_trust_remote_code, ib_device=args.ib_device, prefill_extra_server_args=("--disable-overlap-schedule",), decode_extra_server_args=("--disable-overlap-schedule",), direct_extra_server_args=("--enable-streaming-session",), ) if __name__ == "__main__": main()