lightseekorg--tokenspeed/python/tokenspeed/runtime/utils/server_args.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""The arguments of the server."""

import argparse
import dataclasses
import json
import os
import random
from typing import Literal

from tokenspeed_kernel.ops.attention.triton.linear.chunk_delta_h import (
    CHUNK_SIZE as FLA_CHUNK_SIZE,
)
from tokenspeed_kernel.platform import current_platform

from tokenspeed.runtime.distributed.mapping import Mapping, _resolve_parallelism_sizes
from tokenspeed.runtime.utils import (
    get_amdgpu_memory_capacity,
    get_colorful_logger,
    get_nvgpu_memory_capacity,
    is_valid_ipv6_address,
    maybe_model_redirect,
    nullable_str,
)
from tokenspeed.runtime.utils.network import is_port_available

logger = get_colorful_logger(__name__)

ENABLE_CP = os.environ.get("ENABLE_CP", "false").lower() in ("true", "1")


def str_to_bool(value: str | bool) -> bool:
    if isinstance(value, bool):
        return value
    normalized = value.lower()
    if normalized in {"1", "true", "yes", "on"}:
        return True
    if normalized in {"0", "false", "no", "off"}:
        return False
    raise argparse.ArgumentTypeError(f"invalid boolean value: {value!r}")


@dataclasses.dataclass
class ServerArgs:
    # Model and tokenizer
    model: str
    tokenizer: str | None = None
    tokenizer_mode: str = "auto"
    skip_tokenizer_init: bool = False
    load_format: str = "auto"
    trust_remote_code: bool = True
    dtype: str = "auto"
    kv_cache_dtype: str = "auto"
    kv_cache_quant_method: str = "none"
    quantization: str | None = None
    quantization_param_path: nullable_str = None
    max_model_len: int | None = None
    device: str = "cuda"
    served_model_name: str | None = None
    revision: str | None = None
    language_model_only: bool = False

    # Port for the HTTP server
    host: str = "127.0.0.1"
    port: int = 8000

    # Memory and scheduling
    gpu_memory_utilization: float | None = None
    max_num_seqs: int | None = None
    max_total_tokens: int | None = None
    chunked_prefill_size: int | None = None
    max_prefill_tokens: int = 8192
    enable_mixed_batch: bool = False
    block_size: int = 64
    # special kv cache
    mamba_ssm_dtype: str = "float32"
    mamba_track_interval: int = 256
    max_mamba_cache_size: int | None = None
    mamba_full_memory_ratio: float = 0.9
    enable_mamba_l2: bool = False
    mamba_l2_host_slots: int = 0
    mamba_l2_ratio: float = 2.0
    mamba_l2_layout: str = "layer_first"
    mamba_l2_io_backend: str = "kernel"
    mamba_l2_host_gb: int = 0

    # Other runtime options
    stream_interval: int = 1
    stream_output: bool = False
    # Inline detokenization is the only supported path and is intentionally
    # not configurable from the CLI.
    enable_inline_detokenizer: bool = True
    seed: int | None = None
    distributed_timeout_seconds: int | None = None
    download_dir: str | None = None
    # Used for customizing extensible models
    ext_yaml: str | None = None
    base_gpu_id: int = 0
    gpu_id_step: int = 1

    # Logging
    log_level: str = "info"
    log_level_http: str | None = None
    enable_log_requests: bool = False
    log_requests_level: int = 0
    enable_log_request_stats: bool = False
    enable_metrics: bool = False
    decode_log_interval: int = 40
    metrics_reporters: list[str] | None = None
    app_key: str | None = None

    # API related
    api_key: str | None = None
    enable_cache_report: bool = False
    kv_events_config: str | None = None

    # Data parallelism
    data_parallel_size: int | None = None
    load_balance_method: str = "shortest_queue"
    load_watch_interval: float = 0.02

    # Expert parallelism
    ep_size: int = 1
    init_expert_location: str = "trivial"
    ep_num_redundant_experts: int = 0
    ep_dispatch_algorithm: (
        Literal[
            "static",
            "dynamic",
            "fake",
            "static_with_zero_expert",
            "dynamic_with_zero_expert",
        ]
        | None
    ) = None
    eplb_algorithm: str = "auto"
    expert_distribution_recorder_mode: (
        Literal["stat", "stat_approx", "per_pass", "per_token"] | None
    ) = None
    expert_distribution_recorder_buffer_size: int | None = None
    enable_expert_distribution_metrics: bool = False
    enable_eplb: bool = False

    # MoE backend
    moe_backend: str = "auto"
    draft_moe_backend: str | None = None
    all2all_backend: str = "none"
    deepep_mode: Literal["auto", "normal", "low_latency"] = "auto"
    disable_flashinfer_cutlass_moe_fp4_allgather: bool = False

    # KVStore
    enable_kvstore: bool = False
    kvstore_ratio: float = 2.0
    kvstore_size: int = 0
    kvstore_io_backend: str = "kernel"
    kvstore_mem_layout: str = "layer_first"
    kvstore_storage_backend: str | None = None
    kvstore_storage_backend_extra_config: str | None = None
    enable_mla_l1_5_cache: bool = False

    # Multi-node distributed serving
    dist_init_addr: str | None = None
    nnodes: int = 1
    node_rank: int = 0

    # Hugging Face model config overrides in JSON
    hf_overrides: str = "{}"
    preferred_sampling_params: str | None = None

    # Kernel backend
    attention_backend: str | None = None
    drafter_attention_backend: str | None = None
    sampling_backend: str | None = None
    dp_sampling: bool = False
    dp_sampling_min_bs: int | None = None
    attention_use_fp4_indexer_cache: bool | None = None
    use_trtllm_ragged_deepseek_prefill: bool | None = None

    # DeepSeek V4
    deepseek_v4_mega_moe_max_num_tokens: int = 0
    deepseek_v4_indexer_prefill_max_logits_mb: int = 512
    deepseek_v4_prefill_chunk_size: int = 4

    # Grammar backend
    grammar_backend: str = "none"
    # Used by ``input_processor`` to defer json_schema grammars past the
    # model's reasoning channel.
    reasoning_parser: str | None = None
    grammar_compile_timeout_secs: float = 30.0
    grammar_compile_max_retries: int = 2
    disable_any_whitespace: bool = False
    # Force the synchronous eager grammar fallback even on CUDA. Useful
    # for parity-testing against the captured-grammar path (output should
    # match; throughput will be lower since the sync stalls every step).
    disable_capturable_grammar: bool = False

    # Speculative decoding
    draft_model_path_use_base: bool | None = False
    speculative_config: str | None = None
    speculative_algorithm: str | None = None
    speculative_draft_model_path: str | None = None
    speculative_draft_model_quantization: str | None = "unquant"
    speculative_num_steps: int = 3
    speculative_eagle_topk: int = 1
    speculative_num_draft_tokens: int | None = None
    eagle3_layers_to_capture: str | None = None
    # Logprob support flags — all OFF by default. Enabling extends the
    # captured CUDA-graph footprint; requests asking for logprobs on a
    # server started without the matching flag will receive empty logprobs.
    enable_output_logprobs: bool = False

    # Runtime options
    disable_pdl: bool = False
    enable_prefix_caching: bool = True
    disable_kvstore: bool = False
    enforce_eager: bool = False
    disable_cuda_graph_padding: bool = False
    enable_cudagraph_gc: bool = False
    enable_nccl_nvls: bool = False
    enable_symm_mem: bool = False
    disable_custom_all_reduce: bool = False
    disable_overlap_schedule: bool = False
    disable_tf32: bool = False
    force_deterministic_rsag: bool = False
    disable_sampling_tp_sync: bool = False
    low_latency_max_num_tokens_per_gpu: int = 256
    max_cudagraph_capture_size: int | None = None
    disable_prefill_graph: bool | None = False
    # Breakable prefill graph bucket cap: None = auto min(2048, chunk); 0 disables.
    prefill_graph_max_tokens: int | None = None
    # Explicit prefill bucket list; unset = the relative-stride ladder (see get_prefill_token_buckets).
    prefill_graph_capture_sizes: list[int] | None = None
    cudagraph_capture_sizes: list[int] | None = None
    enable_nan_detection: bool = False
    enable_nvtx: bool = False
    enable_p2p_check: bool = False
    triton_attention_reduce_in_fp32: bool = False
    delete_ckpt_after_loading: bool = False
    weight_loader_prefetch_checkpoints: bool = False
    weight_loader_prefetch_num_threads: int = 4
    enable_memory_saver: bool = False
    enable_custom_logit_processor: bool = False
    mla_disable_ragged: bool = False
    warmups: str | None = None

    # parallel strategy
    nprocs_per_node: int | None = None
    world_size: int | None = None
    attn_tp_size: int | None = None
    dense_tp_size: int | None = None
    moe_tp_size: int | None = None
    mapping: Mapping | None = None

    mla_chunk_multiplier: int = 4
    mm_attention_backend: str | None = None

    # For PD/EPD disaggregation: "null", "prefill", "decode", or "encode" (vision-tower-only).
    disaggregation_mode: str = "null"
    disaggregation_bootstrap_port: int = 8998
    disaggregation_transfer_backend: str = "mooncake"
    disaggregation_ib_device: str | None = None
    disaggregation_layerwise_interval: int = 1
    pdlb_url: str | None = None

    skip_server_warmup: bool = False

    # For communication + norm fusion
    comm_fusion_max_num_tokens: int = 2048
    enable_allreduce_fusion: bool = False

    enable_expert_parallel: bool = False

    @property
    def mamba_cache_chunk_size(self) -> int:
        return max(FLA_CHUNK_SIZE, self.block_size)

    def __post_init__(self):
        self.resolve_basic_defaults()
        self.resolve_parallelism()
        self.resolve_memory_and_scheduling()
        self.resolve_kernel_backends()
        self.resolve_cache()
        self.resolve_speculative_decoding()
        self.resolve_communication()
        self.resolve_disaggregation()
        self.validate()

    def resolve_basic_defaults(self):
        self.model = maybe_model_redirect(self.model)

        if self.kv_cache_dtype == "fp8":
            self.kv_cache_dtype = "fp8_e4m3"

        self.resolve_config_aliases()

        # Set missing default values
        if self.tokenizer is None:
            self.tokenizer = self.model

        if self.served_model_name is None:
            self.served_model_name = self.model

        if self.seed is None:
            self.seed = random.randint(0, 1 << 30)

    def resolve_config_aliases(self):
        if self.use_trtllm_ragged_deepseek_prefill is not None:
            self.mla_disable_ragged = not self.use_trtllm_ragged_deepseek_prefill

        if self.speculative_config is not None:
            try:
                config = json.loads(self.speculative_config)
            except json.JSONDecodeError as exc:
                raise ValueError("--speculative-config must be valid JSON") from exc

            if not isinstance(config, dict):
                raise ValueError("--speculative-config must be a JSON object")

            method = config.get("method")
            if method is not None and self.speculative_algorithm is None:
                self.speculative_algorithm = str(method).upper()

            draft_model = config.get("model")
            if draft_model is not None and self.speculative_draft_model_path is None:
                self.speculative_draft_model_path = str(draft_model)

            num_speculative_tokens = config.get("num_speculative_tokens")
            if num_speculative_tokens is not None:
                num_speculative_tokens = int(num_speculative_tokens)
                if self.speculative_algorithm == "DFLASH":
                    if self.speculative_num_draft_tokens is None:
                        self.speculative_num_draft_tokens = num_speculative_tokens
                    self.speculative_num_steps = max(num_speculative_tokens - 1, 0)
                else:
                    self.speculative_num_steps = num_speculative_tokens

        if self.speculative_num_draft_tokens is None:
            self.speculative_num_draft_tokens = self.speculative_num_steps + 1

    def resolve_memory_and_scheduling(self):
        if current_platform().is_amd:
            gpu_mem = get_amdgpu_memory_capacity()
        elif current_platform().is_nvidia:
            gpu_mem = get_nvgpu_memory_capacity()
        else:
            # GPU memory is not known yet or no GPU is available.
            gpu_mem = None

        # Set GPU memory utilization, which depends on the tensor parallelism size.
        self._gpu_memory_utilization_defaulted = False
        if self.gpu_memory_utilization is None:
            if self.mapping.world_size >= 16:
                self.gpu_memory_utilization = 0.79
            elif self.mapping.world_size >= 8:
                self.gpu_memory_utilization = 0.81
            elif self.mapping.world_size >= 4:
                self.gpu_memory_utilization = 0.95
            elif self.mapping.world_size >= 2:
                self.gpu_memory_utilization = 0.87
            else:
                self.gpu_memory_utilization = 0.88
            self._gpu_memory_utilization_defaulted = True

        # Set the chunked prefill token budget.
        if self.chunked_prefill_size is None:
            self.chunked_prefill_size = 8192

        # Set CUDA graph max capture size.
        if self.max_cudagraph_capture_size is None:
            # Based on detailed statistics, when serving TP1/TP2 models on lower-end GPUs with HBM<25G, you can either disable CUDA graph or set max_cudagraph_capture_size to a very small value to reduce graph memory overhead, with almost no impact on performance. TP4/TP8 serving still needs CUDA graph for high performance, and 80 is enough for lower-end GPUs.
            if gpu_mem is not None and gpu_mem < 25_000:
                if self.mapping.world_size < 4:
                    self.max_cudagraph_capture_size = 8
                else:
                    self.max_cudagraph_capture_size = 80
            elif self.speculative_algorithm:
                self.max_cudagraph_capture_size = 80
            else:
                self.max_cudagraph_capture_size = 160

        # Set max number of sequences.
        if self.max_num_seqs is None:
            if self.speculative_algorithm:
                self.max_num_seqs = 80
            else:
                self.max_num_seqs = 160

    def resolve_kernel_backends(self):
        # Choose kernel backends
        # attention_backend default is NOT set here — deferred to
        # AttnInitializer.modify_args where both hardware and model arch are known.

        if self.sampling_backend is None:
            # ``flashinfer`` is the only built-in backend that respects per-request
            # ``temperature`` / ``top_p`` / ``top_k``. ``greedy`` is argmax-only
            # (see ``GreedySamplingBackend.sample``: *"sampling_info is ignored
            # for single-step (always argmax)"*) — fast for hand-tuned greedy
            # decoding but silently wrong for any serving deployment where
            # requests carry sampling params, since the model collapses into
            # repetition-mode loops within a few hundred steps. Default to the
            # sampling-respecting backend on NVIDIA where flashinfer is
            # available, fall back to greedy elsewhere; users can still opt
            # into greedy explicitly via ``--sampling-backend greedy``.
            if current_platform().is_nvidia:
                self.sampling_backend = "flashinfer"
            else:
                self.sampling_backend = "greedy"

    def resolve_parallelism(self):
        world_size = self.world_size
        nprocs_per_node = self.nprocs_per_node
        nnodes = 1 if self.nnodes is None else self.nnodes

        attn_tp_size = self.attn_tp_size
        attn_dp_size = self.data_parallel_size

        # ``ENABLE_CP`` interprets attention TP size as CP size.
        attn_cp_size = 1
        if ENABLE_CP:
            attn_cp_size, attn_tp_size = attn_tp_size, 1

        if world_size is None:
            world_size = 1
            if attn_tp_size is not None:
                world_size *= attn_tp_size
            if attn_cp_size is not None:
                world_size *= attn_cp_size
            if attn_dp_size is not None:
                world_size *= attn_dp_size
            logger.info(
                "Inferred world_size (%s) from attn_tp_size (%s) x attn_cp_size (%s) x attn_dp_size (%s)",
                world_size,
                attn_tp_size,
                attn_cp_size,
                attn_dp_size,
            )
        else:
            logger.info("Specified world_size (%s)", world_size)

        attn_tp_size, attn_cp_size, attn_dp_size = _resolve_parallelism_sizes(
            world_size, attn_tp_size, attn_cp_size, attn_dp_size
        )

        # Dense layers still default to full TP participation when no
        # dedicated dense_tp_size is provided.
        dense_tp_size = self.dense_tp_size
        if self.dense_tp_size is None:
            # dense always do tp now.
            dense_tp_size = world_size
        dense_dp_size = None

        # --enable-expert-parallel auto-sets ep_size = world_size
        if self.enable_expert_parallel and self.ep_size == 1:
            self.ep_size = world_size
            logger.info("--enable-expert-parallel: auto-setting ep_size=%s", world_size)

        # MoE parallel sizes default to consuming the full world size unless
        # the user overrides them explicitly.
        moe_ep_size = 1 if self.ep_size is None else self.ep_size
        moe_tp_size = (
            world_size // moe_ep_size if self.moe_tp_size is None else self.moe_tp_size
        )
        moe_dp_size = None

        self.mapping = Mapping(
            world_size=world_size,
            attn_tp_size=attn_tp_size,
            attn_cp_size=attn_cp_size,
            attn_dp_size=attn_dp_size,
            dense_tp_size=dense_tp_size,
            dense_dp_size=dense_dp_size,
            moe_tp_size=moe_tp_size,
            moe_ep_size=moe_ep_size,
            moe_dp_size=moe_dp_size,
            nprocs_per_node=nprocs_per_node,
            nnodes=nnodes,
            base_gpu_id=self.base_gpu_id,
            gpu_id_step=self.gpu_id_step,
        )

        # Impl constraints:
        if self.mapping.moe.has_tp and self.mapping.moe.has_ep:
            raise ValueError("MoE TP and EP cannot be both > 1")

        logger.info("Parallelism configuration:\n%s", self.mapping)

    def resolve_cache(self):
        # Handle KVStore settings.
        self._handle_kvstore()
        self.validate_cache_options()

    def resolve_speculative_decoding(self):
        # Keep drafter backend consistent with the main model unless explicitly set.
        if (
            self.speculative_algorithm is not None
            and self.drafter_attention_backend is None
        ):
            self.drafter_attention_backend = self.attention_backend

        if (
            self.speculative_algorithm == "MTP"
            and self.speculative_draft_model_path is None
        ):
            self.draft_model_path_use_base = True

        if self.draft_model_path_use_base:
            self.speculative_draft_model_path = self.model

        if self.speculative_draft_model_path == self.model:
            self.draft_model_path_use_base = True

        if self.speculative_draft_model_quantization == "unquant":
            self.speculative_draft_model_quantization = None

        if self.speculative_algorithm == "DFLASH":
            expected_steps = max(int(self.speculative_num_draft_tokens) - 1, 0)
            if self.speculative_num_steps == ServerArgs.speculative_num_steps:
                self.speculative_num_steps = expected_steps
            elif self.speculative_num_steps != expected_steps:
                raise ValueError(
                    "DFLASH requires speculative_num_steps to equal "
                    "speculative_num_draft_tokens - 1. "
                    f"Got {self.speculative_num_steps=} and "
                    f"{self.speculative_num_draft_tokens=}."
                )

        if self.eagle3_layers_to_capture is not None:
            self.eagle3_layers_to_capture = [
                int(x) for x in self.eagle3_layers_to_capture.split(",")
            ]

        # Hoist the PD-decode topk == 1 check to startup.
        if self.speculative_algorithm is not None and self.speculative_eagle_topk != 1:
            raise ValueError(
                "speculative_eagle_topk > 1 (tree spec) is not currently "
                f"supported: {self.speculative_eagle_topk=}. Only chain spec "
                "(topk=1) is wired end-to-end."
            )

    def resolve_communication(self):
        # Auto-enable allreduce fusion on supported single-node TP configurations.
        platform = current_platform()
        if (
            not self.enable_allreduce_fusion
            and (current_platform().is_hopper_plus or platform.is_amd)
            and self.mapping.nnodes == 1
            and self.mapping.has_attn_tp
            and not self.mapping.has_attn_dp
        ):
            self.enable_allreduce_fusion = True
            logger.info("Auto-enabled allreduce fusion")

        if self.mapping.attn.tp_size != self.mapping.dense.tp_size:
            self.comm_fusion_max_num_tokens = -1
            self.enable_allreduce_fusion = False
            logger.info(
                "allreduce is forbidden due to different attn_tp_size: %s and dense_tp_size: %s!",
                self.mapping.attn.tp_size,
                self.mapping.dense.tp_size,
            )

    def resolve_disaggregation(self):
        # PD disaggregation
        if self.disaggregation_mode == "prefill":
            self.enforce_eager = True
            logger.warning("CUDA graph is disabled for prefill server")
        elif self.disaggregation_mode == "decode":
            # Prefix caching stays configurable for decode servers.
            logger.info(
                "enable_prefix_caching=%r for decode server",
                self.enable_prefix_caching,
            )
        elif self.disaggregation_mode == "encode":
            # Encode server: vision tower only, no LM / KV pool / prefix cache.
            # enforce_eager left as-is (the vision tower keeps its own CUDA graph).
            if self.mapping.has_attn_dp:
                raise ValueError(
                    "disaggregation_mode=encode currently supports "
                    "data_parallel_size == 1 inside one encode server; run "
                    "multiple independent encode servers for horizontal scale."
                )
            self.enable_prefix_caching = False

        # Prefill graph disable logic is handled by AttnInitializer.modify_args
        # after the attention backend is resolved.

        if (
            self.disaggregation_mode == "prefill"
            and self.load_balance_method != "round_robin"
        ):
            if self.mapping.has_attn_dp:
                raise ValueError(
                    "Not supported when "
                    f"{self.disaggregation_mode=} {self.load_balance_method=} "
                    f"{self.mapping.attn.dp_size=}"
                )

    def _handle_kvstore(self):
        if self.disaggregation_mode in ("decode", "encode"):
            self.enable_kvstore = False
            logger.info(
                "%s instance has set enable_kvstore to False!",
                self.disaggregation_mode,
            )
        elif not self.disable_kvstore:
            self.enable_kvstore = True

        if self.kvstore_storage_backend == "mooncake":
            if self.kvstore_mem_layout == "layer_first":
                self.kvstore_mem_layout = "page_first"
                logger.warning(
                    "Mooncake storage backend does not support layer_first layout, switching to %s layout",
                    self.kvstore_mem_layout,
                )

            if self.kvstore_io_backend == "direct":
                self.kvstore_io_backend = "kernel"
                logger.warning(
                    "Mooncake storage backend uses page_first layout, which requires kernel io backend"
                )

    def validate_cache_options(self):
        if self.enable_kvstore and not self.enable_prefix_caching:
            raise ValueError(
                "KVStore and disabled prefix caching are mutually exclusive "
                "and cannot be used at the same time. Please use only one of them."
            )

    def validate(self):
        if (
            self.max_num_seqs is not None
            and self.max_num_seqs < self.mapping.attn.dp_size
        ):
            raise ValueError(
                f"max_num_seqs must be >= attn_dp_size: {self.max_num_seqs=} < {self.mapping.attn.dp_size=}"
            )

        if self.mapping.has_attn_cp and self.max_num_seqs > 1:
            raise ValueError("CP attention is enabled but max_num_seqs > 1")

        if self.mapping.has_attn_dp:
            if self.chunked_prefill_size > self.max_prefill_tokens:
                raise ValueError(
                    f"chunked_prefill_size must be <= max_prefill_tokens: {self.chunked_prefill_size=} > {self.max_prefill_tokens=}"
                )

        if self.deepseek_v4_prefill_chunk_size <= 0:
            raise ValueError("deepseek_v4_prefill_chunk_size must be positive")

        if self.enable_eplb and (self.expert_distribution_recorder_mode is None):
            self.expert_distribution_recorder_mode = "stat"
            logger.info(
                "EPLB is enabled. The expert_distribution_recorder_mode is automatically set."
            )

        if (self.enable_eplb or (self.init_expert_location is not None)) and (
            self.ep_dispatch_algorithm is None
        ):
            self.ep_dispatch_algorithm = "static"
            logger.info(
                "EPLB is enabled or init_expert_location is provided. ep_dispatch_algorithm is configured."
            )

        from tokenspeed.runtime.utils.env import envs

        envs.TOKENSPEED_MAMBA_SSM_DTYPE.set(self.mamba_ssm_dtype)
        if not self.disable_pdl:
            os.environ.setdefault("TORCHINDUCTOR_ENABLE_PDL", "1")
            # Enable PDL for fused attention kernels.
            os.environ.setdefault("TRTLLM_ENABLE_PDL", "1")
        os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")

    @staticmethod
    def add_cli_args(parser: argparse.ArgumentParser):
        parser.allow_abbrev = False

        # Model and port args
        parser.add_argument(
            "model_path",
            nargs="?",
            metavar="model",
            default=None,
            help="The model name or path (positional argument). "
            "Equivalent to --model.",
        )
        parser.add_argument(
            "--model",
            "--model-path",
            metavar="MODEL",
            type=str,
            default=None,
            help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
        )
        parser.add_argument(
            "--tokenizer",
            metavar="TOKENIZER",
            type=str,
            default=ServerArgs.tokenizer,
            help="The path of the tokenizer.",
        )
        parser.add_argument(
            "--host", type=str, default=ServerArgs.host, help="The host of the server."
        )
        parser.add_argument(
            "--port", type=int, default=ServerArgs.port, help="The port of the server."
        )
        parser.add_argument(
            "--tokenizer-mode",
            type=str,
            default=ServerArgs.tokenizer_mode,
            choices=["auto", "slow", "deepseek_v4"],
            help="Tokenizer mode. 'auto' will use the fast "
            "tokenizer and model-specific tokenizer hooks if available, "
            "'slow' will always use the slow tokenizer.",
        )
        parser.add_argument(
            "--skip-tokenizer-init",
            action=argparse.BooleanOptionalAction,
            default=ServerArgs.skip_tokenizer_init,
            help="If set, skip init tokenizer and pass input_ids in generate request",
        )
        parser.add_argument(
            "--language-model-only",
            action="store_true",
            default=ServerArgs.language_model_only,
            help="Skip vision/audio encoders on a multimodal checkpoint and "
            "run text-only. Multimodal requests are rejected.",
        )
        parser.add_argument("--ext-yaml", type=str, default=None)
        parser.add_argument(
            "--load-format",
            type=str,
            default=ServerArgs.load_format,
            choices=[
                "auto",
                "pt",
                "safetensors",
                "npcache",
                "dummy",
                "extensible",
            ],
            help="The format of the model weights to load. "
            '"auto" will try to load the weights in the safetensors format '
            "and fall back to the pytorch bin format if safetensors format "
            "is not available. "
            '"pt" will load the weights in the pytorch bin format. '
            '"safetensors" will load the weights in the safetensors format. '
            '"npcache" will load the weights in pytorch format and store '
            "a numpy cache to speed up the loading. "
            '"dummy" will initialize the weights with random values.',
        )
        parser.add_argument(
            "--trust-remote-code",
            action=argparse.BooleanOptionalAction,
            default=False,
            help="Whether or not to allow for custom models defined on the Hub in their own modeling files.",
        )
        parser.add_argument(
            "--dtype",
            type=str,
            default=ServerArgs.dtype,
            choices=["auto", "half", "float16", "bfloat16", "float", "float32"],
            help="Data type for model weights and activations.\n\n"
            '* "auto" will use FP16 precision for FP32 and FP16 models, and '
            "BF16 precision for BF16 models.\n"
            '* "half" for FP16. Recommended for AWQ quantization.\n'
            '* "float16" is the same as "half".\n'
            '* "bfloat16" for a balance between precision and range.\n'
            '* "float" is shorthand for FP32 precision.\n'
            '* "float32" for FP32 precision.',
        )
        parser.add_argument(
            "--kv-cache-dtype",
            type=str,
            default=ServerArgs.kv_cache_dtype,
            choices=["auto", "fp8", "fp8_e4m3"],
            help='Data type for kv cache storage. "auto" will use model data type. "fp8" is an alias for "fp8_e4m3".',
        )
        parser.add_argument(
            "--kv-cache-quant-method",
            type=str,
            default=ServerArgs.kv_cache_quant_method,
            choices=["none", "per_token_head"],
            help="kv cache quant method",
        )
        parser.add_argument(
            "--quantization",
            type=str,
            default=ServerArgs.quantization,
            choices=[
                "fp8",
                "mxfp4",
                "nvfp4",
                "w8a8_fp8",
                "compressed-tensors",
            ],
            help="The quantization method.",
        )
        parser.add_argument(
            "--quantization-param-path",
            type=nullable_str,
            default=None,
            help="Path to the JSON file containing the KV cache "
            "scaling factors. This should generally be supplied, when "
            "KV cache dtype is FP8. Otherwise, KV cache scaling factors "
            "default to 1.0, which may cause accuracy issues. ",
        )
        parser.add_argument(
            "--max-model-len",
            metavar="MAX_MODEL_LEN",
            type=int,
            default=ServerArgs.max_model_len,
            help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead).",
        )
        parser.add_argument(
            "--device",
            type=str,
            default="cuda",
            choices=["cuda"],
            help="The device type.",
        )
        parser.add_argument(
            "--served-model-name",
            type=str,
            default=ServerArgs.served_model_name,
            help="Override the model name returned by the v1/models endpoint in OpenAI API server.",
        )
        parser.add_argument(
            "--revision",
            type=str,
            default=None,
            help="The specific model version to use. It can be a branch "
            "name, a tag name, or a commit id. If unspecified, will use "
            "the default version.",
        )
        # Memory and scheduling
        parser.add_argument(
            "--gpu-memory-utilization",
            metavar="GPU_MEMORY_UTILIZATION",
            type=float,
            default=ServerArgs.gpu_memory_utilization,
            help="The fraction of GPU memory to use for model weights and KV cache. Use a smaller value if you see out-of-memory errors.",
        )
        parser.add_argument(
            "--max-num-seqs",
            metavar="MAX_NUM_SEQS",
            type=int,
            default=ServerArgs.max_num_seqs,
            help="Maximum number of sequences to process concurrently.",
        )
        parser.add_argument(
            "--max-total-tokens",
            type=int,
            default=ServerArgs.max_total_tokens,
            help="The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. "
            "This overrides the automatically calculated token pool size.",
        )
        parser.add_argument(
            "--chunked-prefill-size",
            metavar="CHUNKED_PREFILL_SIZE",
            type=int,
            default=ServerArgs.chunked_prefill_size,
            help="Maximum number of tokens the scheduler may issue in a single iteration. Setting this to -1 disables chunked prefill.",
        )
        parser.add_argument(
            "--enable-mixed-batch",
            action="store_true",
            dest="enable_mixed_batch",
            default=ServerArgs.enable_mixed_batch,
            help="Allow the scheduler to issue prefill and decode requests in the same iteration.",
        )
        parser.add_argument(
            "--block-size",
            metavar="BLOCK_SIZE",
            type=int,
            default=ServerArgs.block_size,
        )

        # KVStore
        parser.add_argument(
            "--disable-kvstore",
            action="store_true",
            help="Disable KVStore",
        )
        parser.add_argument(
            "--kvstore-ratio",
            type=float,
            default=ServerArgs.kvstore_ratio,
            help="The ratio of the size of the KVStore host memory pool to the size of the device pool.",
        )
        parser.add_argument(
            "--kvstore-size",
            type=int,
            default=ServerArgs.kvstore_size,
            help="The size of the KVStore host memory pool in gigabytes, which will override kvstore_ratio if set.",
        )
        parser.add_argument(
            "--kvstore-io-backend",
            type=str,
            choices=["direct", "kernel"],
            default=ServerArgs.kvstore_io_backend,
            help="The IO backend for KVStore transfer between CPU and GPU.",
        )
        parser.add_argument(
            "--kvstore-mem-layout",
            type=str,
            choices=[
                "layer_first",
                "page_first",
                "page_head",
            ],
            default=ServerArgs.kvstore_mem_layout,
            help="The layout of the KVStore host memory pool.",
        )
        parser.add_argument(
            "--kvstore-storage-backend",
            type=str,
            choices=["mooncake"],
            default=ServerArgs.kvstore_storage_backend,
            help="The storage backend for KVStore. "
            "Built-in backends: mooncake. "
            "For dynamic backend, use --kvstore-storage-backend-extra-config to specify: "
            "backend_name (custom name), module_path (Python module path), class_name (backend class name).",
        )
        parser.add_argument(
            "--kvstore-storage-backend-extra-config",
            type=str,
            default=ServerArgs.kvstore_storage_backend_extra_config,
            help="A dictionary in JSON string format containing extra configuration for the storage backend.",
        )
        parser.add_argument(
            "--enable-mla-l1-5-cache",
            action="store_true",
            help="Enable MLA L1.5 cache in disaggregation paths.",
        )
        # Mamba Cache
        parser.add_argument(
            "--mamba-ssm-dtype",
            type=str,
            default=ServerArgs.mamba_ssm_dtype,
            choices=["float32", "bfloat16"],
            help="It is used to tune mamba ssm dtype",
        )
        parser.add_argument(
            "--mamba-track-interval",
            type=int,
            default=ServerArgs.mamba_track_interval,
            help="The interval to track the mamba state during decode.",
        )
        parser.add_argument(
            "--max-mamba-cache-size",
            type=int,
            default=ServerArgs.max_mamba_cache_size,
            help="The maximum number of Mamba cache chunks. If unset, the pool size is profiled from available memory.",
        )
        parser.add_argument(
            "--mamba-full-memory-ratio",
            type=float,
            default=ServerArgs.mamba_full_memory_ratio,
            help="Memory ratio used to split cache budget between Mamba state chunks and full-attention KV cache.",
        )
        parser.add_argument(
            "--enable-mamba-l2",
            action="store_true",
            help="Enable host-memory L2 cache for Mamba state slots.",
        )
        parser.add_argument(
            "--mamba-l2-host-slots",
            type=int,
            default=ServerArgs.mamba_l2_host_slots,
            help="Number of host Mamba L2 slots. If 0, derive from --mamba-l2-host-gb or --mamba-l2-ratio.",
        )
        parser.add_argument(
            "--mamba-l2-ratio",
            type=float,
            default=ServerArgs.mamba_l2_ratio,
            help="Mamba host L2 slot ratio relative to device Mamba slots when host slots are not explicit.",
        )
        parser.add_argument(
            "--mamba-l2-layout",
            type=str,
            choices=["layer_first"],
            default=ServerArgs.mamba_l2_layout,
            help="Mamba host L2 memory layout.",
        )
        parser.add_argument(
            "--mamba-l2-io-backend",
            type=str,
            choices=["direct", "kernel"],
            default=ServerArgs.mamba_l2_io_backend,
            help="IO backend for Mamba L2 host/device transfers.",
        )
        parser.add_argument(
            "--mamba-l2-host-gb",
            type=int,
            default=ServerArgs.mamba_l2_host_gb,
            help="Mamba L2 host memory budget in GiB. Overrides --mamba-l2-ratio when host slots are not explicit.",
        )

        parser.add_argument(
            "--max-prefill-tokens",
            metavar="MAX_PREFILL_TOKENS",
            type=int,
            default=ServerArgs.max_prefill_tokens,
            help=(
                "Maximum prefill-token budget used when chunked prefill is "
                "disabled. Per-iteration scheduling is controlled by "
                "--chunked-prefill-size."
            ),
        )
        # Other runtime options
        parser.add_argument(
            "--stream-interval",
            type=int,
            default=ServerArgs.stream_interval,
            help="The interval (or buffer size) for streaming in terms of the token length. A smaller value makes streaming smoother, while a larger value makes the throughput higher",
        )
        parser.add_argument(
            "--stream-output",
            action="store_true",
            help="Whether to output as a sequence of disjoint segments.",
        )
        parser.add_argument(
            "--seed",
            metavar="SEED",
            type=int,
            default=ServerArgs.seed,
            help="The random seed.",
        )
        parser.add_argument(
            "--distributed-timeout-seconds",
            metavar="DISTRIBUTED_TIMEOUT_SECONDS",
            type=int,
            default=ServerArgs.distributed_timeout_seconds,
            help="Set timeout for torch.distributed initialization.",
        )
        parser.add_argument(
            "--download-dir",
            type=str,
            default=ServerArgs.download_dir,
            help="Model download directory for huggingface.",
        )
        parser.add_argument(
            "--base-gpu-id",
            type=int,
            default=ServerArgs.base_gpu_id,
            help="The base GPU ID to start allocating GPUs from. Useful when running multiple instances on the same machine.",
        )
        parser.add_argument(
            "--gpu-id-step",
            type=int,
            default=ServerArgs.gpu_id_step,
            help="The delta between consecutive GPU IDs that are used. For example, setting it to 2 will use GPU 0,2,4,...",
        )

        # Logging
        parser.add_argument(
            "--log-level",
            type=str,
            default=ServerArgs.log_level,
            help="The logging level of all loggers.",
        )
        parser.add_argument(
            "--log-level-http",
            type=str,
            default=ServerArgs.log_level_http,
            help="The logging level of HTTP server. If not set, reuse --log-level by default.",
        )
        parser.add_argument(
            "--enable-log-requests",
            action=argparse.BooleanOptionalAction,
            default=ServerArgs.enable_log_requests,
            help="Log metadata, inputs, outputs of all requests. The verbosity is decided by --log-requests-level",
        )
        parser.add_argument(
            "--log-requests-level",
            type=int,
            default=0,
            help="0: Log metadata. 1. Log metadata and partial input/output. 2. Log every input/output.",
            choices=[0, 1, 2],
        )
        parser.add_argument(
            "--enable-log-request-stats",
            action=argparse.BooleanOptionalAction,
            default=ServerArgs.enable_log_request_stats,
            help=(
                "Log a one-line per-request performance summary when each request "
                "finishes or aborts: timings (queue/prefill/ttft/total/preemption), "
                "token counts (prompt/cache/output), cache-hit rate, decode "
                "throughput, and spec-decode acceptance. Measured entirely on the "
                "host (no GPU sync), so it adds no engine slowdown."
            ),
        )
        parser.add_argument(
            "--enable-metrics",
            action="store_true",
            help="Enable log metrics.",
        )
        parser.add_argument(
            "--metrics-reporters",
            action="append",
            choices=["prometheus"],
            default=["prometheus"],
            help="Select metrics reporter(can be specified multiple times)",
        )

        parser.add_argument(
            "--app-key",
            type=str,
            default=ServerArgs.app_key,
            help="Set app key of the server",
        )

        parser.add_argument(
            "--decode-log-interval",
            type=int,
            default=ServerArgs.decode_log_interval,
            help="The log interval of decode batch.",
        )
        # API related
        parser.add_argument(
            "--api-key",
            type=str,
            default=ServerArgs.api_key,
            help="Set API key of the server. It is also used in the OpenAI API compatible server.",
        )
        parser.add_argument(
            "--enable-cache-report",
            action="store_true",
            help="Return number of cached tokens in usage.prompt_tokens_details for each openai request.",
        )
        parser.add_argument(
            "--kv-events-config",
            type=str,
            default=ServerArgs.kv_events_config,
            help=(
                "JSON KV cache event publisher config. Set "
                "'enable_kv_cache_events': true and publisher 'zmq' to "
                "publish device prefix-cache mutations."
            ),
        )

        # Data parallelism
        parser.add_argument(
            "--data-parallel-size",
            metavar="DATA_PARALLEL_SIZE",
            type=int,
            default=ServerArgs.data_parallel_size,
            help="The data parallelism size. If not set, inferred from world_size and attn_tp_size.",
        )
        parser.add_argument(
            "--load-balance-method",
            type=str,
            default=ServerArgs.load_balance_method,
            help="The load balancing strategy for data parallelism.",
            choices=[
                "round_robin",
                "shortest_queue",
                "minimum_cache_usage",
            ],
        )
        parser.add_argument(
            "--load-watch-interval",
            type=float,
            default=ServerArgs.load_watch_interval,
            help="The interval of load watching in seconds.",
        )

        # Expert parallelism
        parser.add_argument(
            "--expert-parallel-size",
            "--ep-size",
            type=int,
            default=ServerArgs.ep_size,
            help="The expert parallelism size.",
        )
        parser.add_argument(
            "--init-expert-location",
            type=str,
            default=ServerArgs.init_expert_location,
            help="Initial location of EP experts.",
        )
        parser.add_argument(
            "--ep-num-redundant-experts",
            type=int,
            default=ServerArgs.ep_num_redundant_experts,
            help="Allocate this number of redundant experts in expert parallel.",
        )
        parser.add_argument(
            "--ep-dispatch-algorithm",
            type=str,
            default=ServerArgs.ep_dispatch_algorithm,
            help="The algorithm to choose ranks for redundant experts in expert parallel.",
        )
        parser.add_argument(
            "--eplb-algorithm",
            type=str,
            default=ServerArgs.eplb_algorithm,
            help="Chosen EPLB algorithm",
        )
        parser.add_argument(
            "--expert-distribution-recorder-mode",
            type=str,
            default=ServerArgs.expert_distribution_recorder_mode,
            help="Mode of expert distribution recorder.",
        )
        parser.add_argument(
            "--expert-distribution-recorder-buffer-size",
            type=int,
            default=ServerArgs.expert_distribution_recorder_buffer_size,
            help="Circular buffer size of expert distribution recorder. Set to -1 to denote infinite buffer.",
        )
        parser.add_argument(
            "--enable-expert-distribution-metrics",
            action="store_true",
            help="Enable logging metrics for expert balancedness",
        )
        parser.add_argument(
            "--enable-eplb",
            action="store_true",
            help="Enable EPLB algorithm",
        )
        parser.add_argument(
            "--moe-backend",
            type=str,
            default=ServerArgs.moe_backend,
            help="MoE runner backend: auto, triton, gluon, flashinfer_trtllm",
        )
        parser.add_argument(
            "--draft-moe-backend",
            type=str,
            default=ServerArgs.draft_moe_backend,
            help="MoE runner backend for the draft model in speculative decoding. "
            "If not set, defaults to --moe-backend.",
        )
        parser.add_argument(
            "--all2all-backend",
            metavar="ALL2ALL_BACKEND",
            type=str,
            default=ServerArgs.all2all_backend,
            help="MoE all-to-all backend: none, deepep, etc.",
        )
        parser.add_argument(
            "--deepep-mode",
            type=str,
            choices=["normal", "low_latency", "auto"],
            default=ServerArgs.deepep_mode,
            help="Select the mode when enable DeepEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch.",
        )
        parser.add_argument(
            "--disable-flashinfer-cutlass-moe-fp4-allgather",
            action="store_true",
            help="Disable flashinfer cutlass MoE FP4 allgather.",
        )

        # Multi-node distributed serving
        parser.add_argument(
            "--dist-init-addr",
            type=str,
            help="The host address for initializing distributed backend (e.g., `192.168.0.2:25000`).",
        )
        parser.add_argument(
            "--nnodes", type=int, default=ServerArgs.nnodes, help="The number of nodes."
        )
        parser.add_argument(
            "--node-rank", type=int, default=ServerArgs.node_rank, help="The node rank."
        )

        # Model override args
        parser.add_argument(
            "--hf-overrides",
            metavar="HF_OVERRIDES",
            type=str,
            help="A dictionary in JSON string format used to override default model configurations.",
            default=ServerArgs.hf_overrides,
        )
        parser.add_argument(
            "--preferred-sampling-params",
            type=str,
            help="json-formatted sampling settings that will be returned in /get_model_info",
        )

        # Kernel backend
        attention_backend_choices = [
            "mha",
            "mla",
            "fa3",
            "fa4",
            "triton",
            "flashinfer",
            "trtllm",
            "trtllm_mla",
            "flashmla",
            "tokenspeed_mla",
            "hybrid_linear_attn",
        ]
        parser.add_argument(
            "--attention-backend",
            type=str,
            choices=attention_backend_choices,
            default=ServerArgs.attention_backend,
            help="Choose the kernels for attention layers.",
        )
        parser.add_argument(
            "--drafter-attention-backend",
            type=str,
            choices=attention_backend_choices,
            help="Attention backend for drafter model in speculative decoding. "
            "If not specified, uses the same backend as the main model (attention_backend).",
        )
        parser.add_argument(
            "--sampling-backend",
            type=str,
            choices=[
                "greedy",
                "flashinfer",
                "flashinfer_full",
                "triton",
                "triton_full",
            ],
            default=ServerArgs.sampling_backend,
            help="Sampling backend. "
            "'greedy': argmax + verify_chain_greedy, zero sampling-param plumbing. "
            "'flashinfer': temperature/top_k/top_p via fused softmax + top_k_top_p_sampling_from_probs; "
            "min_p and penalties silently ignored. "
            "'triton': temperature/top_k/top_p via MRV2-style logits-to-Gumbel-Max; "
            "min_p and penalties silently ignored. "
            "'flashinfer_full': adds min_p plus frequency/presence/repetition penalties and logit_bias "
            "via the softmax+renorm+min_p kernel sequence. "
            "'triton_full': adds min_p plus frequency/presence/repetition penalties and logit_bias "
            "with Triton Gumbel-Max for single-step sampling. "
            "Allocates a counts[max_req_pool_size, vocab_size] int32 buffer (substantial memory). "
            "Finite top_k values must be < 128 or -1.",
        )
        parser.add_argument(
            "--dp-sampling",
            action="store_true",
            default=ServerArgs.dp_sampling,
            help=(
                "Enable Batch-DP spec-verify sampling. Backend selection defaults "
                "to auto; override with TOKENSPEED_DP_SAMPLING_BACKEND."
            ),
        )
        parser.add_argument(
            "--dp-sampling-min-bs",
            type=int,
            default=ServerArgs.dp_sampling_min_bs,
            help="Minimum effective decode batch for Batch-DP spec-verify. "
            "Defaults to 2 * TP size.",
        )
        parser.add_argument(
            "--attention-use-fp4-indexer-cache",
            "--attention-config.use-fp4-indexer-cache",
            "--attention_config.use_fp4_indexer_cache",
            type=str_to_bool,
            nargs="?",
            const=True,
            default=ServerArgs.attention_use_fp4_indexer_cache,
            help="Use the MXFP4 sparse attention indexer cache layout.",
        )
        parser.add_argument(
            "--attention-config.use-trtllm-ragged-deepseek-prefill",
            "--attention-config.use_trtllm_ragged_deepseek_prefill",
            "--attention_config.use_trtllm_ragged_deepseek_prefill",
            dest="use_trtllm_ragged_deepseek_prefill",
            type=str_to_bool,
            nargs="?",
            const=True,
            default=ServerArgs.use_trtllm_ragged_deepseek_prefill,
            help="Use ragged prefill for DeepSeek MLA attention.",
        )
        parser.add_argument(
            "--deepseek-v4-mega-moe-max-num-tokens",
            type=int,
            default=ServerArgs.deepseek_v4_mega_moe_max_num_tokens,
            help=(
                "DeepSeek V4 MegaMoE staging-buffer cap on tokens per forward "
                "(0 = derive from chunked-prefill / cuda-graph budgets)."
            ),
        )
        parser.add_argument(
            "--deepseek-v4-indexer-prefill-max-logits-mb",
            type=int,
            default=ServerArgs.deepseek_v4_indexer_prefill_max_logits_mb,
            help=(
                "DeepSeek V4 sparse indexer prefill workspace cap (MiB) for the "
                "softplus_sqrt logits buffer."
            ),
        )
        parser.add_argument(
            "--deepseek-v4-prefill-chunk-size",
            type=int,
            default=ServerArgs.deepseek_v4_prefill_chunk_size,
            help=(
                "Maximum number of requests per DeepSeek V4 FlashMLA prefill " "chunk."
            ),
        )
        parser.add_argument(
            "--grammar-backend",
            type=str,
            choices=["xgrammar", "none"],
            default=ServerArgs.grammar_backend,
            help="Grammar backend. 'none' disables grammar-guided decoding entirely ",
        )
        parser.add_argument(
            "--reasoning-parser",
            type=str,
            default=ServerArgs.reasoning_parser,
            help=(
                "Reasoning parser name (e.g. 'minimax', 'kimi_k25'). "
                "Used to defer json_schema grammars past the model's "
                "reasoning channel."
            ),
        )
        parser.add_argument(
            "--grammar-compile-timeout-secs",
            type=float,
            default=ServerArgs.grammar_compile_timeout_secs,
            help="Per-compile wallclock budget before the request is aborted.",
        )
        parser.add_argument(
            "--grammar-compile-max-retries",
            type=int,
            default=ServerArgs.grammar_compile_max_retries,
            help="Compile timeouts allowed before a grammar key is permanently rejected.",
        )
        parser.add_argument(
            "--disable-any-whitespace",
            action="store_true",
            default=ServerArgs.disable_any_whitespace,
            help="Compile xgrammar JSON grammars in tight mode (no arbitrary "
            "whitespace between tokens). Mitigates models that wedge into "
            "endless whitespace until length cutoff. xgrammar only.",
        )
        parser.add_argument(
            "--disable-capturable-grammar",
            action="store_true",
            default=ServerArgs.disable_capturable_grammar,
            help="Force the synchronous eager grammar fallback even on CUDA. "
            "For parity-testing the captured-grammar path: output should "
            "match; throughput will be lower (sync stall every step).",
        )
        parser.add_argument(
            "--mla-disable-ragged",
            action="store_true",
            help="Disable the ragged prefill wrapper on MLA kernel backends during EXTEND.",
        )

        # Speculative decoding
        parser.add_argument(
            "--draft-model-path-use-base",
            action="store_true",
            help="The path of the draft model weights use the path of the base model",
        )
        parser.add_argument(
            "--speculative-config",
            "--speculative_config",
            type=str,
            default=ServerArgs.speculative_config,
            help="JSON speculative decoding configuration. Supported keys are method, model, and num_speculative_tokens.",
        )
        parser.add_argument(
            "--speculative-algorithm",
            type=str,
            choices=["EAGLE3", "MTP", "DFLASH"],
            help="Speculative algorithm.",
        )
        parser.add_argument(
            "--speculative-draft-model-path",
            type=str,
            help="The path of the draft model weights. This can be a local folder or a Hugging Face repo ID.",
        )
        parser.add_argument(
            "--speculative-draft-model-quantization",
            type=str,
            default=ServerArgs.speculative_draft_model_quantization,
            help="Quantization method for the draft model. Defaults to 'unquant'.",
        )
        parser.add_argument(
            "--speculative-num-steps",
            type=int,
            help="The number of steps sampled from draft model in Speculative Decoding.",
            default=ServerArgs.speculative_num_steps,
        )
        parser.add_argument(
            "--speculative-eagle-topk",
            type=int,
            help="The number of tokens sampled from the draft model in each speculative step.",
            choices=[1],
            default=ServerArgs.speculative_eagle_topk,
        )
        parser.add_argument(
            "--speculative-num-draft-tokens",
            type=int,
            help="The number of tokens sampled from the draft model in Speculative Decoding.",
            default=ServerArgs.speculative_num_draft_tokens,
        )
        parser.add_argument(
            "--enable-output-logprobs",
            action="store_true",
            default=ServerArgs.enable_output_logprobs,
            help="Enable per-token sampled-token logprobs. OFF by default; enabling extends the captured CUDA-graph footprint. Requests asking for logprobs on a server without this flag receive empty logprobs.",
        )
        parser.add_argument(
            "--eagle3-layers-to-capture",
            type=str,
            help="The layers of Eagle3 to capture.",
            default=ServerArgs.eagle3_layers_to_capture,
        )

        # Runtime options
        parser.add_argument(
            "--disable-pdl",
            action="store_true",
            help="Disable PDL launch.",
        )
        prefix_cache_group = parser.add_mutually_exclusive_group()
        prefix_cache_group.add_argument(
            "--enable-prefix-caching",
            action="store_true",
            default=ServerArgs.enable_prefix_caching,
            help="Enable prefix caching.",
        )
        prefix_cache_group.add_argument(
            "--no-enable-prefix-caching",
            dest="enable_prefix_caching",
            action="store_false",
            help="Disable prefix caching.",
        )
        parser.add_argument(
            "--enforce-eager",
            action="store_true",
            help="Disable CUDA graph.",
        )
        parser.add_argument(
            "--disable-cuda-graph-padding",
            action="store_true",
            help="Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed.",
        )
        parser.add_argument(
            "--enable-cudagraph-gc",
            action="store_true",
            help="Enable garbage collection during CUDA graph capture. If disabled (default), GC is frozen during capture to speed up the process.",
        )
        parser.add_argument(
            "--enable-nccl-nvls",
            action="store_true",
            help="Enable NCCL NVLS for prefill heavy requests when available.",
        )
        parser.add_argument(
            "--enable-symm-mem",
            action="store_true",
            help="Enable NCCL symmetric memory for fast collectives.",
        )
        parser.add_argument(
            "--disable-custom-all-reduce",
            action="store_true",
            help="Disable the custom all-reduce kernel and fall back to NCCL.",
        )
        parser.add_argument(
            "--disable-overlap-schedule",
            action="store_true",
            help="Disable the overlap scheduler, which overlaps the CPU scheduler with GPU model worker.",
        )
        parser.add_argument(
            "--disable-tf32",
            action="store_true",
            help="Disable forcing TF32 on for cuBLAS/cuDNN. By default the server sets "
            "NVIDIA_TF32_OVERRIDE=1 and TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1.",
        )
        parser.add_argument(
            "--max-cudagraph-capture-size",
            metavar="MAX_CUDAGRAPH_CAPTURE_SIZE",
            type=int,
            default=ServerArgs.max_cudagraph_capture_size,
            help="Set the maximum batch size for CUDA graph capture.",
        )
        parser.add_argument(
            "--cudagraph-capture-sizes",
            metavar="CUDAGRAPH_CAPTURE_SIZE",
            type=int,
            nargs="+",
            help="Set the list of batch sizes for CUDA graph capture.",
        )
        parser.add_argument(
            "--disable-prefill-graph",
            action="store_true",
            help="Disable cuda graph for prefill.",
        )
        parser.add_argument(
            "--prefill-graph-max-tokens",
            type=int,
            default=ServerArgs.prefill_graph_max_tokens,
            help="Largest token bucket captured by the breakable prefill CUDA "
            "graph. Default (unset) = min(2048, chunked-prefill size); "
            "0 disables.",
        )
        parser.add_argument(
            "--prefill-graph-capture-sizes",
            metavar="PREFILL_GRAPH_CAPTURE_SIZE",
            type=int,
            nargs="+",
            help="Explicit list of token-bucket sizes to capture for the "
            "breakable prefill graph (like --cudagraph-capture-sizes for "
            "decode). Unset: a relative-stride ladder bounding padded compute "
            "at ~12.5%% of any size.",
        )
        parser.add_argument(
            "--enable-nan-detection",
            action="store_true",
            help="Enable the NaN guard: sanitize non-finite logits before "
            "sampling, detect requests whose logits contained NaN (or whose "
            "sampled token id escaped the vocab range), and terminate only "
            "those requests with a numerical error so corruption cannot "
            "spread to the rest of the batch.",
        )
        parser.add_argument(
            "--enable-nvtx",
            action="store_true",
            help="Emit NVTX ranges around input_prep / target_forward / "
            "sampling / drafter stages for nsys profiling. Off by default "
            "(true no-op — no NVTX calls are made). Also enabled by "
            "TOKENSPEED_NVTX=1.",
        )
        parser.add_argument(
            "--enable-p2p-check",
            action="store_true",
            help="Enable the full GPU P2P access check, otherwise trust the driver's P2P report.",
        )
        parser.add_argument(
            "--triton-attention-reduce-in-fp32",
            action="store_true",
            help="Cast the intermediate attention results to fp32 to avoid possible crashes related to fp16."
            "This only affects Triton attention kernels.",
        )
        parser.add_argument(
            "--delete-ckpt-after-loading",
            action="store_true",
            help="Delete the model checkpoint after loading the model.",
        )
        parser.add_argument(
            "--weight-loader-prefetch-checkpoints",
            action="store_true",
            help=(
                "Prefetch safetensors checkpoint shards into OS page cache before "
                "loading. Local ranks split the shard list to reduce repeated reads "
                "from shared filesystems."
            ),
        )
        parser.add_argument(
            "--weight-loader-prefetch-num-threads",
            type=int,
            default=ServerArgs.weight_loader_prefetch_num_threads,
            help="Number of background threads per rank for checkpoint prefetching.",
        )
        parser.add_argument(
            "--enable-memory-saver",
            action="store_true",
            help="Allow saving memory using release_memory_occupation and resume_memory_occupation",
        )
        parser.add_argument(
            "--enable-custom-logit-processor",
            action="store_true",
            help="Enable users to pass custom logit processors to the server (disabled by default for security)",
        )
        # Server warmups
        parser.add_argument(
            "--skip-server-warmup",
            action="store_true",
            help="If set, skip warmup.",
        )
        parser.add_argument(
            "--warmups",
            type=str,
            required=False,
            help="Specify custom warmup functions (csv) to run before server starts eg. --warmups=warmup_name1,warmup_name2 "
            "will run the functions `warmup_name1` and `warmup_name2` specified in warmup.py before the server starts listening for requests",
        )

        parser.add_argument(
            "--tensor-parallel-size",
            "--tp",
            type=int,
            default=None,
            help="Sets tensor parallelism size uniformly (equivalent to --attn-tp-size). "
            "Cannot be used together with --attn-tp-size.",
        )
        parser.add_argument(
            "--enable-expert-parallel",
            action="store_true",
            help="Enable expert parallelism by automatically setting ep_size to world_size.",
        )

        # Specify different parallel strategies, different combinations correspond to different communication groups and weight partitioning, as well as different communication methods
        parser.add_argument(
            "--attn-tp-size",
            type=int,
            default=ServerArgs.attn_tp_size,
            help="Specify tp size for attn part",
        )
        parser.add_argument(
            "--dense-tp-size",
            type=int,
            default=ServerArgs.dense_tp_size,
            help="Specify tp size for dense part, default equals nprocs-per-node, if non dp_attn && combine_dense mode, this parameter will be overridden by attn_tp_size",
        )
        parser.add_argument(
            "--moe-tp-size",
            type=int,
            default=ServerArgs.moe_tp_size,
            help="Specify tp size for MoE part, default equals nprocs-per-node, if non dp_attn && combine_dense mode, this parameter will be overridden by attn_tp_size",
        )
        parser.add_argument(
            "--nprocs-per-node",
            type=int,
            default=ServerArgs.nprocs_per_node,
            help="Number of processes to start per node",
        )
        parser.add_argument(
            "--world-size",
            type=int,
            default=ServerArgs.world_size,
            help="Total number of processes across all nodes.",
        )
        parser.add_argument(
            "--force-deterministic-rsag",
            action="store_true",
            help="Enable force deterministic rsag.",
        )
        parser.add_argument(
            "--disable-sampling-tp-sync",
            action="store_true",
            help="Skip broadcasting sampler outputs across the attention TP "
            "group. Only safe when the sampling kernels are deterministic.",
        )
        parser.add_argument(
            "--low-latency-max-num-tokens-per-gpu",
            type=int,
            default=ServerArgs.low_latency_max_num_tokens_per_gpu,
            help="Low latency max num tokens per gpu",
        )

        parser.add_argument(
            "--mla-chunk-multiplier",
            type=int,
            default=ServerArgs.mla_chunk_multiplier,
            help=(
                "Per-iter MLA chunked-prefill chunk capacity multiplier; "
                "the actual capacity is chunked_prefill_size * mla_chunk_multiplier."
            ),
        )

        # Multimodal
        mm_attention_backend_choices = [
            "fa3",
            "fa4",
            "triton_attn",
            "flashinfer_cudnn",
        ]
        parser.add_argument(
            "--mm-attention-backend",
            type=str,
            choices=mm_attention_backend_choices,
            default=ServerArgs.mm_attention_backend,
            help="Set multimodal attention backend.",
        )
        # Disaggregation
        parser.add_argument(
            "--disaggregation-mode",
            type=str,
            default="null",
            choices=["null", "prefill", "decode", "encode"],
            help='Used for PD/EPD disaggregation. "prefill" for prefill-only server, "decode" for decode-only server, and "encode" for a vision-tower-only server that ships image embeddings to a prefill server. If not specified, it is not disaggregated',
        )
        parser.add_argument(
            "--comm-fusion-max-num-tokens",
            type=int,
            default=ServerArgs.comm_fusion_max_num_tokens,
            help="Max num tokens for communication fusion workspace",
        )
        parser.add_argument(
            "--enable-allreduce-fusion",
            action="store_true",
            help="Enable allreduce fusion for improved decode performance. Auto-enabled on supported single-node TP configurations.",
        )
        parser.add_argument(
            "--disaggregation-bootstrap-port",
            type=int,
            default=ServerArgs.disaggregation_bootstrap_port,
            help="Bootstrap server port on the prefill server. Default is 8998.",
        )
        parser.add_argument(
            "--disaggregation-transfer-backend",
            type=str,
            default=ServerArgs.disaggregation_transfer_backend,
            choices=["mooncake", "mooncake_async"],
            help="The backend for disaggregation transfer. Default is mooncake.",
        )
        parser.add_argument(
            "--disaggregation-ib-device",
            type=str,
            default=ServerArgs.disaggregation_ib_device,
            help="The InfiniBand devices for disaggregation transfer, accepts single device (e.g., --disaggregation-ib-device mlx5_0) "
            "or multiple comma-separated devices (e.g., --disaggregation-ib-device mlx5_0,mlx5_1). "
            "Default is None, which triggers automatic device detection when mooncake backend is enabled.",
        )
        parser.add_argument(
            "--disaggregation-layerwise-interval",
            type=int,
            default=ServerArgs.disaggregation_layerwise_interval,
            help="The interval of layerwise transfer for disaggregation. Default is 1.",
        )
        parser.add_argument(
            "--pdlb-url",
            type=str,
            default=None,
            help="The URL of the PD disaggregation load balancer. If set, the prefill/decode server will register with the load balancer.",
        )

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace):
        args.ep_size = args.expert_parallel_size

        # Resolve model (positional model arg vs --model)
        positional_model = getattr(args, "model_path", None)
        if positional_model is not None and args.model is not None:
            raise ValueError(
                "Cannot specify model both as a positional argument and --model. "
                "Use one or the other."
            )
        if positional_model is not None:
            args.model = positional_model
        if args.model is None:
            raise ValueError(
                "Model is required. Provide it as a positional argument "
                "(e.g., `tokenspeed serve <model>`) or via --model/--model-path."
            )

        # --tensor-parallel-size → --attn-tp-size
        tensor_parallel_size = getattr(args, "tensor_parallel_size", None)
        if tensor_parallel_size is not None:
            if args.attn_tp_size is not None:
                raise ValueError(
                    "Cannot specify both --tensor-parallel-size and --attn-tp-size. "
                    "--tensor-parallel-size is an alias for --attn-tp-size."
                )
            args.attn_tp_size = tensor_parallel_size

        # Only pass fields that argparse actually produced. Falling back to
        # ``None`` for missing attrs would silently clobber dataclass defaults
        # for non-CLI-exposed fields (e.g. ``enable_inline_detokenizer``).
        attrs = [attr.name for attr in dataclasses.fields(cls)]
        return cls(
            **{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)}
        )

    def url(self):
        if is_valid_ipv6_address(self.host):
            return f"http://[{self.host}]:{self.port}"
        return f"http://{self.host}:{self.port}"


def prepare_server_args(argv: list[str]) -> ServerArgs:
    """
    Prepare the server arguments from the command line arguments.

    Args:
        args: The command line arguments. Typically, it should be `sys.argv[1:]`.

    Returns:
        The server arguments.
    """
    parser = argparse.ArgumentParser(allow_abbrev=False)
    ServerArgs.add_cli_args(parser)
    raw_args = parser.parse_args(argv)
    server_args = ServerArgs.from_cli_args(raw_args)
    return server_args


ZMQ_TCP_PORT_DELTA = 233


@dataclasses.dataclass
class PortArgs:
    # The ipc filename for AsyncLLM to receive BatchTokenIDOut directly
    # from the scheduler (zmq).
    tokenizer_ipc_name: str
    # The ipc filename for scheduler (rank 0) to receive inputs from tokenizer (zmq)
    scheduler_input_ipc_name: str

    # The port for nccl initialization (torch.dist)
    nccl_port: int

    # The ipc filename for rpc call between Engine and Scheduler
    rpc_ipc_name: str

    # The ipc filename for Scheduler to send metrics
    metrics_ipc_name: str

    # The ipc filename for Tokenizer and worker tokenizer
    tokenizer_worker_ipc_name: str | None

    @staticmethod
    def init_new(server_args: ServerArgs, dp_rank: int | None = None) -> "PortArgs":
        port = server_args.port + random.randint(100, 1000)
        while True:
            if is_port_available(port):
                break
            if port < 60000:
                port += 42
            else:
                port -= 43

        # DP attention. Use TCP + port to handle both single-node and multi-node.
        if server_args.mapping.nnodes == 1 and server_args.dist_init_addr is None:
            # Only use default port fallback when dp_size == 1
            # For dp_size > 1, we need explicit dist_init_addr to avoid port conflicts
            if server_args.mapping.has_attn_dp:
                raise ValueError(
                    f"When dp_size > 1 (dp_size={server_args.mapping.attn.dp_size}), you must provide --dist-init-addr. "
                    f"Example: --dist-init-addr 127.0.0.1:4000"
                )
            dist_init_addr = ("127.0.0.1", server_args.port + ZMQ_TCP_PORT_DELTA)
        else:
            dist_init_addr = server_args.dist_init_addr.split(":")
        if len(dist_init_addr) != 2:
            raise ValueError(
                "please provide --dist-init-addr as host:port of head node"
            )

        dist_init_host, dist_init_port = dist_init_addr
        dist_init_port = int(dist_init_port)

        # Scan forward until we find a port cluster where all derived ports are free.
        # This handles the case where a previous engine instance left ports in
        # TIME_WAIT or its child processes haven't fully terminated yet.
        # Note: the port at offset +1 (formerly detokenizer_port) is intentionally
        # skipped so the rest of the port layout stays stable for any external
        # tooling that indexed off the historical port cluster.
        while True:
            port_base = dist_init_port + 1
            rpc_port = port_base + 2
            metrics_ipc_port = port_base + 3
            if dp_rank is None:
                # TokenizerManager to DataParallelController
                scheduler_input_port = port_base + 4
            else:
                scheduler_input_port = port_base + 2 + 1 + dp_rank
            rpc_ipc_port = scheduler_input_port + 1
            if all(
                is_port_available(p)
                for p in [
                    dist_init_port,
                    port_base,
                    rpc_port,
                    metrics_ipc_port,
                    scheduler_input_port,
                    rpc_ipc_port,
                ]
            ):
                break
            dist_init_port += 10

        return PortArgs(
            tokenizer_ipc_name=f"tcp://{dist_init_host}:{port_base}",
            scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}",
            nccl_port=port,
            rpc_ipc_name=f"tcp://{dist_init_host}:{rpc_port}",
            metrics_ipc_name=f"tcp://{dist_init_host}:{metrics_ipc_port}",
            tokenizer_worker_ipc_name=None,
        )