@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 65552b476b

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Misc] Use config definitions from Transformers library (#21913)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Aug 9, 2025 at 14:10 UTC 65552b476b1c475ef433995d2699bb27428693b3

11 files changed +54 -51

vllm/model_executor/models/aimv2.py

+11 -11

@@ -8,7 +8,6 @@ from typing import Optional

9 import torch

10 import torch.nn as nn

11 -from transformers import PretrainedConfig

12 from vllm.attention.layer import MultiHeadAttention

13 from vllm.distributed import get_tensor_model_parallel_world_size

@@ -21,12 +20,13 @@ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,

20 from vllm.model_executor.layers.quantization.base_config import (

21 QuantizationConfig)

22 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

23 +from vllm.transformers_utils.configs.ovis import AIMv2Config

26 class AIMv2SwiGLUFFN(nn.Module):

28 - def __init__(self, config: PretrainedConfig,

29 - quant_config: QuantizationConfig, prefix: str):

28 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,

29 + prefix: str):

30 super().__init__()

31 hidden_features = config.intermediate_size

32 in_features = config.hidden_size

@@ -57,7 +57,7 @@ class AIMv2SwiGLUFFN(nn.Module):

58 class AIMv2PatchEmbed(nn.Module):

60 - def __init__(self, config: PretrainedConfig):

60 + def __init__(self, config: AIMv2Config):

61 super().__init__()

62 self.proj = nn.Conv2d(

63 config.num_channels,

@@ -75,7 +75,7 @@ class AIMv2PatchEmbed(nn.Module):

76 class AIMv2ViTPreprocessor(nn.Module):

78 - def __init__(self, config: PretrainedConfig):

78 + def __init__(self, config: AIMv2Config):

79 super().__init__()

80 num_patches = (config.image_size // config.patch_size)**2

@@ -93,8 +93,8 @@ class AIMv2ViTPreprocessor(nn.Module):

94 class AIMv2Attention(nn.Module):

96 - def __init__(self, config: PretrainedConfig,

97 - quant_config: QuantizationConfig, prefix: str):

96 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,

97 + prefix: str):

98 super().__init__()

99 self.config = config

100 self.embed_dim = config.hidden_size

@@ -141,8 +141,8 @@ class AIMv2Attention(nn.Module):

141

142 class AIMv2Block(nn.Module):

143

144 - def __init__(self, config: PretrainedConfig,

145 - quant_config: QuantizationConfig, prefix: str):

144 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,

145 + prefix: str):

146 super().__init__()

147 self.attn = AIMv2Attention(config,

148 quant_config=quant_config,

@@ -163,7 +163,7 @@ class AIMv2Transformer(nn.Module):

163

164 def __init__(

165 self,

166 - config: PretrainedConfig,

166 + config: AIMv2Config,

167 quant_config: QuantizationConfig,

168 *,

169 require_post_norm: Optional[bool] = None,

@@ -193,7 +193,7 @@ class AIMv2Transformer(nn.Module):

193 class AIMv2Model(torch.nn.Module):

194

195 def __init__(self,

196 - config: PretrainedConfig,

196 + config: AIMv2Config,

197 quant_config: QuantizationConfig,

198 *,

199 require_post_norm: Optional[bool] = None,

vllm/model_executor/models/commandr.py

+4 -4

@@ -27,7 +27,7 @@ from typing import Optional, Union

28 import torch

29 from torch import nn

30 -from transformers import CohereConfig

30 +from transformers import Cohere2Config, CohereConfig

32 from vllm.attention import Attention

33 from vllm.compilation.decorators import support_torch_compile

@@ -89,7 +89,7 @@ class CohereMLP(nn.Module):

90 def __init__(

91 self,

92 - config: CohereConfig,

92 + config: Union[CohereConfig, Cohere2Config],

93 quant_config: Optional[QuantizationConfig] = None,

94 prefix: str = "",

95 ):

@@ -124,7 +124,7 @@ class CohereAttention(nn.Module):

124

125 def __init__(

126 self,

127 - config: CohereConfig,

127 + config: Union[CohereConfig, Cohere2Config],

128 cache_config: Optional[CacheConfig] = None,

129 quant_config: Optional[QuantizationConfig] = None,

130 prefix: str = "",

@@ -242,7 +242,7 @@ class CohereAttention(nn.Module):

242 class CohereDecoderLayer(nn.Module):

243

244 def __init__(self,

245 - config: CohereConfig,

245 + config: Union[CohereConfig, Cohere2Config],

246 cache_config: Optional[CacheConfig] = None,

247 quant_config: Optional[QuantizationConfig] = None,

248 prefix: str = ""):

vllm/model_executor/models/dbrx.py

+7 -7

@@ -6,7 +6,7 @@ from typing import Optional, Union

7 import torch

8 import torch.nn as nn

9 -from transformers import PretrainedConfig

9 +from transformers import DbrxConfig

11 from vllm.attention import Attention

12 from vllm.config import CacheConfig, VllmConfig

@@ -39,7 +39,7 @@ class DbrxRouter(nn.Module):

40 def __init__(

41 self,

42 - config: PretrainedConfig,

42 + config: DbrxConfig,

43 params_dtype: Optional[torch.dtype] = None,

44 ):

45 super().__init__()

@@ -63,7 +63,7 @@ class DbrxExperts(FusedMoE):

64 def __init__(

65 self,

66 - config: PretrainedConfig,

66 + config: DbrxConfig,

67 quant_config: Optional[QuantizationConfig] = None,

68 params_dtype: Optional[torch.dtype] = None,

69 prefix: str = "",

@@ -138,7 +138,7 @@ class DbrxMoE(nn.Module):

138

139 def __init__(

140 self,

141 - config: PretrainedConfig,

141 + config: DbrxConfig,

142 quant_config: Optional[QuantizationConfig] = None,

143 params_dtype: Optional[torch.dtype] = None,

144 prefix: str = "",

@@ -169,7 +169,7 @@ class DbrxAttention(nn.Module):

169

170 def __init__(

171 self,

172 - config: PretrainedConfig,

172 + config: DbrxConfig,

173 cache_config: Optional[CacheConfig] = None,

174 quant_config: Optional[QuantizationConfig] = None,

175 prefix: str = "",

@@ -249,7 +249,7 @@ class DbrxFusedNormAttention(nn.Module):

249

250 def __init__(

251 self,

252 - config: PretrainedConfig,

252 + config: DbrxConfig,

253 cache_config: Optional[CacheConfig] = None,

254 quant_config: Optional[QuantizationConfig] = None,

255 prefix: str = "",

@@ -284,7 +284,7 @@ class DbrxBlock(nn.Module):

284

285 def __init__(

286 self,

287 - config: PretrainedConfig,

287 + config: DbrxConfig,

288 cache_config: Optional[CacheConfig] = None,

289 quant_config: Optional[QuantizationConfig] = None,

290 prefix: str = "",

vllm/model_executor/models/deepseek_v2.py

+9 -6

@@ -29,7 +29,7 @@ from typing import Any, Optional, Union

30 import torch

31 from torch import nn

32 -from transformers import PretrainedConfig

32 +from transformers import DeepseekV2Config, DeepseekV3Config

34 from vllm.attention import Attention

35 from vllm.compilation.decorators import support_torch_compile

@@ -100,7 +100,7 @@ class DeepseekV2MoE(nn.Module):

100

101 def __init__(

102 self,

103 - config: PretrainedConfig,

103 + config: Union[DeepseekV2Config, DeepseekV3Config],

104 quant_config: Optional[QuantizationConfig] = None,

105 prefix: str = "",

106 enable_eplb: bool = False,

@@ -221,7 +221,7 @@ class DeepseekV2Attention(nn.Module):

221

222 def __init__(

223 self,

224 - config: PretrainedConfig,

224 + config: Union[DeepseekV2Config, DeepseekV3Config],

225 hidden_size: int,

226 num_heads: int,

227 qk_nope_head_dim: int,

@@ -373,7 +373,7 @@ class DeepseekV2MLAAttention(nn.Module):

373

374 def __init__(

375 self,

376 - config: PretrainedConfig,

376 + config: Union[DeepseekV2Config, DeepseekV3Config],

377 hidden_size: int,

378 num_heads: int,

379 qk_nope_head_dim: int,

@@ -538,7 +538,7 @@ class DeepseekV2DecoderLayer(nn.Module):

538

539 def __init__(

540 self,

541 - config: PretrainedConfig,

541 + config: Union[DeepseekV2Config, DeepseekV3Config],

542 prefix: str,

543 model_config: ModelConfig,

544 cache_config: Optional[CacheConfig] = None,

@@ -973,7 +973,10 @@ class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):

973 pass

974

975

976 -def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,

976 +# Compatibility with

977 +# https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py

978 +def get_spec_layer_idx_from_weight_name(config: Union[DeepseekV2Config,

979 + DeepseekV3Config],

980 weight_name: str) -> Optional[int]:

981 if (hasattr(config, "num_nextn_predict_layers")

982 and config.num_nextn_predict_layers > 0):

vllm/model_executor/models/dots1.py

+4 -4

@@ -29,7 +29,7 @@ from typing import Any, Optional, Union

30 import torch

31 from torch import nn

32 -from transformers import PretrainedConfig

32 +from transformers import Dots1Config

34 from vllm.attention import Attention

35 from vllm.compilation.decorators import support_torch_compile

@@ -99,7 +99,7 @@ class Dots1MoE(nn.Module):

100 def __init__(

101 self,

102 - config: PretrainedConfig,

102 + config: Dots1Config,

103 quant_config: Optional[QuantizationConfig] = None,

104 prefix: str = "",

105 ):

@@ -174,7 +174,7 @@ class Dots1Attention(nn.Module):

174 hidden_size: int,

175 num_heads: int,

176 num_kv_heads: int,

177 - config: PretrainedConfig,

177 + config: Dots1Config,

178 rope_theta: float = 10000,

179 rope_scaling: Optional[dict[str, Any]] = None,

180 max_position_embeddings: int = 8192,

@@ -260,7 +260,7 @@ class Dots1DecoderLayer(nn.Module):

260

261 def __init__(

262 self,

263 - config: PretrainedConfig,

263 + config: Dots1Config,

264 prefix: str,

265 model_config: ModelConfig,

266 cache_config: Optional[CacheConfig] = None,

vllm/model_executor/models/exaone4.py

+3 -3

@@ -26,7 +26,7 @@ from typing import Any, Optional, Union

27 import torch

28 from torch import nn

29 -from transformers import PretrainedConfig

29 +from transformers import Exaone4Config

31 from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

@@ -96,7 +96,7 @@ class Exaone4Attention(nn.Module):

97 def __init__(

98 self,

99 - config: PretrainedConfig,

99 + config: Exaone4Config,

100 hidden_size: int,

101 num_heads: int,

102 num_kv_heads: int,

@@ -224,7 +224,7 @@ class Exaone4DecoderLayer(nn.Module):

224

225 def __init__(

226 self,

227 - config: PretrainedConfig,

227 + config: Exaone4Config,

228 cache_config: Optional[CacheConfig] = None,

229 quant_config: Optional[QuantizationConfig] = None,

230 prefix: str = "",

vllm/model_executor/models/glm4_moe.py

+5 -5

@@ -28,7 +28,7 @@ from typing import Any, Optional, Union

29 import torch

30 from torch import nn

31 -from transformers import PretrainedConfig

31 +from transformers.models.glm4_moe import Glm4MoeConfig

33 from vllm.attention import Attention

34 from vllm.compilation.decorators import support_torch_compile

@@ -100,7 +100,7 @@ class Glm4MoE(nn.Module):

100

101 def __init__(

102 self,

103 - config: PretrainedConfig,

103 + config: Glm4MoeConfig,

104 quant_config: Optional[QuantizationConfig] = None,

105 prefix: str = "",

106 enable_eplb: bool = False,

@@ -198,7 +198,7 @@ class Glm4MoeAttention(nn.Module):

198

199 def __init__(

200 self,

201 - config: PretrainedConfig,

201 + config: Glm4MoeConfig,

202 hidden_size: int,

203 num_heads: int,

204 num_kv_heads: int,

@@ -297,7 +297,7 @@ class Glm4MoeDecoderLayer(nn.Module):

297

298 def __init__(

299 self,

300 - config: PretrainedConfig,

300 + config: Glm4MoeConfig,

301 cache_config: Optional[CacheConfig] = None,

302 quant_config: Optional[QuantizationConfig] = None,

303 prefix: str = "",

@@ -681,7 +681,7 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

681 return self.model.get_expert_mapping()

682

683

684 -def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,

684 +def get_spec_layer_idx_from_weight_name(config: Glm4MoeConfig,

685 weight_name: str) -> Optional[int]:

686 if hasattr(config,

687 "num_nextn_predict_layers") and (config.num_nextn_predict_layers

vllm/model_executor/models/minimax_text_01.py

+3 -3

@@ -12,7 +12,7 @@ import torch.distributed

12 import torch.nn.functional as F

13 from einops import rearrange

14 from torch import nn

15 -from transformers.configuration_utils import PretrainedConfig

15 +from transformers import MiniMaxConfig

17 from vllm import envs

18 from vllm.attention import Attention, AttentionMetadata

@@ -656,7 +656,7 @@ class MiniMaxText01DecoderLayer(nn.Module):

656

657 def __init__(

658 self,

659 - config: PretrainedConfig,

659 + config: MiniMaxConfig,

660 cache_config: Optional[CacheConfig] = None,

661 quant_config: Optional[QuantizationConfig] = None,

662 expert_num: int = 1,

@@ -860,7 +860,7 @@ class MiniMaxText01Model(nn.Module):

860

861 def __init__(

862 self,

863 - config: PretrainedConfig,

863 + config: MiniMaxConfig,

864 quant_config: Optional[QuantizationConfig] = None,

865 cache_config: Optional[CacheConfig] = None,

866 scheduler_config=None,

vllm/model_executor/models/olmoe.py

+2 -2

@@ -19,7 +19,7 @@ from typing import Any, Optional, Union

20 import torch

21 from torch import nn

22 -from transformers import PretrainedConfig

22 +from transformers import OlmoeConfig

24 from vllm.attention import Attention

25 from vllm.compilation.decorators import support_torch_compile

@@ -205,7 +205,7 @@ class OlmoeDecoderLayer(nn.Module):

205

206 def __init__(

207 self,

208 - config: PretrainedConfig,

208 + config: OlmoeConfig,

209 cache_config: Optional[CacheConfig] = None,

210 quant_config: Optional[QuantizationConfig] = None,

211 prefix: str = "",

vllm/model_executor/models/qwen2_moe.py

+3 -3

@@ -30,7 +30,7 @@ from typing import Any, Optional, Union

30 import torch

31 import torch.nn.functional as F

32 from torch import nn

33 -from transformers import PretrainedConfig

33 +from transformers import Qwen2MoeConfig

35 from vllm.attention import Attention

36 from vllm.compilation.decorators import support_torch_compile

@@ -98,7 +98,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):

99 def __init__(

100 self,

101 - config: PretrainedConfig,

101 + config: Qwen2MoeConfig,

102 quant_config: Optional[QuantizationConfig] = None,

103 prefix: str = "",

104 ):

@@ -256,7 +256,7 @@ class Qwen2MoeDecoderLayer(nn.Module):

256

257 def __init__(

258 self,

259 - config: PretrainedConfig,

259 + config: Qwen2MoeConfig,

260 cache_config: Optional[CacheConfig] = None,

261 quant_config: Optional[QuantizationConfig] = None,

262 prefix: str = "",

vllm/model_executor/models/qwen3_moe.py

+3 -3

@@ -28,7 +28,7 @@ from typing import Any, Optional, Union

29 import torch

30 from torch import nn

31 -from transformers import PretrainedConfig

31 +from transformers import Qwen3MoeConfig

33 from vllm.attention import Attention

34 from vllm.compilation.decorators import support_torch_compile

@@ -101,7 +101,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):

101

102 def __init__(

103 self,

104 - config: PretrainedConfig,

104 + config: Qwen3MoeConfig,

105 quant_config: Optional[QuantizationConfig] = None,

106 prefix: str = "",

107 enable_eplb: bool = False,

@@ -278,7 +278,7 @@ class Qwen3MoeDecoderLayer(nn.Module):

278

279 def __init__(

280 self,

281 - config: PretrainedConfig,

281 + config: Qwen3MoeConfig,

282 cache_config: Optional[CacheConfig] = None,

283 quant_config: Optional[QuantizationConfig] = None,

284 prefix: str = "",