@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 706ff13224

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 15, 03:30 UTC

[Model] Adds support for SlimMoE models Phi-tiny-MoE-instruct (#20286)

Signed-off-by: Zichong Li <t-lizichong@microsoft.com@Reasoning-H100-VM3.drbuo4tcjzruhloch3eo0b25ef.cx.internal.cloudapp.net> Co-authored-by: Zichong Li <t-lizichong@microsoft.com@Reasoning-H100-VM3.drbuo4tcjzruhloch3eo0b25ef.cx.internal.cloudapp.net> Co-authored-by: Isotr0py <2037008807@qq.com>

zichongli5 committed Jul 2, 2025 at 05:54 UTC 706ff1322468e8c445fe166a6bdc3136fba4698a

1 file changed +10 -1

vllm/model_executor/models/phimoe.py

+10 -1

@@ -68,6 +68,7 @@ class PhiMoEConfig(PretrainedConfig):

68 num_hidden_layers=32,

69 num_attention_heads=32,

70 num_key_value_heads=8,

71 + head_dim=None,

72 hidden_act="silu",

73 max_position_embeddings=4096 * 32,

74 initializer_range=0.02,

@@ -101,8 +102,11 @@ class PhiMoEConfig(PretrainedConfig):

102 # for backward compatibility

103 if num_key_value_heads is None:

104 num_key_value_heads = num_attention_heads

105 + if head_dim is None:

106 + head_dim = hidden_size // num_attention_heads

107

108 self.num_key_value_heads = num_key_value_heads

109 + self.head_dim = head_dim

110 self.hidden_act = hidden_act

111 self.initializer_range = initializer_range

112 self.rms_norm_eps = rms_norm_eps

@@ -294,6 +298,7 @@ class PhiMoEAttention(nn.Module):

298 hidden_size: int,

299 num_heads: int,

300 num_kv_heads: int,

301 + head_dim: Optional[int] = None,

302 max_position: int = 4096 * 32,

303 rope_theta: float = 10000,

304 cache_config: Optional[CacheConfig] = None,

@@ -317,7 +322,9 @@ class PhiMoEAttention(nn.Module):

322 # the KV heads across multiple tensor parallel GPUs.

323 assert tp_size % self.total_num_kv_heads == 0

324 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

320 - self.head_dim = hidden_size // self.total_num_heads

325 + if head_dim is None:

326 + head_dim = hidden_size // num_heads

327 + self.head_dim = head_dim

328 self.q_size = self.num_heads * self.head_dim

329 self.kv_size = self.num_kv_heads * self.head_dim

330 self.scaling = self.head_dim**-0.5

@@ -387,6 +394,8 @@ class PhiMoEDecoderLayer(nn.Module):

394 num_heads=config.num_attention_heads,

395 max_position=config.max_position_embeddings,

396 num_kv_heads=config.num_key_value_heads,

397 + head_dim=getattr(config, "head_dim",

398 + self.hidden_size // config.num_attention_heads),

399 rope_theta=rope_theta,

400 cache_config=cache_config,

401 quant_config=quant_config,