# SPDX-License-Identifier: Apache-2.0 # Derived from vLLM; modified by DeepSeek. import torch from deepinfer import envs from deepinfer.engine.engine_type import EngineType from deepinfer.model_config_loader import load_model_config def _get_runtime_max_model_len(max_model_len: int) -> int: """Return the exclusive upper bound of every runtime position.""" nextn = envs.DEEPINFER_USE_NEXTN_DECODE scheduler_overhead = nextn + (nextn + 1 if envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING else 0) draft_overhead = nextn return max_model_len + scheduler_overhead + draft_overhead class ModelConfig: """Configuration for the model. Args: model: Name or path of the huggingface model to use. tokenizer: Name or path of the huggingface tokenizer to use. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. load_format: Weight initialization mode: "auto" loads the HF safetensors index and shards from the model directory. "dummy" initializes synthetic weights for profiling, using the same HF model configuration. seed: Random seed for reproducibility. revision: 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. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. max_model_len: Maximum length of a sequence (including prompt and output). If None, will be derived from the model. """ def __init__( self, model: str, tokenizer: str, trust_remote_code: bool, load_format: str, seed: int, max_num_seqs: int, revision: str | None = None, tokenizer_revision: str | None = None, max_model_len: int | None = None, max_num_batched_tokens: int | None = None, ignore_len: int = -1, served_model_name: str = "deepinfer-model", engine_type: EngineType = EngineType.NORMAL, ) -> None: self.model = model self.tokenizer = tokenizer self.trust_remote_code = trust_remote_code self.load_format = load_format self._verify_load_format() self.seed = seed self.revision = revision self.tokenizer_revision = tokenizer_revision self.max_num_seqs = max_num_seqs self.max_model_len = max_model_len if max_model_len is not None else 4096 self.runtime_max_model_len = _get_runtime_max_model_len(self.max_model_len) self.max_num_batched_tokens = _get_max_num_batched_tokens(self.max_model_len, max_num_batched_tokens, max_num_seqs) self.ignore_len = ignore_len self.served_model_name = served_model_name self.dtype = torch.bfloat16 self.use_compressed_tokenizer = False self.engram_tokenizer = self.tokenizer self.engine_type = engine_type self._load_model_cfg() _GPT_RUNTIME_ATTRS = ( "max_num_seqs", "max_num_batched_tokens", "runtime_max_model_len", "engine_type", ) def _load_model_cfg(self) -> None: self.model_cfg = load_model_config(self.model) self.architecture = self.model_cfg.architecture for attr in self._GPT_RUNTIME_ATTRS: setattr(self.model_cfg.gpt, attr, getattr(self, attr)) def __getstate__(self) -> dict: # model_cfg 不可 pickle，改为只传递必要的运行时标量， # 由 worker 进程在 __setstate__ 中重新 load model_cfg 后回填。 state = self.__dict__.copy() model_cfg = state.pop("model_cfg", None) overrides: dict = {} if model_cfg is not None: overrides = {k: model_cfg.gpt[k] for k in self._GPT_RUNTIME_ATTRS} state["_gpt_runtime_overrides"] = overrides return state def __setstate__(self, newstate: dict) -> None: overrides = newstate.pop("_gpt_runtime_overrides", {}) self.__dict__.update(newstate) self.model_cfg = load_model_config(self.model) self.architecture = self.model_cfg.architecture for key, value in overrides.items(): setattr(self.model_cfg.gpt, key, value) def _verify_load_format(self) -> None: load_format = self.load_format.lower() supported_load_format = ["auto", "dummy"] assert load_format in supported_load_format self.load_format = load_format def get_vocab_size(self) -> int: return self.model_cfg.gpt.vocab_size class CacheConfig: def __init__( self, block_size: int, gpu_memory_utilization: float, lookahead_tokens: int = 0, ) -> None: self.block_size = block_size self.gpu_memory_utilization = gpu_memory_utilization self._verify_args() # Will be set after profiling. self.num_gpu_blocks = None # Will be set after init cache. self.num_state_seqs = None self.bytes_of_state_per_seqs = None self.bytes_per_block = None self.lookahead_tokens = lookahead_tokens def _verify_args(self) -> None: if self.gpu_memory_utilization > 1.0: raise ValueError(f"GPU memory utilization must be less than 1.0. Got {self.gpu_memory_utilization}.") class ParallelConfig: def __init__( self, data_parallel_size: int, context_parallel_size: int, expert_parallel_size: int, engram_parallel_size: int, engine_rank: int, engine_group_rank: int, engine_type: EngineType = EngineType.NORMAL, ) -> None: if data_parallel_size == 0: data_parallel_size = max(1, expert_parallel_size // context_parallel_size) self.data_parallel_size = data_parallel_size self.context_parallel_size = context_parallel_size self.expert_parallel_size = expert_parallel_size self.engine_rank = engine_rank self.data_parallel_rank = engine_rank % data_parallel_size self.engine_group_rank = engine_group_rank self.engine_type = engine_type self.model_parallel_size = self.context_parallel_size self.engram_parallel_size = engram_parallel_size or (self.model_parallel_size * data_parallel_size) self.world_size = max(self.model_parallel_size * data_parallel_size, self.engram_parallel_size) self._verify_args() def _verify_args(self) -> None: assert self.expert_parallel_size > 0 assert (self.model_parallel_size * self.data_parallel_size) % self.expert_parallel_size == 0 assert self.engram_parallel_size > 0 assert self.world_size % (self.model_parallel_size * self.data_parallel_size) == 0 assert self.world_size % self.engram_parallel_size == 0 assert 0 <= self.engine_rank < self.world_size // self.model_parallel_size class SchedulerConfig: def __init__( self, max_num_batched_tokens: int | None, max_num_seqs: int, max_model_len: int, stop_preempted_request: bool = False, chunked_prefill: bool = False, prefill_attempt_interval: int = envs.DEEPINFER_PREFILL_ATTEMPT_INTERVAL, prefill_any_attempt_interval: int = envs.DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL, ) -> None: self.max_num_batched_tokens = _get_max_num_batched_tokens(max_model_len, max_num_batched_tokens, max_num_seqs) self.max_num_seqs = max_num_seqs self.max_model_len = max_model_len self.runtime_max_model_len = _get_runtime_max_model_len(max_model_len) self.stop_preempted_request = stop_preempted_request self.chunked_prefill = chunked_prefill # NORMAL 引擎每隔多少个 decode step 尝试一次 prefill（ALL 投票点：全员有活才做） self.prefill_attempt_interval = prefill_attempt_interval # NORMAL 引擎 ANY 保底投票点间隔：任一 rank 有活即集体 prefill，防饿死 self.prefill_any_attempt_interval = prefill_any_attempt_interval self._verify_args() def _verify_args(self) -> None: if (not self.chunked_prefill) and self.max_num_batched_tokens < self.max_model_len: raise ValueError( f"max_num_batched_tokens ({self.max_num_batched_tokens}) is smaller than max_model_len ({self.max_model_len}). This effectively limits the maximum sequence length to max_num_batched_tokens and makes DeepInfer reject longer sequences. Please increase max_num_batched_tokens or decrease max_model_len." ) if self.prefill_attempt_interval < 1: raise ValueError(f"prefill_attempt_interval ({self.prefill_attempt_interval}) should be >= 1") if self.prefill_any_attempt_interval < 1: raise ValueError(f"prefill_any_attempt_interval ({self.prefill_any_attempt_interval}) should be >= 1") def _get_max_num_batched_tokens( max_model_len: int, max_num_batched_tokens: int | None, max_num_seqs: int, ) -> int: max_num_decoding_tokens = max_num_seqs * (1 + envs.DEEPINFER_USE_NEXTN_DECODE) if max_num_batched_tokens is not None: assert max_num_batched_tokens >= max_num_decoding_tokens return max_num_batched_tokens return max(max_model_len, 4096, max_num_decoding_tokens)