@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 07c4353057

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

[Model] Support Grok1 (#13795)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Feb 25, 2025 at 20:07 UTC 07c4353057c2abb65db87104d44973c8cf5833bf

11 files changed +634 -17

docs/source/models/supported_models.md

@@ -286,6 +286,11 @@ See [this page](#generative-models) for more information on how to use generativ

286 * `parasail-ai/GritLM-7B-vllm`.

287 * ✅︎

288 * ✅︎

289 +- * `Grok1ModelForCausalLM`

290 + * Grok1

291 + * `hpcai-tech/grok-1`.

292 + * ✅︎

293 + * ✅︎

294 - * `InternLMForCausalLM`

295 * InternLM

296 * `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc.

tests/models/registry.py

@@ -130,6 +130,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

130 "GPTNeoXForCausalLM": _HfExamplesInfo("EleutherAI/pythia-160m"),

131 "GraniteForCausalLM": _HfExamplesInfo("ibm/PowerLM-3b"),

132 "GraniteMoeForCausalLM": _HfExamplesInfo("ibm/PowerMoE-3b"),

133 + "Grok1ModelForCausalLM": _HfExamplesInfo("hpcai-tech/grok-1",

134 + trust_remote_code=True),

135 "InternLMForCausalLM": _HfExamplesInfo("internlm/internlm-chat-7b",

136 trust_remote_code=True),

137 "InternLM2ForCausalLM": _HfExamplesInfo("internlm/internlm2-chat-7b",

vllm/model_executor/layers/fused_moe/fused_moe.py

+31 -12

@@ -1040,6 +1040,7 @@ def inplace_fused_experts(hidden_states: torch.Tensor,

1040 w2: torch.Tensor,

1041 topk_weights: torch.Tensor,

1042 topk_ids: torch.Tensor,

1043 + activation: str = "silu",

1044 use_fp8_w8a8: bool = False,

1045 use_int8_w8a16: bool = False,

1046 use_int4_w4a16: bool = False,

@@ -1053,9 +1054,10 @@ def inplace_fused_experts(hidden_states: torch.Tensor,

1054 a2_scale: Optional[torch.Tensor] = None,

1055 block_shape: Optional[List[int]] = None) -> None:

1056 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,

1056 - use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16,

1057 - global_num_experts, expert_map, w1_scale, w2_scale,

1058 - w1_zp, w2_zp, a1_scale, a2_scale, block_shape)

1057 + activation, use_fp8_w8a8, use_int8_w8a16,

1058 + use_int4_w4a16, global_num_experts, expert_map,

1059 + w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,

1060 + block_shape)

1061

1062

1063 def inplace_fused_experts_fake(

@@ -1064,6 +1066,7 @@ def inplace_fused_experts_fake(

1066 w2: torch.Tensor,

1067 topk_weights: torch.Tensor,

1068 topk_ids: torch.Tensor,

1069 + activation: str = "silu",

1070 use_fp8_w8a8: bool = False,

1071 use_int8_w8a16: bool = False,

1072 use_int4_w4a16: bool = False,

@@ -1093,6 +1096,7 @@ def outplace_fused_experts(

1096 w2: torch.Tensor,

1097 topk_weights: torch.Tensor,

1098 topk_ids: torch.Tensor,

1099 + activation: str = "silu",

1100 use_fp8_w8a8: bool = False,

1101 use_int8_w8a16: bool = False,

1102 use_int4_w4a16: bool = False,

@@ -1106,7 +1110,7 @@ def outplace_fused_experts(

1110 a2_scale: Optional[torch.Tensor] = None,

1111 block_shape: Optional[List[int]] = None) -> torch.Tensor:

1112 return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,

1109 - False, use_fp8_w8a8, use_int8_w8a16,

1113 + False, activation, use_fp8_w8a8, use_int8_w8a16,

1114 use_int4_w4a16, global_num_experts, expert_map,

1115 w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,

1116 a2_scale, block_shape)

@@ -1118,6 +1122,7 @@ def outplace_fused_experts_fake(

1122 w2: torch.Tensor,

1123 topk_weights: torch.Tensor,

1124 topk_ids: torch.Tensor,

1125 + activation: str = "silu",

1126 use_fp8_w8a8: bool = False,

1127 use_int8_w8a16: bool = False,

1128 use_int4_w4a16: bool = False,

@@ -1147,6 +1152,7 @@ def fused_experts(hidden_states: torch.Tensor,

1152 topk_weights: torch.Tensor,

1153 topk_ids: torch.Tensor,

1154 inplace: bool = False,

1155 + activation: str = "silu",

1156 use_fp8_w8a8: bool = False,

1157 use_int8_w8a16: bool = False,

1158 use_int4_w4a16: bool = False,

@@ -1162,15 +1168,17 @@ def fused_experts(hidden_states: torch.Tensor,

1168

1169 if inplace:

1170 torch.ops.vllm.inplace_fused_experts(

1165 - hidden_states, w1, w2, topk_weights, topk_ids, use_fp8_w8a8,

1166 - use_int8_w8a16, use_int4_w4a16, global_num_experts, expert_map,

1167 - w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape)

1171 + hidden_states, w1, w2, topk_weights, topk_ids, activation,

1172 + use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, global_num_experts,

1173 + expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,

1174 + block_shape)

1175 return hidden_states

1176 else:

1177 return torch.ops.vllm.outplace_fused_experts(

1171 - hidden_states, w1, w2, topk_weights, topk_ids, use_fp8_w8a8,

1172 - use_int8_w8a16, use_int4_w4a16, global_num_experts, expert_map,

1173 - w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape)

1178 + hidden_states, w1, w2, topk_weights, topk_ids, activation,

1179 + use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, global_num_experts,

1180 + expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,

1181 + block_shape)

1182

1183

1184 def fused_experts_impl(hidden_states: torch.Tensor,

@@ -1179,6 +1187,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,

1187 topk_weights: torch.Tensor,

1188 topk_ids: torch.Tensor,

1189 inplace: bool = False,

1190 + activation: str = "silu",

1191 use_fp8_w8a8: bool = False,

1192 use_int8_w8a16: bool = False,

1193 use_int4_w4a16: bool = False,

@@ -1303,8 +1312,14 @@ def fused_experts_impl(hidden_states: torch.Tensor,

1312 use_int4_w4a16=use_int4_w4a16,

1313 block_shape=block_shape)

1314

1306 - torch.ops._C.silu_and_mul(intermediate_cache2,

1307 - intermediate_cache1.view(-1, N))

1315 + if activation == "silu":

1316 + torch.ops._C.silu_and_mul(intermediate_cache2,

1317 + intermediate_cache1.view(-1, N))

1318 + elif activation == "gelu":

1319 + torch.ops._C.gelu_and_mul(intermediate_cache2,

1320 + intermediate_cache1.view(-1, N))

1321 + else:

1322 + raise ValueError(f"Unsupported FusedMoe activation: {activation}")

1323

1324 invoke_fused_moe_kernel(intermediate_cache2,

1325 w2,

@@ -1339,6 +1354,7 @@ def fused_moe(

1354 topk: int,

1355 renormalize: bool,

1356 inplace: bool = False,

1357 + activation: str = "silu",

1358 use_grouped_topk: bool = False,

1359 num_expert_group: Optional[int] = None,

1360 topk_group: Optional[int] = None,

@@ -1370,6 +1386,8 @@ def fused_moe(

1386 - renormalize (bool): If True, renormalize the top-k weights to sum to 1.

1387 - inplace (bool): If True, perform the operation in-place.

1388 Defaults to False.

1389 + - activation (str): The activation function to apply after the first

1390 + MoE layer.

1391 - num_expert_group: Optional[int]: additional parameter for grouped_topk

1392 - topk_group: Optional[int]: additional parameter for grouped_topk

1393 - use_grouped_topk: If True, use grouped_topk instead of fused_topk

@@ -1420,6 +1438,7 @@ def fused_moe(

1438 topk_weights,

1439 topk_ids,

1440 inplace=inplace,

1441 + activation=activation,

1442 use_fp8_w8a8=use_fp8_w8a8,

1443 use_int8_w8a16=use_int8_w8a16,

1444 use_int4_w4a16=use_int4_w4a16,

vllm/model_executor/layers/fused_moe/layer.py

+17 -5

@@ -120,7 +120,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

120 expert_map: Optional[torch.Tensor] = None,

121 custom_routing_function: Optional[Callable] = None,

122 scoring_func: str = "softmax",

123 - e_score_correction_bias: Optional[torch.Tensor] = None

123 + e_score_correction_bias: Optional[torch.Tensor] = None,

124 + activation: str = "silu",

125 ) -> torch.Tensor:

126 return self.forward(x=x,

127 layer=layer,

@@ -134,7 +135,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

135 expert_map=expert_map,

136 custom_routing_function=custom_routing_function,

137 scoring_func=scoring_func,

137 - e_score_correction_bias=e_score_correction_bias)

138 + e_score_correction_bias=e_score_correction_bias,

139 + activation=activation)

140

141 def forward_cuda(

142 self,

@@ -150,7 +152,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

152 expert_map: Optional[torch.Tensor] = None,

153 custom_routing_function: Optional[Callable] = None,

154 scoring_func: str = "softmax",

153 - e_score_correction_bias: Optional[torch.Tensor] = None

155 + e_score_correction_bias: Optional[torch.Tensor] = None,

156 + activation: str = "silu",

157 ) -> torch.Tensor:

158 topk_weights, topk_ids = FusedMoE.select_experts(

159 hidden_states=x,

@@ -170,6 +173,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

173 topk_weights=topk_weights,

174 topk_ids=topk_ids,

175 inplace=True,

176 + activation=activation,

177 global_num_experts=global_num_experts,

178 expert_map=expert_map)

179

@@ -186,9 +190,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

190 global_num_experts: int = -1,

191 expert_map: Optional[torch.Tensor] = None,

192 custom_routing_function: Optional[Callable] = None,

193 + activation: str = "silu",

194 **kwargs,

195 ):

196 assert custom_routing_function is None

197 + assert activation == "silu", f"{activation} is not supported."

198 return layer.ipex_fusion(

199 x,

200 use_grouped_topk,

@@ -213,7 +219,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

219 expert_map: Optional[torch.Tensor] = None,

220 custom_routing_function: Optional[Callable] = None,

221 scoring_func: str = "softmax",

216 - e_score_correction_bias: Optional[torch.Tensor] = None

222 + e_score_correction_bias: Optional[torch.Tensor] = None,

223 + activation: str = "silu",

224 ) -> torch.Tensor:

225 assert not use_grouped_topk

226 assert num_expert_group is None

@@ -225,6 +232,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

232 if e_score_correction_bias is not None:

233 raise NotImplementedError(

234 "Expert score correction bias is not supported for TPU.")

235 + assert activation == "silu", f"{activation} is not supported for TPU."

236 return fused_moe_pallas(hidden_states=x,

237 w1=layer.w13_weight,

238 w2=layer.w2_weight,

@@ -277,6 +285,7 @@ class FusedMoE(torch.nn.Module):

285 custom_routing_function: Optional[Callable] = None,

286 scoring_func: str = "softmax",

287 e_score_correction_bias: Optional[torch.Tensor] = None,

288 + activation: str = "silu",

289 ):

290 super().__init__()

291

@@ -305,6 +314,7 @@ class FusedMoE(torch.nn.Module):

314 self.custom_routing_function = custom_routing_function

315 self.scoring_func = scoring_func

316 self.e_score_correction_bias = e_score_correction_bias

317 + self.activation = activation

318 self.expert_map = None

319

320 if self.ep_size > 1:

@@ -653,7 +663,9 @@ class FusedMoE(torch.nn.Module):

663 num_expert_group=self.num_expert_group,

664 custom_routing_function=self.custom_routing_function,

665 scoring_func=self.scoring_func,

656 - e_score_correction_bias=self.e_score_correction_bias)

666 + e_score_correction_bias=self.e_score_correction_bias,

667 + activation=self.activation,

668 + )

669

670 if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):

671 # Default set to False. (May have to add shared expert outputs.)

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -469,7 +469,9 @@ class AWQMoEMethod(FusedMoEMethodBase):

469 custom_routing_function: Optional[Callable] = None,

470 scoring_func: str = "softmax",

471 e_score_correction_bias: Optional[torch.Tensor] = None,

472 + activation: str = "silu",

473 ) -> torch.Tensor:

474 + assert activation == "silu", "Only SiLU activation is supported."

475 if expert_map is not None:

476 raise NotImplementedError(

477 "Expert Parallelism is not supported for "

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -219,6 +219,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

219 custom_routing_function: Optional[Callable] = None,

220 scoring_func: str = "softmax",

221 e_score_correction_bias: Optional[torch.Tensor] = None,

222 + activation: str = "silu",

223 ) -> torch.Tensor:

224 from vllm.model_executor.layers.fused_moe import fused_experts

225

@@ -240,6 +241,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

241 topk_weights=topk_weights,

242 topk_ids=topk_ids,

243 inplace=True,

244 + activation=activation,

245 use_fp8_w8a8=True,

246 global_num_experts=global_num_experts,

247 expert_map=expert_map,

@@ -550,7 +552,9 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

552 custom_routing_function: Optional[Callable] = None,

553 scoring_func: str = "softmax",

554 e_score_correction_bias: Optional[torch.Tensor] = None,

555 + activation: str = "silu",

556 ) -> torch.Tensor:

557 + assert activation == "silu", "Only SiLU activation is supported."

558 if expert_map is not None:

559 raise NotImplementedError(

560 "Expert Parallelism is not supported for "

vllm/model_executor/layers/quantization/experts_int8.py

@@ -113,6 +113,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

113 custom_routing_function: Optional[Callable] = None,

114 scoring_func: str = "softmax",

115 e_score_correction_bias: Optional[torch.Tensor] = None,

116 + activation: str = "silu",

117 ) -> torch.Tensor:

118 from vllm.model_executor.layers.fused_moe import fused_experts

119

@@ -134,6 +135,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

135 topk_weights=topk_weights,

136 topk_ids=topk_ids,

137 inplace=True,

138 + activation=activation,

139 use_int8_w8a16=True,

140 global_num_experts=global_num_experts,

141 expert_map=expert_map,

vllm/model_executor/layers/quantization/fp8.py

@@ -675,6 +675,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

675 custom_routing_function: Optional[Callable] = None,

676 scoring_func: str = "softmax",

677 e_score_correction_bias: Optional[torch.Tensor] = None,

678 + activation: str = "silu",

679 ) -> torch.Tensor:

680 from vllm.model_executor.layers.fused_moe import fused_experts

681

@@ -698,6 +699,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

699 topk_weights=topk_weights,

700 topk_ids=topk_ids,

701 inplace=True,

702 + activation=activation,

703 use_fp8_w8a8=True,

704 global_num_experts=global_num_experts,

705 expert_map=expert_map,

vllm/model_executor/layers/quantization/gptq_marlin.py

@@ -590,7 +590,10 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

590 custom_routing_function: Optional[Callable] = None,

591 scoring_func: str = "softmax",

592 e_score_correction_bias: Optional[torch.Tensor] = None,

593 + activation: str = "silu",

594 ) -> torch.Tensor:

595 + assert activation == "silu", "Only SiLU activation is supported."

596 +

597 # The input must currently be float16

598 orig_dtype = x.dtype

599 x = x.half()

vllm/model_executor/models/grok1.py new

+565

@@ -0,0 +1,565 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# Adapted from

3 +# https://github.com/ROCm/vllm/blob/cea7419f151cc50293a05b7fac8547f8f887c9f6/vllm/model_executor/models/grok1.py

6 +#

7 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

8 +# and OPT implementations in this library. It has been modified from its

9 +# original forms to accommodate minor architectural differences compared

10 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

11 +#

12 +# Licensed under the Apache License, Version 2.0 (the "License");

13 +# you may not use this file except in compliance with the License.

14 +# You may obtain a copy of the License at

15 +#

16 +# http://www.apache.org/licenses/LICENSE-2.0

17 +#

18 +# Unless required by applicable law or agreed to in writing, software

19 +# distributed under the License is distributed on an "AS IS" BASIS,

20 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

21 +# See the License for the specific language governing permissions and

22 +# limitations under the License.

23 +"""Inference-only Grok1 model."""

24 +from typing import Iterable, List, Optional, Set, Tuple, Union

25 +

26 +import torch

27 +import torch.nn.functional as F

28 +from torch import nn

29 +

30 +from vllm.attention import Attention, AttentionMetadata

31 +from vllm.compilation.decorators import support_torch_compile

32 +from vllm.config import CacheConfig, VllmConfig

33 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

34 +from vllm.model_executor.layers.fused_moe import FusedMoE

35 +from vllm.model_executor.layers.layernorm import RMSNorm

36 +from vllm.model_executor.layers.linear import (QKVParallelLinear,

37 + ReplicatedLinear,

38 + RowParallelLinear)

39 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

40 +from vllm.model_executor.layers.quantization import QuantizationConfig

41 +from vllm.model_executor.layers.rotary_embedding import get_rope

42 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

43 +from vllm.model_executor.layers.vocab_parallel_embedding import (

44 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

45 +from vllm.model_executor.model_loader.weight_utils import (

46 + default_weight_loader, maybe_remap_kv_scale_name)

47 +from vllm.model_executor.sampling_metadata import SamplingMetadata

48 +from vllm.sequence import IntermediateTensors

49 +

50 +from .interfaces import SupportsLoRA, SupportsPP

51 +from .utils import (is_pp_missing_parameter,

52 + make_empty_intermediate_tensors_factory, make_layers,

53 + maybe_prefix)

54 +

55 +# Default Grok1-specific constants, overridden by config values if present

56 +DEFAULT_ATTN_OUTPUT_MULTIPLIER = 0.08838834764831845

57 +DEFAULT_OUTPUT_MULTIPLIER_SCALE = 0.5773502691896257

58 +DEFAULT_EMBEDDING_MULTIPLIER_SCALE = 78.38367176906169

59 +

60 +

61 +class Grok1MoE(nn.Module):

62 + """A tensor-parallel MoE implementation for Grok1 that shards each expert

63 + across all ranks.

64 +

65 + Each expert's weights are sharded across all ranks and a fused MoE

66 + kernel is used for the forward pass, and finally we reduce the outputs

67 + across ranks.

68 + """

69 +

70 + def __init__(self,

71 + num_experts: int,

72 + top_k: int,

73 + hidden_size: int,

74 + intermediate_size: int,

75 + params_dtype: Optional[torch.dtype] = None,

76 + quant_config: Optional[QuantizationConfig] = None,

77 + tp_size: Optional[int] = None,

78 + prefix: str = ""):

79 + super().__init__()

80 + self.hidden_size = hidden_size

81 +

82 + # Gate always runs at half / full precision for now.

83 + self.gate = ReplicatedLinear(hidden_size,

84 + num_experts,

85 + bias=False,

86 + params_dtype=params_dtype,

87 + quant_config=None,

88 + prefix=f"{prefix}.gate")

89 +

90 + self.experts = FusedMoE(num_experts=num_experts,

91 + top_k=top_k,

92 + hidden_size=hidden_size,

93 + intermediate_size=intermediate_size,

94 + params_dtype=params_dtype,

95 + reduce_results=True,

96 + renormalize=True,

97 + quant_config=quant_config,

98 + tp_size=tp_size,

99 + activation="gelu",

100 + prefix=f"{prefix}.experts")

101 +

102 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

103 + # NOTE: hidden_states can have either 1D or 2D shape.

104 + orig_shape = hidden_states.shape

105 + hidden_states = hidden_states.view(-1, self.hidden_size)

106 + # router_logits: (num_tokens, n_experts)

107 + router_logits, _ = self.gate(hidden_states)

108 + router_logits = 30.0 * F.tanh(router_logits / 30.0)

109 + final_hidden_states = self.experts(hidden_states, router_logits)

110 + return final_hidden_states.view(orig_shape)

111 +

112 +

113 +class Grok1Attention(nn.Module):

114 +

115 + def __init__(

116 + self,

117 + hidden_size: int,

118 + num_heads: int,

119 + num_kv_heads: int,

120 + max_position: int = 4096 * 32,

121 + rope_theta: float = 10000,

122 + cache_config: Optional[CacheConfig] = None,

123 + quant_config: Optional[QuantizationConfig] = None,

124 + prefix: str = "",

125 + config=None, # Added config parameter

126 + ) -> None:

127 + super().__init__()

128 + self.hidden_size = hidden_size

129 + self.config = config # Store config reference

130 + tp_size = get_tensor_model_parallel_world_size()

131 + self.total_num_heads = num_heads

132 + assert self.total_num_heads % tp_size == 0

133 + self.num_heads = self.total_num_heads // tp_size

134 + self.total_num_kv_heads = num_kv_heads

135 + if self.total_num_kv_heads >= tp_size:

136 + # Number of KV heads is greater than TP size, so we partition

137 + # the KV heads across multiple tensor parallel GPUs.

138 + assert self.total_num_kv_heads % tp_size == 0

139 + else:

140 + # Number of KV heads is less than TP size, so we replicate

141 + # the KV heads across multiple tensor parallel GPUs.

142 + assert tp_size % self.total_num_kv_heads == 0

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

144 + self.head_dim = hidden_size // self.total_num_heads

145 + self.q_size = self.num_heads * self.head_dim

146 + self.kv_size = self.num_kv_heads * self.head_dim

147 + self.scaling = self.head_dim**-0.5

148 + self.rope_theta = rope_theta

149 +

150 + self.qkv_proj = QKVParallelLinear(

151 + hidden_size,

152 + self.head_dim,

153 + self.total_num_heads,

154 + self.total_num_kv_heads,

155 + bias=False,

156 + quant_config=quant_config,

157 + prefix=f"{prefix}.qkv_proj",

158 + )

159 + self.o_proj = RowParallelLinear(

160 + self.total_num_heads * self.head_dim,

161 + hidden_size,

162 + bias=False,

163 + quant_config=quant_config,

164 + prefix=f"{prefix}.o_proj",

165 + )

166 + self.rotary_emb = get_rope(

167 + self.head_dim,

168 + rotary_dim=self.head_dim,

169 + max_position=max_position,

170 + base=int(self.rope_theta),

171 + is_neox_style=True,

172 + )

173 +

174 + attn_logits_soft_cap = max(

175 + getattr(config, "attn_logit_softcapping", 30.0), 0.0)

176 +

177 + self.attn = Attention(self.num_heads,

178 + self.head_dim,

179 + self.scaling,

180 + num_kv_heads=self.num_kv_heads,

181 + cache_config=cache_config,

182 + quant_config=quant_config,

183 + logits_soft_cap=attn_logits_soft_cap,

184 + prefix=f"{prefix}.attn")

185 +

186 + def forward(

187 + self,

188 + positions: torch.Tensor,

189 + hidden_states: torch.Tensor,

190 + kv_cache: torch.Tensor,

191 + attn_metadata: AttentionMetadata,

192 + ) -> torch.Tensor:

193 + qkv, _ = self.qkv_proj(hidden_states)

194 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

195 + q, k = self.rotary_emb(positions, q, k)

196 + attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

197 + output, _ = self.o_proj(attn_output)

198 +

199 + # Apply attention output multiplier if specified in config

200 + attn_multiplier = getattr(self.config, "attn_output_multiplier",

201 + None) if self.config else None

202 + if attn_multiplier is not None:

203 + output = output * attn_multiplier

204 + return output

205 +

206 +

207 +class Grok1DecoderLayer(nn.Module):

208 +

209 + def __init__(

210 + self,

211 + config,

212 + cache_config: Optional[CacheConfig] = None,

213 + quant_config: Optional[QuantizationConfig] = None,

214 + prefix: str = "",

215 + ) -> None:

216 + super().__init__()

217 + self.hidden_size = config.hidden_size

218 + # Check for fp8 quantization

219 + self.use_fp8 = False

220 + if quant_config is not None:

221 + self.use_fp8 = getattr(quant_config, "is_fp8_w8a8",

222 + lambda: False)()

223 + if not self.use_fp8 and hasattr(quant_config, "is_fp8"):

224 + self.use_fp8 = quant_config.is_fp8

225 +

226 + # Requires transformers > 4.32.0

227 + # Default rope_theta value if not in config

228 + rope_theta = 10000

229 + self.attn = Grok1Attention(

230 + hidden_size=self.hidden_size,

231 + num_heads=config.num_attention_heads,

232 + max_position=config.max_position_embeddings,

233 + num_kv_heads=config.num_key_value_heads,

234 + rope_theta=rope_theta,

235 + cache_config=cache_config,

236 + quant_config=quant_config,

237 + prefix=f"{prefix}.attn",

238 + config=config) # Pass config to Grok1Attention

239 +

240 + # Grok1 uses "num_experts" in its config

241 + num_experts = getattr(config, "num_experts", 8)

242 + num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)

243 +

244 + self.moe_block = Grok1MoE(num_experts=num_experts,

245 + top_k=num_experts_per_tok,

246 + hidden_size=config.hidden_size,

247 + intermediate_size=config.intermediate_size,

248 + quant_config=quant_config,

249 + prefix=f"{prefix}.moe_block")

250 +

251 + self.pre_attn_norm = RMSNorm(config.hidden_size,

252 + eps=config.rms_norm_eps)

253 + self.post_attn_norm = RMSNorm(config.hidden_size,

254 + eps=config.rms_norm_eps)

255 + self.pre_moe_norm = RMSNorm(config.hidden_size,

256 + eps=config.rms_norm_eps)

257 + self.post_moe_norm = RMSNorm(config.hidden_size,

258 + eps=config.rms_norm_eps)

259 +

260 + def forward(

261 + self,

262 + positions: torch.Tensor,

263 + hidden_states: torch.Tensor,

264 + kv_cache: torch.Tensor,

265 + attn_metadata: AttentionMetadata,

266 + residual: Optional[torch.Tensor],

267 + ) -> Tuple[torch.Tensor, torch.Tensor]:

268 + # Self Attention

269 + if residual is None:

270 + residual = hidden_states

271 + hidden_states = self.pre_attn_norm(hidden_states)

272 + else:

273 + hidden_states, residual = self.pre_attn_norm(

274 + hidden_states, residual)

275 +

276 + hidden_states = self.attn(

277 + positions=positions,

278 + hidden_states=hidden_states,

279 + kv_cache=kv_cache,

280 + attn_metadata=attn_metadata,

281 + )

282 +

283 + # Post attention normalization

284 + hidden_states = self.post_attn_norm(hidden_states)

285 +

286 + # MoE block with normalization

287 + hidden_states, residual = self.pre_moe_norm(hidden_states, residual)

288 + hidden_states = self.moe_block(hidden_states)

289 + hidden_states = self.post_moe_norm(hidden_states)

290 +

291 + return hidden_states, residual

292 +

293 +

294 +@support_torch_compile

295 +class Grok1Model(nn.Module):

296 +

297 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

298 + super().__init__()

299 +

300 + config = vllm_config.model_config.hf_config

301 + cache_config = vllm_config.cache_config

302 + quant_config = vllm_config.quant_config

303 + lora_config = vllm_config.lora_config

304 +

305 + self.padding_idx = config.pad_token_id

306 + lora_vocab = (lora_config.lora_extra_vocab_size *

307 + (lora_config.max_loras or 1)) if lora_config else 0

308 + self.vocab_size = config.vocab_size + lora_vocab

309 + self.org_vocab_size = config.vocab_size

310 + self.embedding_multiplier_scale = getattr(

311 + config, "embedding_multiplier_scale",

312 + DEFAULT_EMBEDDING_MULTIPLIER_SCALE)

313 +

314 + self.embed_tokens = VocabParallelEmbedding(

315 + self.vocab_size,

316 + config.hidden_size,

317 + org_num_embeddings=config.vocab_size,

318 + quant_config=quant_config,

319 + )

320 +

321 + self.start_layer, self.end_layer, self.layers = make_layers(

322 + config.num_hidden_layers,

323 + lambda prefix: Grok1DecoderLayer(

324 + config, cache_config, quant_config=quant_config, prefix=prefix

325 + ),

326 + prefix=f"{prefix}.layers")

327 +

328 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

329 + self.make_empty_intermediate_tensors = (

330 + make_empty_intermediate_tensors_factory(

331 + ["hidden_states", "residual"], config.hidden_size))

332 +

333 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

334 + hidden_states = self.embed_tokens(input_ids)

335 + hidden_states = hidden_states * self.embedding_multiplier_scale

336 + return hidden_states

337 +

338 + def forward(

339 + self,

340 + input_ids: torch.Tensor,

341 + positions: torch.Tensor,

342 + kv_caches: List[torch.Tensor],

343 + attn_metadata: AttentionMetadata,

344 + intermediate_tensors: Optional[IntermediateTensors],

345 + inputs_embeds: Optional[torch.Tensor] = None,

346 + ) -> Union[torch.Tensor, IntermediateTensors]:

347 + if get_pp_group().is_first_rank:

348 + if inputs_embeds is not None:

349 + hidden_states = inputs_embeds

350 + else:

351 + hidden_states = self.get_input_embeddings(input_ids)

352 + residual = None

353 + else:

354 + assert intermediate_tensors is not None

355 + hidden_states = intermediate_tensors["hidden_states"]

356 + residual = intermediate_tensors["residual"]

357 +

358 + for i in range(self.start_layer, self.end_layer):

359 + layer = self.layers[i]

360 + hidden_states, residual = layer(positions, hidden_states,

361 + kv_caches[i - self.start_layer],

362 + attn_metadata, residual)

363 +

364 + if not get_pp_group().is_last_rank:

365 + return IntermediateTensors({

366 + "hidden_states": hidden_states,

367 + "residual": residual

368 + })

369 +

370 + hidden_states, _ = self.norm(hidden_states, residual)

371 + return hidden_states

372 +

373 +

374 +class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

375 + fall_back_to_pt_during_load = False

376 +

377 + packed_modules_mapping = {

378 + "qkv_proj": [

379 + "q_proj",

380 + "k_proj",

381 + "v_proj",

382 + ],

383 + }

384 +

385 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

386 + super().__init__()

387 +

388 + config = vllm_config.model_config.hf_config

389 + quant_config = vllm_config.quant_config

390 + lora_config = vllm_config.lora_config

391 +

392 + self.config = config

393 + self.lora_config = lora_config

394 + self.quant_config = quant_config

395 +

396 + self.model = Grok1Model(vllm_config=vllm_config,

397 + prefix=maybe_prefix(prefix, "model"))

398 +

399 + self.unpadded_vocab_size = config.vocab_size

400 + if lora_config:

401 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

402 +

403 + self.lm_head = ParallelLMHead(

404 + self.unpadded_vocab_size,

405 + config.hidden_size,

406 + org_num_embeddings=config.vocab_size,

407 + padding_size=DEFAULT_VOCAB_PADDING_SIZE

408 + # We need bigger padding if using lora for kernel compatibility

409 + if not lora_config else lora_config.lora_vocab_padding_size,

410 + quant_config=quant_config,

411 + prefix=maybe_prefix(prefix, "lm_head"),

412 + )

413 +

414 + if self.config.tie_word_embeddings:

415 + self.lm_head.weight = self.model.embed_tokens.weight

416 +

417 + self.output_multiplier_scale = getattr(

418 + config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE)

419 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

420 + config.vocab_size,

421 + self.output_multiplier_scale)

422 +

423 + self.sampler = get_sampler()

424 + self.make_empty_intermediate_tensors = (

425 + self.model.make_empty_intermediate_tensors)

426 +

427 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

428 + return self.model.get_input_embeddings(input_ids)

429 +

430 + def forward(

431 + self,

432 + input_ids: torch.Tensor,

433 + positions: torch.Tensor,

434 + kv_caches: List[torch.Tensor],

435 + attn_metadata: AttentionMetadata,

436 + intermediate_tensors: Optional[IntermediateTensors] = None,

437 + inputs_embeds: Optional[torch.Tensor] = None,

438 + ) -> Union[torch.Tensor, IntermediateTensors]:

439 + hidden_states = self.model(input_ids, positions, kv_caches,

440 + attn_metadata, intermediate_tensors,

441 + inputs_embeds)

442 + return hidden_states

443 +

444 + def compute_logits(

445 + self,

446 + hidden_states: torch.Tensor,

447 + sampling_metadata: SamplingMetadata,

448 + ) -> Optional[torch.Tensor]:

449 + logits = self.logits_processor(self.lm_head, hidden_states,

450 + sampling_metadata)

451 + return logits

452 +

453 + def sample(

454 + self,

455 + logits: Optional[torch.Tensor],

456 + sampling_metadata: SamplingMetadata,

457 + ) -> Optional[SamplerOutput]:

458 + next_tokens = self.sampler(logits, sampling_metadata)

459 + return next_tokens

460 +

461 + def load_weights(self, weights: Iterable[Tuple[str,

462 + torch.Tensor]]) -> Set[str]:

463 + stacked_params_mapping = [

464 + # (param_name, shard_name, shard_id)

465 + ("qkv_proj", "q_proj", "q"),

466 + ("qkv_proj", "k_proj", "k"),

467 + ("qkv_proj", "v_proj", "v"),

468 + ]

469 +

470 + # Map Grok1's unique expert parameter names to standard names

471 + # Grok1 uses "num_experts" in its config

472 + num_experts = getattr(self.config, "num_experts", 8)

473 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

474 + ckpt_gate_proj_name="linear", # Grok1 specific

475 + ckpt_down_proj_name="linear_1", # Grok1 specific

476 + ckpt_up_proj_name="linear_v", # Grok1 specific

477 + num_experts=num_experts)

478 +

479 + params_dict = dict(self.named_parameters())

480 + loaded_params: Set[str] = set()

481 +

482 + for name, loaded_weight in weights:

483 + if "rotary_emb.inv_freq" in name:

484 + continue

485 +

486 + if (self.quant_config is not None and

487 + (scale_name := self.quant_config.get_cache_scale(name))):

488 + # Loading kv cache quantization scales

489 + param = params_dict[scale_name]

490 + weight_loader = getattr(param, "weight_loader",

491 + default_weight_loader)

492 + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else

493 + loaded_weight[0])

494 + weight_loader(param, loaded_weight)

495 + loaded_params.add(scale_name)

496 + continue

497 +

498 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

499 + if weight_name not in name:

500 + continue

501 + name = name.replace(weight_name, param_name)

502 + # Skip loading extra bias for GPTQ models.

503 + if ((name.endswith(".bias") or name.endswith("_bias"))

504 + and name not in params_dict):

505 + continue

506 + # Skip layers on other devices.

507 + if is_pp_missing_parameter(name, self):

508 + continue

509 + if name.endswith("scale"):

510 + # Remapping the name of FP8 kv-scale.

511 + name = maybe_remap_kv_scale_name(name, params_dict)

512 + if name is None:

513 + continue

514 + param = params_dict[name]

515 + weight_loader = param.weight_loader

516 + weight_loader(param, loaded_weight, shard_id)

517 + break

518 + else:

519 + for mapping in expert_params_mapping:

520 + param_name, weight_name, expert_id, shard_id = mapping

521 + if weight_name not in name:

522 + continue

523 + name = name.replace(weight_name, param_name)

524 + # Skip layers on other devices.

525 + if is_pp_missing_parameter(name, self):

526 + continue

527 + if ((name.endswith(".bias") or name.endswith("_bias"))

528 + and name not in params_dict):

529 + continue

530 + param = params_dict[name]

531 + weight_loader = param.weight_loader

532 + weight_loader(param,

533 + loaded_weight,

534 + name,

535 + shard_id=shard_id,

536 + expert_id=expert_id)

537 + break

538 + else:

539 + # Skip loading extra bias for GPTQ models.

540 + if ((name.endswith(".bias") or name.endswith("_bias"))

541 + and name not in params_dict):

542 + continue

543 + # Skip layers on other devices.

544 + if is_pp_missing_parameter(name, self):

545 + continue

546 +

547 + # Remapping the name of FP8 kv-scale.

548 + name = maybe_remap_kv_scale_name(name, params_dict)

549 + if name is None:

550 + continue

551 +

552 + # Handle Grok1-specific norm.scale naming

553 + if "norm.scale" in name:

554 + name = name.replace("scale", "weight")

555 +

556 + # Skip lm_head when tie_word_embeddings is True

557 + if "lm_head" in name and self.config.tie_word_embeddings:

558 + continue

559 +

560 + param = params_dict[name]

561 + weight_loader = getattr(param, "weight_loader",

562 + default_weight_loader)

563 + weight_loader(param, loaded_weight)

564 + loaded_params.add(name)

565 + return loaded_params

vllm/model_executor/models/registry.py

@@ -60,6 +60,7 @@ _TEXT_GENERATION_MODELS = {

60 "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),

61 "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),

62 "GritLM": ("gritlm", "GritLM"),

63 + "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),

64 "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),

65 "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),

66 "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),