gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
c9d350b37217294000af3ee41dbcdfc9481bc08a
agentic-pd-hybrid/src/agentic_pd_hybrid/cli.py
kzlin c9d350b372 docs: KVC v1-v4 debug journey + raise session soft_cap to 16 
Document the iterative debugging from v1 (broken KVC) through v4
(routing fixed + session cap raised), with code-level analysis of
the two main bugs encountered:

1. v2 root cause (mis-diagnosed previously as `allow_local_prefill`):
   `--policy default` for KVC mechanism caused replay's round-robin
   policy and the PD router's round-robin to diverge, sending requests
   with `session_params` to a D worker that did not have the session
   open. Resulted in 56-61% truncation with finish_reason
   "session id X does not exist".
   Fix: use `--policy kv-aware` (sweep_tp1_v3_kvaware.sh) so replay
   emits `x-smg-target-worker` and PD router uses consistent_hashing.

2. v3 new bottleneck: `pd-router-fallback-large-append-session-cap`
   dominated 52-65% of requests. Root cause was hardcoded
   `min(4, ...)` in `_decode_session_soft_cap`. With 7 D workers x 4
   sessions = 28 slots for 52 trace sessions, ~24 sessions starved
   permanently (bimodal direct-to-D rate of 0% or 99%).
   Fix: raise the cap to 16 (replay.py).

Also includes the v3 finding that direct-to-d-session path P50=0.495s
and TTFT P50=0.043s already beats the 8-way DP baseline (0.65s/0.093s)
- the KVC core mechanism works when fallback paths are avoided.

Files:
- docs/KVC_DEBUG_JOURNEY_V1_TO_V4.md: full journey + code location index
- docs/SWEBENCH_EXPERIMENT_{PROGRESS,RESULTS}.md: prior session notes
- scripts/sweep_tp1_v{2,3,4}*.sh: experiment driver scripts
- src/agentic_pd_hybrid/replay.py: cap 4 -> 16, audit fields
- src/agentic_pd_hybrid/pd_router.py: strip session_params from prefill
- src/agentic_pd_hybrid/metrics.py: truncated_request_count

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-28 21:10:41 +08:00

765 lines
30 KiB
Python
Raw Blame History

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)
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(
"--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,
)
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
),
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,
direct_extra_server_args=("--enable-streaming-session",),
)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink