@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c7645bce04

Remove grok model arch from vllm (#46706)

Signed-off-by: Xianbao QIAN <xianbao.qian@gmail.com>

Tiezhen WANG committed Jun 26, 2026 at 14:02 UTC c7645bce044be01c3c30ceead99388f01aa6d496

13 files changed +3 -1401

docs/configuration/optimization.md

+1 -1

@@ -297,7 +297,7 @@ The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't,

297 vLLM raises a clear `ImportError` at tokenizer load. The override applies to

298 any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`,

299 `deepseek_v32`, `deepseek_v4`, …). Models that don't use the HF

300 -fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag.

300 +fast tokenizer (`mistral`, `kimi_audio`) ignore the flag.

301

302 Tokenizer-bound workloads — long shared prefixes, bursty short prompts,

303 batch detokenization — see the largest wins. If your bottleneck is GPU

docs/models/supported_models.md

-5

@@ -414,8 +414,6 @@ th {

414 | `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | ✅︎ | ✅︎ |

415 | `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ |

416 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |

417 -| `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ |

418 -| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |

418 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |

419 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |

@@ -488,9 +486,6 @@ th {

486 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |

488

491 -!!! note

492 - Grok2 requires `tokenizer.tok.json` with `tiktoken` installed. You can optionally override MoE router renormalization with `moe_router_renormalize`.

493 -

489 Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it!

490

tests/models/language/generation/test_grok.py deleted

-43

@@ -1,43 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import pytest

4 -

5 -from ...utils import dummy_hf_overrides

6 -

7 -MODELS = ["xai-org/grok-2"]

8 -

9 -

10 -def _grok2_dummy_overrides(hf_config):

11 - hf_config = dummy_hf_overrides(hf_config, model_arch="Grok1ForCausalLM")

12 - text_config = hf_config.get_text_config()

13 - text_config.update(

14 - {

15 - "hidden_size": 256,

16 - "intermediate_size": 512,

17 - "moe_intermediate_size": 256,

18 - "num_attention_heads": 4,

19 - "num_key_value_heads": 2,

20 - "head_dim": 64,

21 - }

22 - )

23 - return hf_config

24 -

25 -

26 -@pytest.mark.parametrize("model", MODELS)

27 -def test_dummy_generate(vllm_runner, monkeypatch, model: str) -> None:

28 - with monkeypatch.context() as m:

29 - m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

30 - with vllm_runner(

31 - model,

32 - load_format="dummy",

33 - max_model_len=128,

34 - hf_overrides=_grok2_dummy_overrides,

35 - enforce_eager=True,

36 - ) as llm:

37 - prompt = "Hello from Grok-2"

38 - tokenizer = llm.get_llm().get_tokenizer()

39 - prompt_len = len(tokenizer.encode(prompt))

40 - outputs = llm.generate_greedy([prompt], max_tokens=1)

41 - output_ids, output_str = outputs[0]

42 - assert len(output_ids) > prompt_len

43 - assert output_str is not None

tests/models/registry.py

-4

@@ -319,10 +319,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

319 "GraniteMoeSharedForCausalLM": _HfExamplesInfo(

320 "ibm-research/moe-7b-1b-active-shared-experts"

321 ),

322 - "Grok1ModelForCausalLM": _HfExamplesInfo(

323 - "hpcai-tech/grok-1", trust_remote_code=True

324 - ),

325 - "Grok1ForCausalLM": _HfExamplesInfo("xai-org/grok-2", trust_remote_code=True),

322 "HrmTextForCausalLM": _HfExamplesInfo(

323 "sapientinc/HRM-Text-1B",

324 min_transformers_version="5.9.0",

tests/tokenizers_/test_basic.py

-5

@@ -9,7 +9,6 @@ from transformers import (

9 )

11 from vllm.tokenizers import TokenizerLike, get_tokenizer

12 -from vllm.tokenizers.grok2 import Grok2Tokenizer

12 from vllm.tokenizers.hf import HfTokenizer

13 from vllm.tokenizers.mistral import MistralTokenizer

@@ -35,10 +34,6 @@ def test_tokenizer_like_protocol():

34 assert isinstance(tokenizer, MistralTokenizer)

35 _assert_tokenizer_like(tokenizer)

38 - tokenizer = get_tokenizer("xai-org/grok-2", tokenizer_mode="grok2")

39 - assert isinstance(tokenizer, Grok2Tokenizer)

40 - _assert_tokenizer_like(tokenizer)

41 -

37 tokenizer = get_tokenizer("deepseek-ai/DeepSeek-V3", tokenizer_mode="deepseek_v32")

38 assert isinstance(tokenizer, HfTokenizer)

vllm/config/model.py

-2

@@ -601,8 +601,6 @@ class ModelConfig:

601 if self.tokenizer_mode == "auto":

602 if self.model_impl == "terratorch":

603 self.tokenizer_mode = "terratorch"

604 - elif arch == "Grok1ForCausalLM":

605 - self.tokenizer_mode = "grok2"

604 elif arch == "MoonshotKimiaForCausalLM":

605 self.tokenizer_mode = "kimi_audio"

606 elif arch == "DeepseekV32ForCausalLM":

vllm/model_executor/models/grok1.py deleted

-792

@@ -1,792 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -# Adapted from

4 -# https://github.com/ROCm/vllm/blob/cea7419f151cc50293a05b7fac8547f8f887c9f6/vllm/model_executor/models/grok1.py

7 -#

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

9 -# and OPT implementations in this library. It has been modified from its

10 -# original forms to accommodate minor architectural differences compared

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

12 -#

13 -# Licensed under the Apache License, Version 2.0 (the "License");

14 -# you may not use this file except in compliance with the License.

15 -# You may obtain a copy of the License at

16 -#

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

18 -#

19 -# Unless required by applicable law or agreed to in writing, software

20 -# distributed under the License is distributed on an "AS IS" BASIS,

21 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

22 -# See the License for the specific language governing permissions and

23 -# limitations under the License.

24 -"""Inference-only Grok (Grok1/Grok2) model."""

25 -

26 -import math

27 -from collections.abc import Iterable

28 -from itertools import islice

29 -from typing import Any

30 -

31 -import torch

32 -import torch.nn.functional as F

33 -from torch import nn

34 -

35 -from vllm.compilation.decorators import support_torch_compile

36 -from vllm.config import CacheConfig, VllmConfig

37 -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

38 -from vllm.logger import init_logger

39 -from vllm.model_executor.layers.activation import GeluAndMul

40 -from vllm.model_executor.layers.attention import Attention

41 -from vllm.model_executor.layers.fused_moe import (

42 - FusedMoE,

43 - fused_moe_make_expert_params_mapping,

44 -)

45 -from vllm.model_executor.layers.layernorm import RMSNorm

46 -from vllm.model_executor.layers.linear import (

47 - MergedColumnParallelLinear,

48 - QKVParallelLinear,

49 - ReplicatedLinear,

50 - RowParallelLinear,

51 -)

52 -from vllm.model_executor.layers.logits_processor import LogitsProcessor

53 -from vllm.model_executor.layers.quantization import QuantizationConfig

54 -from vllm.model_executor.layers.rotary_embedding import get_rope

55 -from vllm.model_executor.layers.vocab_parallel_embedding import (

56 - ParallelLMHead,

57 - VocabParallelEmbedding,

58 -)

59 -from vllm.model_executor.model_loader.weight_utils import (

60 - default_weight_loader,

61 - maybe_remap_kv_scale_name,

62 -)

63 -from vllm.sequence import IntermediateTensors

64 -

65 -from .interfaces import SupportsLoRA, SupportsPP

66 -from .utils import (

67 - AutoWeightsLoader,

68 - is_pp_missing_parameter,

69 - make_empty_intermediate_tensors_factory,

70 - make_layers,

71 - maybe_prefix,

72 -)

73 -

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

75 -DEFAULT_ATTN_OUTPUT_MULTIPLIER = 0.08838834764831845

76 -DEFAULT_OUTPUT_MULTIPLIER_SCALE = 0.5773502691896257

77 -DEFAULT_EMBEDDING_MULTIPLIER_SCALE = 78.38367176906169

78 -DEFAULT_ROUTER_LOGIT_SOFTCAP = 30.0

79 -

80 -logger = init_logger(__name__)

81 -

82 -

83 -def _get_num_experts(config) -> int:

84 - return getattr(config, "num_experts", getattr(config, "num_local_experts", 8))

85 -

86 -

87 -def _get_moe_intermediate_size(config) -> int:

88 - return getattr(config, "moe_intermediate_size", config.intermediate_size)

89 -

90 -

91 -def _get_grok_version(config) -> str:

92 - """Detect Grok version from HF config using multiple heuristics."""

93 - # Check for Grok2-specific attributes (both for robust detection)

94 - has_residual_moe = getattr(config, "residual_moe", False)

95 - has_moe_intermediate_size = hasattr(config, "moe_intermediate_size")

96 -

97 - if has_residual_moe or has_moe_intermediate_size:

98 - return "grok2"

99 -

100 - return "grok1" # Default to Grok1

101 -

102 -

103 -def _get_rope_parameters(config) -> dict[str, Any] | None:

104 - rope_parameters = getattr(config, "rope_parameters", None)

105 - if rope_parameters is None:

106 - rope_type = getattr(config, "rope_type", None)

107 - if rope_type is None:

108 - return None

109 - rope_parameters = {"rope_type": rope_type}

110 - rope_theta = getattr(config, "rope_theta", None)

111 - if rope_theta is not None:

112 - rope_parameters["rope_theta"] = rope_theta

113 - scaling_factor = getattr(config, "scaling_factor", None)

114 - if scaling_factor is not None:

115 - rope_parameters["factor"] = scaling_factor

116 - for name in (

117 - "original_max_position_embeddings",

118 - "extrapolation_factor",

119 - "attn_factor",

120 - "beta_fast",

121 - "beta_slow",

122 - ):

123 - value = getattr(config, name, None)

124 - if value is not None:

125 - rope_parameters[name] = value

126 -

127 - if rope_parameters.get("rope_type") == "original":

128 - rope_parameters = dict(rope_parameters)

129 - rope_parameters["rope_type"] = "default"

130 - return rope_parameters

131 -

132 -

133 -def _get_moe_renormalize(config) -> bool:

134 - explicit_value = getattr(

135 - config, "moe_router_renormalize", getattr(config, "moe_renormalize", None)

136 - )

137 - if explicit_value is not None:

138 - return bool(explicit_value)

139 - return not getattr(config, "residual_moe", False)

140 -

141 -

142 -class Grok1MLP(nn.Module):

143 - def __init__(

144 - self,

145 - hidden_size: int,

146 - intermediate_size: int,

147 - quant_config: QuantizationConfig | None = None,

148 - prefix: str = "",

149 - ) -> None:

150 - super().__init__()

151 - self.gate_up_proj = MergedColumnParallelLinear(

152 - input_size=hidden_size,

153 - output_sizes=[intermediate_size] * 2,

154 - bias=False,

155 - quant_config=quant_config,

156 - prefix=f"{prefix}.gate_up_proj",

157 - )

158 - self.down_proj = RowParallelLinear(

159 - input_size=intermediate_size,

160 - output_size=hidden_size,

161 - bias=False,

162 - quant_config=quant_config,

163 - prefix=f"{prefix}.down_proj",

164 - )

165 - self.act_fn = GeluAndMul()

166 -

167 - def forward(self, x: torch.Tensor) -> torch.Tensor:

168 - x, _ = self.gate_up_proj(x)

169 - x = self.act_fn(x)

170 - x, _ = self.down_proj(x)

171 - return x

172 -

173 -

174 -class Grok1MoE(nn.Module):

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

176 - across all ranks.

177 -

178 - Each expert's weights are sharded across all ranks and a fused MoE

179 - kernel is used for the forward pass, and finally we reduce the outputs

180 - across ranks.

181 - """

182 -

183 - def __init__(

184 - self,

185 - num_experts: int,

186 - top_k: int,

187 - hidden_size: int,

188 - intermediate_size: int,

189 - router_logit_soft_cap: float = 0.0,

190 - params_dtype: torch.dtype | None = None,

191 - quant_config: QuantizationConfig | None = None,

192 - tp_size: int | None = None,

193 - renormalize: bool = False,

194 - prefix: str = "",

195 - ):

196 - super().__init__()

197 - self.hidden_size = hidden_size

198 -

199 - # Gate always runs at half / full precision for now.

200 - self.gate = ReplicatedLinear(

201 - hidden_size,

202 - num_experts,

203 - bias=False,

204 - params_dtype=params_dtype,

205 - quant_config=None,

206 - prefix=f"{prefix}.gate",

207 - )

208 -

209 - self.experts = FusedMoE(

210 - num_experts=num_experts,

211 - top_k=top_k,

212 - hidden_size=hidden_size,

213 - intermediate_size=intermediate_size,

214 - params_dtype=params_dtype,

215 - renormalize=renormalize,

216 - quant_config=quant_config,

217 - tp_size=tp_size,

218 - activation="gelu",

219 - prefix=f"{prefix}.experts",

220 - )

221 - self.router_logit_soft_cap = router_logit_soft_cap

222 -

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

224 - # NOTE: hidden_states can have either 1D or 2D shape.

225 - orig_shape = hidden_states.shape

226 - hidden_states = hidden_states.view(-1, self.hidden_size)

227 - # router_logits: (num_tokens, n_experts)

228 - router_logits, _ = self.gate(hidden_states)

229 - if self.router_logit_soft_cap > 0:

230 - router_logits = self.router_logit_soft_cap * F.tanh(

231 - router_logits / self.router_logit_soft_cap

232 - )

233 - final_hidden_states = self.experts(hidden_states, router_logits)

234 - return final_hidden_states.view(orig_shape)

235 -

236 -

237 -class Grok1Attention(nn.Module):

238 - def __init__(

239 - self,

240 - hidden_size: int,

241 - num_heads: int,

242 - num_kv_heads: int,

243 - max_position: int = 4096 * 32,

244 - rope_parameters: dict[str, Any] | None = None,

245 - cache_config: CacheConfig | None = None,

246 - quant_config: QuantizationConfig | None = None,

247 - prefix: str = "",

248 - config=None, # Added config parameter

249 - ) -> None:

250 - super().__init__()

251 - self.hidden_size = hidden_size

252 - self.config = config # Store config reference

253 - tp_size = get_tensor_model_parallel_world_size()

254 - self.total_num_heads = num_heads

255 - assert self.total_num_heads % tp_size == 0

256 - self.num_heads = self.total_num_heads // tp_size

257 - self.total_num_kv_heads = num_kv_heads

258 - if self.total_num_kv_heads >= tp_size:

259 - # Number of KV heads is greater than TP size, so we partition

260 - # the KV heads across multiple tensor parallel GPUs.

261 - assert self.total_num_kv_heads % tp_size == 0

262 - else:

263 - # Number of KV heads is less than TP size, so we replicate

264 - # the KV heads across multiple tensor parallel GPUs.

265 - assert tp_size % self.total_num_kv_heads == 0

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

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

268 - self.q_size = self.num_heads * self.head_dim

269 - self.kv_size = self.num_kv_heads * self.head_dim

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

271 -

272 - self.qkv_proj = QKVParallelLinear(

273 - hidden_size,

274 - self.head_dim,

275 - self.total_num_heads,

276 - self.total_num_kv_heads,

277 - bias=False,

278 - quant_config=quant_config,

279 - prefix=f"{prefix}.qkv_proj",

280 - )

281 - self.o_proj = RowParallelLinear(

282 - self.total_num_heads * self.head_dim,

283 - hidden_size,

284 - bias=False,

285 - quant_config=quant_config,

286 - prefix=f"{prefix}.o_proj",

287 - )

288 - self.rotary_emb = get_rope(

289 - self.head_dim,

290 - max_position=max_position,

291 - rope_parameters=rope_parameters,

292 - is_neox_style=True,

293 - )

294 -

295 - attn_logits_soft_cap = max(getattr(config, "attn_logit_softcapping", 30.0), 0.0)

296 - attn_logit_softcapping_method = getattr(

297 - config, "attn_logit_softcapping_method", None

298 - )

299 - if attn_logit_softcapping_method not in (None, "tanh"):

300 - logger.warning_once(

301 - "Grok attention logit softcapping method '%s' is not "

302 - "supported; falling back to default behavior.",

303 - attn_logit_softcapping_method,

304 - )

305 -

306 - self.attn = Attention(

307 - self.num_heads,

308 - self.head_dim,

309 - self.scaling,

310 - num_kv_heads=self.num_kv_heads,

311 - cache_config=cache_config,

312 - quant_config=quant_config,

313 - logits_soft_cap=attn_logits_soft_cap,

314 - prefix=f"{prefix}.attn",

315 - )

316 - self.attn_multiplier = (

317 - getattr(self.config, "attn_output_multiplier", 1.0) if self.config else 1.0

318 - )

319 -

320 - def forward(

321 - self,

322 - positions: torch.Tensor,

323 - hidden_states: torch.Tensor,

324 - ) -> torch.Tensor:

325 - qkv, _ = self.qkv_proj(hidden_states)

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

327 - q, k = self.rotary_emb(positions, q, k)

328 - attn_output = self.attn(q, k, v)

329 - output, _ = self.o_proj(attn_output)

330 - output *= self.attn_multiplier

331 - return output

332 -

333 -

334 -class Grok1DecoderLayer(nn.Module):

335 - def __init__(

336 - self,

337 - config,

338 - cache_config: CacheConfig | None = None,

339 - quant_config: QuantizationConfig | None = None,

340 - prefix: str = "",

341 - ) -> None:

342 - super().__init__()

343 - self.hidden_size = config.hidden_size

344 - # Check for fp8 quantization

345 - self.use_fp8 = False

346 - if quant_config is not None:

347 - self.use_fp8 = getattr(quant_config, "is_fp8_w8a8", lambda: False)()

348 - if not self.use_fp8 and hasattr(quant_config, "is_fp8"):

349 - self.use_fp8 = quant_config.is_fp8

350 -

351 - self.attn = Grok1Attention(

352 - hidden_size=self.hidden_size,

353 - num_heads=config.num_attention_heads,

354 - max_position=config.max_position_embeddings,

355 - num_kv_heads=config.num_key_value_heads,

356 - rope_parameters=_get_rope_parameters(config),

357 - cache_config=cache_config,

358 - quant_config=quant_config,

359 - prefix=f"{prefix}.attn",

360 - config=config,

361 - ) # Pass config to Grok1Attention

362 -

363 - num_experts = _get_num_experts(config)

364 - num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)

365 - moe_intermediate_size = _get_moe_intermediate_size(config)

366 - moe_renormalize = _get_moe_renormalize(config)

367 -

368 - self.moe_block = Grok1MoE(

369 - num_experts=num_experts,

370 - top_k=num_experts_per_tok,

371 - hidden_size=config.hidden_size,

372 - intermediate_size=moe_intermediate_size,

373 - router_logit_soft_cap=max(

374 - getattr(

375 - config,

376 - "router_logit_softcapping",

377 - DEFAULT_ROUTER_LOGIT_SOFTCAP,

378 - ),

379 - 0.0,

380 - ),

381 - quant_config=quant_config,

382 - renormalize=moe_renormalize,

383 - prefix=f"{prefix}.moe_block",

384 - )

385 - self.residual_moe = getattr(config, "residual_moe", False)

386 - self.residual_moe_scale = 1.0 / math.sqrt(2.0)

387 -

388 - self.pre_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

389 - self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

390 - self.pre_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

391 - self.post_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

392 - self.mlp = None

393 - if self.residual_moe:

394 - self.mlp = Grok1MLP(

395 - hidden_size=config.hidden_size,

396 - intermediate_size=config.intermediate_size,

397 - quant_config=quant_config,

398 - prefix=f"{prefix}.mlp",

399 - )

400 -

401 - def forward(

402 - self,

403 - positions: torch.Tensor,

404 - hidden_states: torch.Tensor,

405 - residual: torch.Tensor | None,

406 - ) -> tuple[torch.Tensor, torch.Tensor]:

407 - # Self Attention

408 - if residual is None:

409 - residual = hidden_states

410 - hidden_states = self.pre_attn_norm(hidden_states)

411 - else:

412 - hidden_states, residual = self.pre_attn_norm(hidden_states, residual)

413 -

414 - hidden_states = self.attn(

415 - positions=positions,

416 - hidden_states=hidden_states,

417 - )

418 -

419 - # Post attention normalization

420 - hidden_states = self.post_attn_norm(hidden_states)

421 -

422 - # MoE block with normalization

423 - hidden_states, residual = self.pre_moe_norm(hidden_states, residual)

424 - if self.residual_moe:

425 - assert self.mlp is not None

426 - hidden_states = (

427 - self.moe_block(hidden_states) + self.mlp(hidden_states)

428 - ) * self.residual_moe_scale

429 - else:

430 - hidden_states = self.moe_block(hidden_states)

431 - hidden_states = self.post_moe_norm(hidden_states)

432 -

433 - return hidden_states, residual

434 -

435 -

436 -@support_torch_compile

437 -class Grok1Model(nn.Module):

438 - def __init__(

439 - self,

440 - *,

441 - vllm_config: VllmConfig,

442 - prefix: str = "",

443 - ckpt_gate_proj_name: str = "linear",

444 - ckpt_down_proj_name: str = "linear_1",

445 - ckpt_up_proj_name: str = "linear_v",

446 - weight_name_remapping: dict[str, str] | None = None,

447 - ):

448 - super().__init__()

449 -

450 - config = vllm_config.model_config.hf_config

451 - cache_config = vllm_config.cache_config

452 - quant_config = vllm_config.quant_config

453 -

454 - self.config = config

455 - self.quant_config = quant_config

456 -

457 - # Store expert naming for weight loading

458 - self.ckpt_gate_proj_name = ckpt_gate_proj_name

459 - self.ckpt_down_proj_name = ckpt_down_proj_name

460 - self.ckpt_up_proj_name = ckpt_up_proj_name

461 - self.weight_name_remapping = weight_name_remapping or {}

462 -

463 - self.vocab_size = config.vocab_size

464 -

465 - self.embedding_multiplier_scale = getattr(

466 - config, "embedding_multiplier_scale", DEFAULT_EMBEDDING_MULTIPLIER_SCALE

467 - )

468 -

469 - self.embed_tokens = VocabParallelEmbedding(

470 - self.vocab_size,

471 - config.hidden_size,

472 - quant_config=quant_config,

473 - )

474 -

475 - self.start_layer, self.end_layer, self.layers = make_layers(

476 - config.num_hidden_layers,

477 - lambda prefix: Grok1DecoderLayer(

478 - config, cache_config, quant_config=quant_config, prefix=prefix

479 - ),

480 - prefix=f"{prefix}.layers",

481 - )

482 -

483 - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

484 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

485 - ["hidden_states", "residual"], config.hidden_size

486 - )

487 -

488 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

489 - hidden_states = self.embed_tokens(input_ids)

490 - hidden_states = hidden_states * self.embedding_multiplier_scale

491 - return hidden_states

492 -

493 - def forward(

494 - self,

495 - input_ids: torch.Tensor | None,

496 - positions: torch.Tensor,

497 - intermediate_tensors: IntermediateTensors | None,

498 - inputs_embeds: torch.Tensor | None = None,

499 - ) -> torch.Tensor | IntermediateTensors:

500 - if get_pp_group().is_first_rank:

501 - if inputs_embeds is not None:

502 - hidden_states = inputs_embeds

503 - else:

504 - hidden_states = self.embed_input_ids(input_ids)

505 - residual = None

506 - else:

507 - assert intermediate_tensors is not None

508 - hidden_states = intermediate_tensors["hidden_states"]

509 - residual = intermediate_tensors["residual"]

510 -

511 - for layer in islice(self.layers, self.start_layer, self.end_layer):

512 - hidden_states, residual = layer(positions, hidden_states, residual)

513 -

514 - if not get_pp_group().is_last_rank:

515 - return IntermediateTensors(

516 - {"hidden_states": hidden_states, "residual": residual}

517 - )

518 -

519 - hidden_states, _ = self.norm(hidden_states, residual)

520 - return hidden_states

521 -

522 - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

523 - # Map expert parameter names to standard names

524 - num_experts = _get_num_experts(self.config)

525 - return fused_moe_make_expert_params_mapping(

526 - self,

527 - ckpt_gate_proj_name=self.ckpt_gate_proj_name,

528 - ckpt_down_proj_name=self.ckpt_down_proj_name,

529 - ckpt_up_proj_name=self.ckpt_up_proj_name,

530 - num_experts=num_experts,

531 - )

532 -

533 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

534 - stacked_params_mapping = [

535 - # (param_name, shard_name, shard_id)

536 - ("qkv_proj", "q_proj", "q"),

537 - ("qkv_proj", "k_proj", "k"),

538 - ("qkv_proj", "v_proj", "v"),

539 - ("mlp.gate_up_proj", "mlp.gate_proj", 0),

540 - ("mlp.gate_up_proj", "mlp.up_proj", 1),

541 - ]

542 -

543 - params_dict = dict(self.named_parameters())

544 - loaded_params: set[str] = set()

545 - expert_params_mapping = self.get_expert_mapping()

546 - for name, loaded_weight in weights:

547 - # Apply version-specific weight name remapping

548 - for old_pattern, new_pattern in self.weight_name_remapping.items():

549 - if old_pattern in name:

550 - name = name.replace(old_pattern, new_pattern)

551 -

552 - for param_name, weight_name, shard_id in stacked_params_mapping:

553 - if weight_name not in name:

554 - continue

555 - name = name.replace(weight_name, param_name)

556 - # Skip loading extra bias for GPTQ models.

557 - if (

558 - name.endswith(".bias") or name.endswith("_bias")

559 - ) and name not in params_dict:

560 - continue

561 - # Skip layers on other devices.

562 - if is_pp_missing_parameter(name, self):

563 - continue

564 - if name.endswith("scale"):

565 - # Remapping the name of FP8 kv-scale.

566 - name = maybe_remap_kv_scale_name(name, params_dict)

567 - if name is None:

568 - continue

569 - if name not in params_dict:

570 - continue

571 - param = params_dict[name]

572 - weight_loader = param.weight_loader

573 - weight_loader(param, loaded_weight, shard_id)

574 - break

575 - else:

576 - for mapping in expert_params_mapping:

577 - param_name, weight_name, expert_id, shard_id = mapping

578 - if weight_name not in name:

579 - continue

580 - name = name.replace(weight_name, param_name)

581 - # Skip layers on other devices.

582 - if is_pp_missing_parameter(name, self):

583 - continue

584 - if (

585 - name.endswith(".bias") or name.endswith("_bias")

586 - ) and name not in params_dict:

587 - continue

588 - param = params_dict[name]

589 - weight_loader = param.weight_loader

590 - weight_loader(

591 - param,

592 - loaded_weight,

593 - name,

594 - shard_id=shard_id,

595 - expert_id=expert_id,

596 - )

597 - break

598 - else:

599 - # Skip loading extra bias for GPTQ models.

600 - if (

601 - name.endswith(".bias") or name.endswith("_bias")

602 - ) and name not in params_dict:

603 - continue

604 - # Skip layers on other devices.

605 - if is_pp_missing_parameter(name, self):

606 - continue

607 -

608 - # Remapping the name of FP8 kv-scale.

609 - name = maybe_remap_kv_scale_name(name, params_dict)

610 - if name is None:

611 - continue

612 -

613 - # Handle Grok1-specific norm.scale naming

614 - if "norm.scale" in name:

615 - name = name.replace("scale", "weight")

616 -

617 - if name not in params_dict:

618 - continue

619 - param = params_dict[name]

620 - weight_loader = getattr(

621 - param, "weight_loader", default_weight_loader

622 - )

623 - weight_loader(param, loaded_weight)

624 - loaded_params.add(name)

625 - return loaded_params

626 -

627 -

628 -class GrokBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

629 - """Base class for Grok models with shared logic."""

630 -

631 - fall_back_to_pt_during_load = False

632 -

633 - # Subclasses should override these

634 - packed_modules_mapping = {

635 - "qkv_proj": [

636 - "q_proj",

637 - "k_proj",

638 - "v_proj",

639 - ],

640 - }

641 -

642 - # Expert weight naming - subclasses override these

643 - ckpt_gate_proj_name: str = "linear"

644 - ckpt_down_proj_name: str = "linear_1"

645 - ckpt_up_proj_name: str = "linear_v"

646 -

647 - def get_weight_name_remapping(self) -> dict[str, str]:

648 - """Return weight name remapping for this version. Override in subclasses."""

649 - return {}

650 -

651 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

652 - super().__init__()

653 -

654 - config = vllm_config.model_config.hf_config

655 - quant_config = vllm_config.quant_config

656 -

657 - self.config = config

658 - self.quant_config = quant_config

659 -

660 - self.model = Grok1Model(

661 - vllm_config=vllm_config,

662 - prefix=maybe_prefix(prefix, "model"),

663 - ckpt_gate_proj_name=self.ckpt_gate_proj_name,

664 - ckpt_down_proj_name=self.ckpt_down_proj_name,

665 - ckpt_up_proj_name=self.ckpt_up_proj_name,

666 - weight_name_remapping=self.get_weight_name_remapping(),

667 - )

668 -

669 - self.lm_head = ParallelLMHead(

670 - config.vocab_size,

671 - config.hidden_size,

672 - quant_config=quant_config,

673 - prefix=maybe_prefix(prefix, "lm_head"),

674 - )

675 -

676 - if self.config.tie_word_embeddings:

677 - self.lm_head.weight = self.model.embed_tokens.weight

678 -

679 - self.output_multiplier_scale = getattr(

680 - config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE

681 - )

682 - self.logits_processor = LogitsProcessor(

683 - config.vocab_size,

684 - scale=self.output_multiplier_scale,

685 - soft_cap=getattr(config, "final_logit_softcapping", None),

686 - )

687 -

688 - self.make_empty_intermediate_tensors = (

689 - self.model.make_empty_intermediate_tensors

690 - )

691 -

692 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

693 - return self.model.embed_input_ids(input_ids)

694 -

695 - def forward(

696 - self,

697 - input_ids: torch.Tensor | None,

698 - positions: torch.Tensor,

699 - intermediate_tensors: IntermediateTensors | None = None,

700 - inputs_embeds: torch.Tensor | None = None,

701 - ) -> torch.Tensor | IntermediateTensors:

702 - hidden_states = self.model(

703 - input_ids, positions, intermediate_tensors, inputs_embeds

704 - )

705 - return hidden_states

706 -

707 - def compute_logits(

708 - self,

709 - hidden_states: torch.Tensor,

710 - ) -> torch.Tensor | None:

711 - logits = self.logits_processor(self.lm_head, hidden_states)

712 - return logits

713 -

714 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

715 - # Skip lm_head when tie_word_embeddings is True

716 - skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else None

717 -

718 - loader = AutoWeightsLoader(

719 - self,

720 - skip_prefixes=skip_prefixes,

721 - )

722 - return loader.load_weights(weights)

723 -

724 - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

725 - return self.model.get_expert_mapping()

726 -

727 -

728 -class Grok1ForCausalLM(GrokBaseForCausalLM):

729 - """Grok1-specific implementation."""

730 -

731 - # Grok1 expert weight naming

732 - ckpt_gate_proj_name = "linear"

733 - ckpt_down_proj_name = "linear_1"

734 - ckpt_up_proj_name = "linear_v"

735 -

736 - def get_weight_name_remapping(self) -> dict[str, str]:

737 - # Grok1 uses standard naming, no remapping needed

738 - return {}

739 -

740 -

741 -class Grok2ForCausalLM(GrokBaseForCausalLM):

742 - """Grok2-specific implementation."""

743 -

744 - # Grok2 has additional packed modules for MLP

745 - packed_modules_mapping = {

746 - "qkv_proj": [

747 - "q_proj",

748 - "k_proj",

749 - "v_proj",

750 - ],

751 - "gate_up_proj": [

752 - "gate_proj",

753 - "up_proj",

754 - ],

755 - }

756 -

757 - # Grok2 expert weight naming

758 - ckpt_gate_proj_name = "w1"

759 - ckpt_down_proj_name = "w2"

760 - ckpt_up_proj_name = "w3"

761 -

762 - def get_weight_name_remapping(self) -> dict[str, str]:

763 - # Grok2 checkpoint uses different naming conventions

764 - return {

765 - ".self_attn.": ".attn.",

766 - ".block_sparse_moe.": ".moe_block.",

767 - }

768 -

769 -

770 -# Version dispatch mapping

771 -_GROK_VERSIONS: dict[str, type[GrokBaseForCausalLM]] = {

772 - "grok1": Grok1ForCausalLM,

773 - "grok2": Grok2ForCausalLM,

774 -}

775 -

776 -

777 -class GrokForCausalLM(GrokBaseForCausalLM):

778 - """Factory class that dispatches to version-specific implementation."""

779 -

780 - def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):

781 - config = vllm_config.model_config.hf_config

782 - version = _get_grok_version(config)

783 -

784 - instance_cls = _GROK_VERSIONS.get(version)

785 - if instance_cls is None:

786 - raise ValueError(f"Unsupported Grok version: {version}")

787 -

788 - # Merge class attributes for LoRA/quantization compatibility

789 - cls.packed_modules_mapping = dict(cls.packed_modules_mapping)

790 - cls.packed_modules_mapping.update(instance_cls.packed_modules_mapping)

791 -

792 - return instance_cls(vllm_config=vllm_config, prefix=prefix)

vllm/model_executor/models/registry.py

+2 -2

@@ -124,8 +124,6 @@ _TEXT_GENERATION_MODELS = {

124 "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),

125 "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),

126 "GritLM": ("gritlm", "GritLM"),

127 - "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),

128 - "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),

127 "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"),

128 "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),

129 "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),

@@ -731,6 +729,8 @@ _PREVIOUSLY_SUPPORTED_MODELS = {

729 "BaichuanForCausalLM": "0.23.0",

730 "AquilaModel": "0.24.0",

731 "AquilaForCausalLM": "0.24.0",

732 + "Grok1ModelForCausalLM": "0.24.0",

733 + "Grok1ForCausalLM": "0.24.0",

734 }

735

736 _OOT_SUPPORTED_MODELS = {

vllm/model_executor/models/transformers/moe.py

-3

@@ -180,7 +180,6 @@ class MoEMixin(MixtureOfExperts):

180 # (ckpt_gate_proj_name, ckpt_down_proj_name, ckpt_up_proj_name)

181 ("gate_proj", "down_proj", "up_proj"), # Most common MoE style

182 ("w1", "w2", "w3"), # Granite, Mixtral, Phi MoE style

183 - ("linear", "linear_1", "linear_v"), # Grok1 style

183 ]

184 num_experts = self.model_config.get_num_experts()

185 num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts

@@ -238,8 +237,6 @@ class MoEMixin(MixtureOfExperts):

237 wrapped_arch = self.config.architectures[0].lower()

238 if "gptoss" in wrapped_arch:

239 activation = "swigluoai"

241 - elif "grok1" in wrapped_arch:

242 - activation = "gelu"

240

241 # Expert mapping for `AutoWeightsLoader`

242 expert_mapping = self.get_expert_mapping()

vllm/renderers/grok2.py deleted

-90

@@ -1,90 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -from vllm.config import VllmConfig

5 -from vllm.entrypoints.chat_utils import (

6 - ChatCompletionMessageParam,

7 - ConversationMessage,

8 - parse_chat_messages,

9 - parse_chat_messages_async,

10 -)

11 -from vllm.logger import init_logger

12 -from vllm.tokenizers.grok2 import Grok2Tokenizer

13 -from vllm.utils.async_utils import make_async

14 -

15 -from .base import BaseRenderer

16 -from .inputs import DictPrompt

17 -from .inputs.preprocess import parse_dec_only_prompt

18 -from .params import ChatParams

19 -

20 -logger = init_logger(__name__)

21 -

22 -

23 -class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):

24 - def __init__(

25 - self,

26 - config: VllmConfig,

27 - tokenizer: Grok2Tokenizer | None,

28 - ) -> None:

29 - super().__init__(config, tokenizer)

30 -

31 - self._apply_chat_template_async = make_async(

32 - self._apply_chat_template, executor=self._executor

33 - )

34 -

35 - def _apply_chat_template(self, *args, **kwargs):

36 - return self.get_tokenizer().apply_chat_template(*args, **kwargs)

37 -

38 - def render_messages(

39 - self,

40 - messages: list[ChatCompletionMessageParam],

41 - params: ChatParams,

42 - ) -> tuple[list[ConversationMessage], DictPrompt]:

43 - conversation, mm_data, mm_uuids = parse_chat_messages(

44 - messages,

45 - self.model_config,

46 - content_format="string",

47 - media_io_kwargs=params.media_io_kwargs,

48 - mm_processor_kwargs=params.mm_processor_kwargs,

49 - )

50 -

51 - prompt_raw = self._apply_chat_template(

52 - conversation=conversation,

53 - messages=messages,

54 - **params.get_apply_chat_template_kwargs(),

55 - )

56 -

57 - prompt = parse_dec_only_prompt(prompt_raw)

58 - if mm_data is not None:

59 - prompt["multi_modal_data"] = mm_data

60 - if mm_uuids is not None:

61 - prompt["multi_modal_uuids"] = mm_uuids

62 -

63 - return conversation, prompt

64 -

65 - async def render_messages_async(

66 - self,

67 - messages: list[ChatCompletionMessageParam],

68 - params: ChatParams,

69 - ) -> tuple[list[ConversationMessage], DictPrompt]:

70 - conversation, mm_data, mm_uuids = await parse_chat_messages_async(

71 - messages,

72 - self.model_config,

73 - content_format="string",

74 - media_io_kwargs=params.media_io_kwargs,

75 - mm_processor_kwargs=params.mm_processor_kwargs,

76 - )

77 -

78 - prompt_raw = await self._apply_chat_template_async(

79 - conversation=conversation,

80 - messages=messages,

81 - **params.get_apply_chat_template_kwargs(),

82 - )

83 -

84 - prompt = parse_dec_only_prompt(prompt_raw)

85 - if mm_data is not None:

86 - prompt["multi_modal_data"] = mm_data

87 - if mm_uuids is not None:

88 - prompt["multi_modal_uuids"] = mm_uuids

89 -

90 - return conversation, prompt

vllm/renderers/registry.py

-1

@@ -22,7 +22,6 @@ logger = init_logger(__name__)

22 _VLLM_RENDERERS = {

23 "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"),

24 "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"),

25 - "grok2": ("grok2", "Grok2Renderer"),

25 "hf": ("hf", "HfRenderer"),

26 "kimi_audio": ("hf", "HfRenderer"),

27 "mistral": ("mistral", "MistralRenderer"),

vllm/tokenizers/grok2.py deleted

-452

@@ -1,452 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Tokenizer for Grok-2 .tok.json format."""

4 -

5 -import functools

6 -import json

7 -from collections.abc import Collection, Sequence, Set

8 -from pathlib import Path

9 -from typing import Any, Literal, overload

10 -

11 -from huggingface_hub.utils import (

12 - EntryNotFoundError,

13 - HfHubHTTPError,

14 - RepositoryNotFoundError,

15 - RevisionNotFoundError,

16 -)

17 -from transformers import BatchEncoding

18 -from transformers.utils import chat_template_utils as hf_chat_utils

19 -

20 -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

21 -from vllm.logger import init_logger

22 -from vllm.transformers_utils.repo_utils import hf_api

23 -

24 -from .protocol import TokenizerLike

25 -

26 -logger = init_logger(__name__)

27 -

28 -PAD = "<|pad|>"

29 -EOS = "<|eos|>"

30 -SEP = "<|separator|>"

31 -RESERVED_TOKEN_TEXTS = [f"<|reserved_{i}|>" for i in range(3, 128)]

32 -CONTROL_TOKEN_TEXTS = [f"<|control{i}|>" for i in range(1, 705)]

33 -DEFAULT_SPECIAL_TOKENS = [PAD, SEP, EOS]

34 -DEFAULT_CONTROL_TOKENS = {"pad": PAD, "sep": SEP, "eos": EOS}

35 -DEFAULT_CHAT_TEMPLATE = (

36 - "{% for message in messages %}"

37 - "{% if message['role'] == 'user' %}"

38 - "{{ 'Human: ' + message['content'].strip() + '<|separator|>\\n\\n' }}"

39 - "{% elif message['role'] == 'system' %}"

40 - "{{ 'System: ' + message['content'].strip() + '<|separator|>\\n\\n' }}"

41 - "{% elif message['role'] == 'assistant' %}"

42 - "{{ 'Assistant: ' + message['content'] + '<|separator|>\\n\\n' }}"

43 - "{% endif %}"

44 - "{% endfor %}"

45 - "{% if add_generation_prompt %}"

46 - "{{ 'Assistant:' }}"

47 - "{% endif %}"

48 -)

49 -

50 -# Default + separate each single digit.

51 -PAT_STR_B = (

52 - r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|"""

53 - r""" ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""

54 -)

55 -

56 -

57 -def _maybe_load_tokenizer_config(

58 - model_path: Path,

59 - *,

60 - repo_id: str | None,

61 - revision: str | None,

62 - download_dir: str | None,

63 -) -> dict[str, Any]:

64 - config_path = model_path / "tokenizer_config.json"

65 - if config_path.is_file():

66 - with config_path.open("r", encoding="utf-8") as f:

67 - return json.load(f)

68 -

69 - if repo_id is None:

70 - return {}

71 -

72 - try:

73 - config_file = hf_api().hf_hub_download(

74 - repo_id=repo_id,

75 - filename="tokenizer_config.json",

76 - revision=revision,

77 - cache_dir=download_dir,

78 - )

79 - except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError):

80 - # If the repo, revision, or file does not exist, fall back silently.

81 - return {}

82 - except HfHubHTTPError as exc:

83 - logger.warning(

84 - "Failed to download tokenizer_config.json from %s. "

85 - "This may be due to a network or authentication issue. "

86 - "The default chat template will be used. Error: %s",

87 - repo_id,

88 - exc,

89 - )

90 - return {}

91 -

92 - try:

93 - with Path(config_file).open("r", encoding="utf-8") as f:

94 - return json.load(f)

95 - except json.JSONDecodeError as exc:

96 - logger.warning(

97 - "Failed to parse tokenizer_config.json. "

98 - "The default chat template will be used. Error: %s",

99 - exc,

100 - )

101 - return {}

102 - except OSError as exc:

103 - logger.warning(

104 - "Failed to open tokenizer_config.json. "

105 - "The default chat template will be used. Error: %s",

106 - exc,

107 - )

108 - return {}

109 -

110 -

111 -def _load_tiktoken_encoding(

112 - vocab_file: Path,

113 -) -> tuple[Any, dict[str, int]]:

114 - try:

115 - import tiktoken

116 - except ImportError as exc:

117 - raise ImportError("Grok-2 tokenizer requires the `tiktoken` package.") from exc

118 -

119 - with vocab_file.open("rb") as f:

120 - xtok_dict = json.load(f)

121 -

122 - mergeable_ranks = {

123 - bytes(item["bytes"]): item["token"]

124 - for item in xtok_dict.get("regular_tokens", [])

125 - }

126 - special_tokens = {

127 - bytes(item["bytes"]).decode("utf-8", errors="replace"): item["token"]

128 - for item in xtok_dict.get("special_tokens", [])

129 - }

130 -

131 - if xtok_dict.get("word_split") == "V1":

132 - pat_str = PAT_STR_B

133 - else:

134 - raise ValueError(f"Unknown word_split: {xtok_dict.get('word_split')!r}")

135 -

136 - pat_str = xtok_dict.get("pat_str", pat_str)

137 -

138 - kwargs = {

139 - "name": str(vocab_file),

140 - "pat_str": pat_str,

141 - "mergeable_ranks": mergeable_ranks,

142 - "special_tokens": special_tokens,

143 - }

144 -

145 - if "vocab_size" in xtok_dict:

146 - kwargs["explicit_n_vocab"] = xtok_dict["vocab_size"]

147 -

148 - tokenizer = tiktoken.Encoding(**kwargs)

149 -

150 - default_allowed_special: set[str] | None = None

151 - if "default_allowed_special" in xtok_dict:

152 - default_allowed_special = {

153 - bytes(bytes_list).decode("utf-8", errors="replace")

154 - for bytes_list in xtok_dict["default_allowed_special"]

155 - }

156 -

157 - tokenizer._default_allowed_special = default_allowed_special or set()

158 - tokenizer._control_tokens = DEFAULT_CONTROL_TOKENS

159 -

160 - def encode_patched(

161 - self,

162 - text: str,

163 - *,

164 - allowed_special: Literal["all"] | Set[str] = set(),

165 - disallowed_special: Literal["all"] | Collection[str] = "all",

166 - ) -> list[int]:

167 - del disallowed_special

168 - if isinstance(allowed_special, set):

169 - allowed_special |= self._default_allowed_special

170 - return tiktoken.Encoding.encode(

171 - self,

172 - text,

173 - allowed_special=allowed_special,

174 - disallowed_special=(),

175 - )

176 -

177 - tokenizer.encode = functools.partial(encode_patched, tokenizer)

178 - tokenizer._default_allowed_special |= set(DEFAULT_CONTROL_TOKENS.values())

179 - tokenizer._default_allowed_special |= set(

180 - CONTROL_TOKEN_TEXTS + RESERVED_TOKEN_TEXTS

181 - )

182 -

183 - return tokenizer, special_tokens

184 -

185 -

186 -class Grok2Tokenizer(TokenizerLike):

187 - @classmethod

188 - def from_pretrained(

189 - cls,

190 - path_or_repo_id: str | Path,

191 - *args,

192 - trust_remote_code: bool = False,

193 - revision: str | None = None,

194 - download_dir: str | None = None,

195 - **kwargs,

196 - ) -> "Grok2Tokenizer":

197 - if args:

198 - logger.debug_once("Ignoring extra positional args for Grok2Tokenizer.")

199 -

200 - path = Path(path_or_repo_id)

201 - if path.is_file():

202 - vocab_file = path

203 - model_path = path.parent

204 - repo_id = None

205 - elif path.is_dir():

206 - vocab_file = path / "tokenizer.tok.json"

207 - model_path = path

208 - repo_id = None

209 - else:

210 - vocab_file = Path(

211 - hf_api().hf_hub_download(

212 - repo_id=str(path_or_repo_id),

213 - filename="tokenizer.tok.json",

214 - revision=revision,

215 - cache_dir=download_dir,

216 - )

217 - )

218 - model_path = vocab_file.parent

219 - repo_id = str(path_or_repo_id)

220 -

221 - if not vocab_file.is_file():

222 - raise FileNotFoundError(f"tokenizer.tok.json not found at {vocab_file}.")

223 -

224 - config = _maybe_load_tokenizer_config(

225 - model_path,

226 - repo_id=repo_id,

227 - revision=revision,

228 - download_dir=download_dir,

229 - )

230 -

231 - return cls(

232 - vocab_file=vocab_file,

233 - name_or_path=str(path_or_repo_id),

234 - truncation_side=kwargs.get("truncation_side", "left"),

235 - chat_template=config.get("chat_template"),

236 - init_kwargs=config,

237 - )

238 -

239 - def __init__(

240 - self,

241 - *,

242 - vocab_file: Path,

243 - name_or_path: str,

244 - truncation_side: str,

245 - chat_template: str | None,

246 - init_kwargs: dict[str, Any] | None = None,

247 - ) -> None:

248 - super().__init__()

249 - self.name_or_path = name_or_path

250 - self._truncation_side = truncation_side

251 - self.init_kwargs = init_kwargs or {}

252 - self._chat_template = chat_template or DEFAULT_CHAT_TEMPLATE

253 -

254 - self._tokenizer, self._special_tokens = _load_tiktoken_encoding(vocab_file)

255 -

256 - self._token_to_id: dict[str, int] = {}

257 - self._id_to_token: dict[int, str] = {}

258 - for token, token_id in self._tokenizer._mergeable_ranks.items():

259 - token_str = token.decode("utf-8", errors="replace")

260 - self._token_to_id[token_str] = token_id

261 - self._id_to_token[token_id] = token_str

262 -

263 - for token, token_id in self._special_tokens.items():

264 - self._token_to_id[token] = token_id

265 - self._id_to_token[token_id] = token

266 -

267 - bos_token_id = self._special_tokens.get(SEP)

268 - if bos_token_id is None:

269 - bos_token_id = self._special_tokens.get(PAD)

270 - if bos_token_id is None:

271 - bos_token_id = self._special_tokens.get(EOS)

272 - if bos_token_id is None:

273 - bos_token_id = 0

274 - self._bos_token_id = bos_token_id

275 -

276 - self._eos_token_id = self._special_tokens.get(EOS, self._bos_token_id)

277 - self._pad_token_id = self._special_tokens.get(PAD, self._eos_token_id)

278 - self._unk_token_id = self._pad_token_id

279 -

280 - self._max_chars_per_token = max(len(tok) for tok in self._token_to_id)

281 -

282 - def num_special_tokens_to_add(self) -> int:

283 - return 0

284 -

285 - @property

286 - def all_special_tokens(self) -> list[str]:

287 - return list(self._special_tokens.keys())

288 -

289 - @property

290 - def all_special_ids(self) -> list[int]:

291 - return list(self._special_tokens.values())

292 -

293 - @property

294 - def bos_token_id(self) -> int:

295 - return self._bos_token_id

296 -

297 - @property

298 - def eos_token_id(self) -> int:

299 - return self._eos_token_id

300 -

301 - @property

302 - def pad_token_id(self) -> int:

303 - return self._pad_token_id

304 -

305 - @property

306 - def is_fast(self) -> bool:

307 - return False

308 -

309 - @property

310 - def vocab_size(self) -> int:

311 - return self._tokenizer.n_vocab

312 -

313 - @property

314 - def max_token_id(self) -> int:

315 - return self._tokenizer.n_vocab - 1

316 -

317 - @property

318 - def max_chars_per_token(self) -> int:

319 - return self._max_chars_per_token

320 -

321 - @property

322 - def truncation_side(self) -> str:

323 - return self._truncation_side

324 -

325 - def get_vocab(self) -> dict[str, int]:

326 - return dict(self._token_to_id)

327 -

328 - def get_added_vocab(self) -> dict[str, int]:

329 - return dict(self._special_tokens)

330 -

331 - def _maybe_truncate(self, tokens: list[int], max_length: int | None) -> list[int]:

332 - if max_length is None or len(tokens) <= max_length:

333 - return tokens

334 - if self.truncation_side == "left":

335 - return tokens[-max_length:]

336 - return tokens[:max_length]

337 -

338 - def encode(

339 - self,

340 - text: str,

341 - truncation: bool | None = None,

342 - max_length: int | None = None,

343 - add_special_tokens: bool = True,

344 - ) -> list[int]:

345 - del add_special_tokens

346 - tokens = self._tokenizer.encode(text)

347 - if truncation:

348 - tokens = self._maybe_truncate(tokens, max_length)

349 - return tokens

350 -

351 - def decode(

352 - self, ids: Sequence[int] | int, skip_special_tokens: bool = False

353 - ) -> str:

354 - if isinstance(ids, int):

355 - ids = [ids]

356 - if skip_special_tokens:

357 - ids = [

358 - token_id

359 - for token_id in ids

360 - if token_id not in self._special_tokens.values()

361 - ]

362 - return self._tokenizer.decode(ids)

363 -

364 - @overload

365 - def convert_tokens_to_ids(self, tokens: str) -> int: ...

366 -

367 - @overload

368 - def convert_tokens_to_ids(self, tokens: list[str]) -> list[int]: ...

369 -

370 - def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:

371 - if isinstance(tokens, str):

372 - return self._token_to_id.get(tokens, self._unk_token_id)

373 - return [self._token_to_id.get(token, self._unk_token_id) for token in tokens]

374 -

375 - def convert_ids_to_tokens(

376 - self, ids: Sequence[int], skip_special_tokens: bool = False

377 - ) -> list[str]:

378 - tokens = []

379 - for token_id in ids:

380 - if skip_special_tokens and token_id in self._special_tokens.values():

381 - continue

382 - tokens.append(self._id_to_token.get(token_id, "<|unk|>"))

383 - return tokens

384 -

385 - def convert_tokens_to_string(self, tokens: list[str]) -> str:

386 - token_ids = self.convert_tokens_to_ids(tokens)

387 - return self.decode(token_ids, skip_special_tokens=False)

388 -

389 - def __call__(

390 - self,

391 - text: str | list[str],

392 - text_pair: str | None = None,

393 - add_special_tokens: bool = True,

394 - truncation: bool = False,

395 - max_length: int | None = None,

396 - ) -> BatchEncoding:

397 - if text_pair is not None:

398 - raise NotImplementedError("text_pair is not supported for Grok2Tokenizer.")

399 -

400 - if isinstance(text, list):

401 - input_ids_batch: list[list[int]] = [

402 - self.encode(

403 - item,

404 - truncation=truncation,

405 - max_length=max_length,

406 - add_special_tokens=add_special_tokens,

407 - )

408 - for item in text

409 - ]

410 - attention_mask_batch = [[1] * len(ids) for ids in input_ids_batch]

411 - return BatchEncoding(

412 - {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch}

413 - )

414 -

415 - input_ids = self.encode(

416 - text,

417 - truncation=truncation,

418 - max_length=max_length,

419 - add_special_tokens=add_special_tokens,

420 - )

421 - attention_mask = [1] * len(input_ids)

422 - return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})

423 -

424 - def get_chat_template(

425 - self, chat_template: str | None, tools: list[dict[str, Any]] | None = None

426 - ) -> str | None:

427 - del tools

428 - return chat_template or self._chat_template

429 -

430 - def apply_chat_template(

431 - self,

432 - messages: list[ChatCompletionMessageParam],

433 - tools: list[dict[str, Any]] | None = None,

434 - chat_template: str | None = None,

435 - tokenize: bool = False,

436 - **kwargs,

437 - ) -> str | list[int]:

438 - template = self.get_chat_template(chat_template, tools=tools)

439 - if template is None:

440 - raise ValueError(

441 - "No chat template available. Provide `chat_template` explicitly."

442 - )

443 - kwargs["return_dict"] = False

444 - prompt = hf_chat_utils.apply_chat_template(

445 - conversation=messages,

446 - chat_template=template,

447 - tools=tools,

448 - **kwargs,

449 - )

450 - if tokenize:

451 - return self.encode(prompt, add_special_tokens=False)

452 - return prompt

vllm/tokenizers/registry.py

-1

@@ -36,7 +36,6 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"}

36 _VLLM_TOKENIZERS = {

37 "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),

38 "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"),

39 - "grok2": ("grok2", "Grok2Tokenizer"),

39 "hf": ("hf", "CachedHfTokenizer"),

40 "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"),

41 "mistral": ("mistral", "MistralTokenizer"),