from __future__ import annotations import argparse from dataclasses import dataclass from typing import Sequence def parse_device_ids(value: str | None) -> list[int]: if value is None or value.strip() == "": return [] return [int(item.strip()) for item in value.split(",") if item.strip()] @dataclass class ServerArgs: model: str host: str = "0.0.0.0" port: int = 8000 log_level: str = "info" zmq_command_addr: str = "" zmq_event_addr: str = "" model_name: str = "dream" decoding_strategy: str = "d2f" sampling_mode: str = "naive" tensor_parallel_size: int = 1 data_parallel_size: int = 1 master_addr: str = "localhost" master_port: int = 2333 distributed_timeout_seconds: int = 600 device_ids: list[int] | None = None block_size: int = 32 buffer_size: int = 4 page_size: int = 32 max_num_batched_tokens: int = 4096 max_num_reqs: int = 128 max_model_len: int = 2048 enable_prefill_cudagraph: bool = True enable_full_static_runner: bool = True prefill_cudagraph_max_len: int = 0 enable_torch_compile: bool = True enable_cudagraph_torch_compile: bool = False torch_compile_mode: str = "reduce-overhead" auto_max_nfe_warmup_steps: int = 8 auto_max_nfe_tpf_floor: float = 1.0 gpu_memory_utilization: float = 0.9 enforce_eager: bool = False attn_impl: str = "triton" enable_prefix_caching: bool = True kv_cache_layout: str = "unified" moe_dispatcher_backend: str = "standard" moe_gemm_impl: str = "triton" deepep_mode: str = "auto" deepep_num_max_dispatch_tokens_per_rank: int = 256 add_block_threshold: float | None = None semi_complete_threshold: float | None = None accept_threshold: float | None = None remask_threshold: float | None = None token_stability_threshold: float | None = None use_lora: bool = False lora_path: str = "" pre_merge_lora: bool = False def engine_kwargs(self) -> dict: return { "model_name": self.model_name, "decoding_strategy": self.decoding_strategy, "sampling_mode": self.sampling_mode, "tensor_parallel_size": self.tensor_parallel_size, "data_parallel_size": self.data_parallel_size, "master_addr": self.master_addr, "master_port": self.master_port, "distributed_timeout_seconds": self.distributed_timeout_seconds, "device_ids": self.device_ids or [], "block_size": self.block_size, "buffer_size": self.buffer_size, "page_size": self.page_size, "max_num_batched_tokens": self.max_num_batched_tokens, "max_num_reqs": self.max_num_reqs, "max_model_len": self.max_model_len, "enable_prefill_cudagraph": self.enable_prefill_cudagraph, "enable_full_static_runner": self.enable_full_static_runner, "prefill_cudagraph_max_len": self.prefill_cudagraph_max_len, "enable_torch_compile": self.enable_torch_compile, "enable_cudagraph_torch_compile": self.enable_cudagraph_torch_compile, "torch_compile_mode": self.torch_compile_mode, "auto_max_nfe_warmup_steps": self.auto_max_nfe_warmup_steps, "auto_max_nfe_tpf_floor": self.auto_max_nfe_tpf_floor, "gpu_memory_utilization": self.gpu_memory_utilization, "enforce_eager": self.enforce_eager, "attn_impl": self.attn_impl, "enable_prefix_caching": self.enable_prefix_caching, "kv_cache_layout": self.kv_cache_layout, "moe_dispatcher_backend": self.moe_dispatcher_backend, "moe_gemm_impl": self.moe_gemm_impl, "deepep_mode": self.deepep_mode, "deepep_num_max_dispatch_tokens_per_rank": self.deepep_num_max_dispatch_tokens_per_rank, "decoding_thresholds": { "add_block_threshold": 0.1 if self.add_block_threshold is None else self.add_block_threshold, "semi_complete_threshold": 0.9 if self.semi_complete_threshold is None else self.semi_complete_threshold, "accept_threshold": 0.9 if self.accept_threshold is None else self.accept_threshold, "remask_threshold": 0.4 if self.remask_threshold is None else self.remask_threshold, "token_stability_threshold": 0.0 if self.token_stability_threshold is None else self.token_stability_threshold, }, "use_lora": self.use_lora, "lora_path": self.lora_path, "pre_merge_lora": self.pre_merge_lora, } def build_arg_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Launch a Diffulex HTTP server") parser.add_argument("--model", required=True, help="Path to the local model directory") parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) parser.add_argument("--log-level", default="info") parser.add_argument("--zmq-command-addr", default="", help="Frontend-to-backend ZMQ PUSH/PULL address") parser.add_argument("--zmq-event-addr", default="", help="Backend-to-frontend ZMQ PUSH/PULL address") parser.add_argument("--model-name", default="dream") parser.add_argument("--decoding-strategy", default="d2f") parser.add_argument("--sampling-mode", default="naive", choices=["naive", "edit"]) parser.add_argument("--tensor-parallel-size", type=int, default=1) parser.add_argument("--data-parallel-size", type=int, default=1) parser.add_argument("--master-addr", default="localhost") parser.add_argument("--master-port", type=int, default=2333) parser.add_argument("--distributed-timeout-seconds", type=int, default=600) parser.add_argument("--device-ids", default="", help="Comma-separated logical CUDA device ids") parser.add_argument("--block-size", type=int, default=32) parser.add_argument("--buffer-size", type=int, default=4) parser.add_argument("--page-size", type=int, default=32) parser.add_argument("--max-num-batched-tokens", type=int, default=4096) parser.add_argument("--max-num-reqs", type=int, default=128) parser.add_argument("--max-model-len", type=int, default=2048) parser.add_argument("--disable-prefill-cudagraph", action="store_true") parser.add_argument("--disable-full-static-runner", action="store_true") parser.add_argument("--prefill-cudagraph-max-len", type=int, default=0) parser.add_argument("--disable-torch-compile", action="store_true") parser.add_argument("--enable-cudagraph-torch-compile", action="store_true") parser.add_argument("--torch-compile-mode", default="reduce-overhead") parser.add_argument("--auto-max-nfe-warmup-steps", type=int, default=8) parser.add_argument("--auto-max-nfe-tpf-floor", type=float, default=1.0) parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) parser.add_argument("--enforce-eager", action="store_true") parser.add_argument("--attn-impl", default="triton", choices=["triton", "naive"]) parser.add_argument("--disable-prefix-caching", action="store_true") parser.add_argument("--kv-cache-layout", default="unified", choices=["unified", "distinct"]) parser.add_argument("--moe-dispatcher-backend", default="standard", choices=["standard", "naive", "deepep"]) parser.add_argument("--moe-gemm-impl", default="triton", choices=["triton", "vllm", "vllm_modular", "naive"]) parser.add_argument("--deepep-mode", default="auto", choices=["normal", "low_latency", "auto"]) parser.add_argument("--deepep-num-max-dispatch-tokens-per-rank", type=int, default=256) parser.add_argument("--add-block-threshold", type=float, default=None) parser.add_argument("--semi-complete-threshold", type=float, default=None) parser.add_argument("--accept-threshold", type=float, default=None) parser.add_argument("--remask-threshold", type=float, default=None) parser.add_argument("--token-stability-threshold", type=float, default=None) parser.add_argument("--use-lora", action="store_true") parser.add_argument("--lora-path", default="") parser.add_argument("--pre-merge-lora", action="store_true") return parser def parse_args(argv: Sequence[str] | None = None) -> ServerArgs: ns = build_arg_parser().parse_args(argv) return ServerArgs( model=ns.model, host=ns.host, port=ns.port, log_level=ns.log_level, zmq_command_addr=ns.zmq_command_addr, zmq_event_addr=ns.zmq_event_addr, model_name=ns.model_name, decoding_strategy=ns.decoding_strategy, sampling_mode=ns.sampling_mode, tensor_parallel_size=ns.tensor_parallel_size, data_parallel_size=ns.data_parallel_size, master_addr=ns.master_addr, master_port=ns.master_port, distributed_timeout_seconds=ns.distributed_timeout_seconds, device_ids=parse_device_ids(ns.device_ids), block_size=ns.block_size, buffer_size=ns.buffer_size, page_size=ns.page_size, max_num_batched_tokens=ns.max_num_batched_tokens, max_num_reqs=ns.max_num_reqs, max_model_len=ns.max_model_len, enable_prefill_cudagraph=not ns.disable_prefill_cudagraph, enable_full_static_runner=not ns.disable_full_static_runner, prefill_cudagraph_max_len=ns.prefill_cudagraph_max_len, enable_torch_compile=not ns.disable_torch_compile, enable_cudagraph_torch_compile=ns.enable_cudagraph_torch_compile, torch_compile_mode=ns.torch_compile_mode, auto_max_nfe_warmup_steps=ns.auto_max_nfe_warmup_steps, auto_max_nfe_tpf_floor=ns.auto_max_nfe_tpf_floor, gpu_memory_utilization=ns.gpu_memory_utilization, enforce_eager=ns.enforce_eager, attn_impl=ns.attn_impl, enable_prefix_caching=not ns.disable_prefix_caching, kv_cache_layout=ns.kv_cache_layout, moe_dispatcher_backend=ns.moe_dispatcher_backend, moe_gemm_impl=ns.moe_gemm_impl, deepep_mode=ns.deepep_mode, deepep_num_max_dispatch_tokens_per_rank=ns.deepep_num_max_dispatch_tokens_per_rank, add_block_threshold=ns.add_block_threshold, semi_complete_threshold=ns.semi_complete_threshold, accept_threshold=ns.accept_threshold, remask_threshold=ns.remask_threshold, token_stability_threshold=ns.token_stability_threshold, use_lora=ns.use_lora, lora_path=ns.lora_path, pre_merge_lora=ns.pre_merge_lora, )