| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | from pydantic import model_validator |
| 7 | from typing_extensions import Self |
| 8 | |
| 9 | from vllm.config.utils import config, get_hash_factors, hash_factors |
| 10 | |
| 11 | if TYPE_CHECKING: |
| 12 | from vllm.config.model import ModelConfig |
| 13 | from vllm.config.parallel import ParallelConfig |
| 14 | |
| 15 | # Architecture -> the hf_text_config field naming its n-gram layers. A model is |
| 16 | # only configurable here if it actually has such layers to store. |
| 17 | _NGRAM_LAYER_FIELDS = { |
| 18 | "DeepseekV41ForCausalLM": "engram_layer_ids", |
| 19 | "Qwen4ExpForCausalLM": "ple_layer_ids", |
| 20 | "Qwen4ExpForConditionalGeneration": "ple_layer_ids", |
| 21 | } |
| 22 | |
| 23 | |
| 24 | def model_has_engram_layers(model_config: "ModelConfig | None") -> bool: |
| 25 | """Whether the model carries n-gram embedding layers.""" |
| 26 | if model_config is None: |
| 27 | return False |
| 28 | field = _NGRAM_LAYER_FIELDS.get(model_config.architecture) |
| 29 | if field is None: |
| 30 | return False |
| 31 | return bool(getattr(model_config.hf_text_config, field, None)) |
| 32 | |
| 33 | |
| 34 | @config |
| 35 | class EngramConfig: |
| 36 | """Configuration for Engram embedding storage and sharding.""" |
| 37 | |
| 38 | cpu_offload: bool = True |
| 39 | """Store embedding weights in pinned CPU memory for UVA lookup.""" |
| 40 | |
| 41 | embedding_across_dp: bool = False |
| 42 | """Shard embeddings across TP and all DP ranks when enabled. |
| 43 | Otherwise, each DP rank has a separate TP-sharded embedding replica.""" |
| 44 | |
| 45 | dp_shared_memory: bool | None = None |
| 46 | """Share CPU-offloaded embedding weights between co-located |
| 47 | DP replicas. Each node stores one copy of every TP shard, reducing host |
| 48 | memory without per-step Engram DP collectives. Requires sufficient |
| 49 | /dev/shm capacity and a shared IPC namespace. Defaults to enabled whenever |
| 50 | the other settings allow it, falling back to per-replica tables when DP |
| 51 | replicas are not co-located on one node or /dev/shm cannot hold them.""" |
| 52 | |
| 53 | use_thp: bool = False |
| 54 | """Back private CPU-offloaded tables with transparent huge pages (best |
| 55 | effort, falls back to ordinary pinned pages). Prefaulting the tables at |
| 56 | startup takes longer. Requires cpu_offload without dp_shared_memory.""" |
| 57 | |
| 58 | @model_validator(mode="after") |
| 59 | def _validate_shared_memory(self) -> Self: |
| 60 | if self.dp_shared_memory and not self.cpu_offload: |
| 61 | raise ValueError("dp_shared_memory requires cpu_offload=True") |
| 62 | if self.use_thp and (not self.cpu_offload or self.dp_shared_memory): |
| 63 | raise ValueError( |
| 64 | "use_thp requires cpu_offload=True and dp_shared_memory=False" |
| 65 | ) |
| 66 | return self |
| 67 | |
| 68 | def verify_model_config(self, model_config: "ModelConfig | None") -> None: |
| 69 | """Reject Engram configuration for models without n-gram embeddings.""" |
| 70 | field = ( |
| 71 | _NGRAM_LAYER_FIELDS.get(model_config.architecture) |
| 72 | if model_config is not None |
| 73 | else None |
| 74 | ) |
| 75 | if ( |
| 76 | model_config is None |
| 77 | or field is None |
| 78 | or not getattr(model_config.hf_text_config, field, None) |
| 79 | ): |
| 80 | raise ValueError( |
| 81 | "EngramConfig requires a model with supported Engram " |
| 82 | "embeddings and non-empty n-gram layer ids." |
| 83 | ) |
| 84 | |
| 85 | def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None: |
| 86 | """Share host tables by default wherever the configuration permits.""" |
| 87 | if self.dp_shared_memory is None: |
| 88 | self.dp_shared_memory = ( |
| 89 | self.cpu_offload |
| 90 | and not self.use_thp |
| 91 | and parallel_config.data_parallel_size > 1 |
| 92 | and not parallel_config.enable_elastic_ep |
| 93 | ) |
| 94 | |
| 95 | def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: |
| 96 | """Reject unsupported embedding parallel topologies.""" |
| 97 | if self.dp_shared_memory: |
| 98 | if parallel_config.data_parallel_size <= 1: |
| 99 | raise ValueError("dp_shared_memory requires data_parallel_size > 1.") |
| 100 | if parallel_config.enable_elastic_ep: |
| 101 | raise ValueError("dp_shared_memory is not supported with elastic EP.") |
| 102 | if ( |
| 103 | self.embedding_across_dp |
| 104 | and parallel_config.data_parallel_size > 1 |
| 105 | and parallel_config.enable_elastic_ep |
| 106 | ): |
| 107 | raise ValueError( |
| 108 | "Engram embedding_across_dp is not supported with elastic EP yet." |
| 109 | ) |
| 110 | |
| 111 | def get_parallel_size(self, parallel_config: "ParallelConfig") -> int: |
| 112 | """Derive the embedding group size from the parallel configuration.""" |
| 113 | size = parallel_config.tensor_parallel_size |
| 114 | if self.embedding_across_dp and parallel_config.data_parallel_size > 1: |
| 115 | size *= parallel_config.data_parallel_size |
| 116 | return size |
| 117 | |
| 118 | def compute_hash(self) -> str: |
| 119 | """Hash settings that affect embedding execution and graph structure.""" |
| 120 | return hash_factors(get_hash_factors(self, set())) |