@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 59d260f5e4

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 10, 12:45 UTC

[Model] Add Grok-2 (#31847)

Signed-off-by: dangoldbj <dangoldbj23@gmail.com>

Bijaya Dangol committed Jan 8, 2026 at 13:59 UTC 59d260f5e4b07eaffb34c79579959f3c517e606f

8 files changed +777 -20

docs/models/supported_models.md

@@ -399,6 +399,7 @@ th {

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

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

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

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

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

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

@@ -459,6 +460,9 @@ th {

461 | `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |

462

463 +!!! note

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

465 +

466 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!

467

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

+43

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

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

@@ -289,6 +289,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

289 "Grok1ModelForCausalLM": _HfExamplesInfo(

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

291 ),

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

293 "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"),

294 "HunYuanMoEV1ForCausalLM": _HfExamplesInfo(

295 "tencent/Hunyuan-A13B-Instruct", trust_remote_code=True

tests/tokenizers_/test_basic.py

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

10 )

12 from vllm.tokenizers import TokenizerLike, get_tokenizer

13 +from vllm.tokenizers.grok2 import Grok2Tokenizer

14 from vllm.tokenizers.mistral import MistralTokenizer

@@ -37,6 +38,10 @@ def test_tokenizer_like_protocol():

38 assert isinstance(tokenizer, MistralTokenizer)

39 _assert_tokenizer_like(tokenizer)

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

42 + assert isinstance(tokenizer, Grok2Tokenizer)

43 + _assert_tokenizer_like(tokenizer)

44 +

46 @pytest.mark.parametrize("tokenizer_name", ["facebook/opt-125m", "gpt2"])

47 def test_tokenizer_revision(tokenizer_name: str):

vllm/model_executor/models/grok1.py

+267 -19

@@ -21,8 +21,9 @@

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 Grok1 model."""

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

26 +import math

27 from collections.abc import Iterable

28 from itertools import islice

29 from typing import Any

@@ -35,9 +36,12 @@ from vllm.attention.layer import Attention

36 from vllm.compilation.decorators import support_torch_compile

37 from vllm.config import CacheConfig, VllmConfig

38 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

39 +from vllm.logger import init_logger

40 +from vllm.model_executor.layers.activation import GeluAndMul

41 from vllm.model_executor.layers.fused_moe import FusedMoE

42 from vllm.model_executor.layers.layernorm import RMSNorm

43 from vllm.model_executor.layers.linear import (

44 + MergedColumnParallelLinear,

45 QKVParallelLinear,

46 ReplicatedLinear,

47 RowParallelLinear,

@@ -68,6 +72,100 @@ from .utils import (

72 DEFAULT_ATTN_OUTPUT_MULTIPLIER = 0.08838834764831845

73 DEFAULT_OUTPUT_MULTIPLIER_SCALE = 0.5773502691896257

74 DEFAULT_EMBEDDING_MULTIPLIER_SCALE = 78.38367176906169

75 +DEFAULT_ROUTER_LOGIT_SOFTCAP = 30.0

76 +

77 +logger = init_logger(__name__)

78 +

79 +

80 +def _get_num_experts(config) -> int:

81 + return getattr(config, "num_experts", getattr(config, "num_local_experts", 8))

82 +

83 +

84 +def _get_moe_intermediate_size(config) -> int:

85 + return getattr(config, "moe_intermediate_size", config.intermediate_size)

86 +

87 +

88 +def _get_grok_version(config) -> str:

89 + """Detect Grok version from HF config using multiple heuristics."""

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

91 + has_residual_moe = getattr(config, "residual_moe", False)

92 + has_moe_intermediate_size = hasattr(config, "moe_intermediate_size")

93 +

94 + if has_residual_moe or has_moe_intermediate_size:

95 + return "grok2"

96 +

97 + return "grok1" # Default to Grok1

98 +

99 +

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

101 + rope_parameters = getattr(config, "rope_parameters", None)

102 + if rope_parameters is None:

103 + rope_type = getattr(config, "rope_type", None)

104 + if rope_type is None:

105 + return None

106 + rope_parameters = {"rope_type": rope_type}

107 + rope_theta = getattr(config, "rope_theta", None)

108 + if rope_theta is not None:

109 + rope_parameters["rope_theta"] = rope_theta

110 + scaling_factor = getattr(config, "scaling_factor", None)

111 + if scaling_factor is not None:

112 + rope_parameters["factor"] = scaling_factor

113 + for name in (

114 + "original_max_position_embeddings",

115 + "extrapolation_factor",

116 + "attn_factor",

117 + "beta_fast",

118 + "beta_slow",

119 + ):

120 + value = getattr(config, name, None)

121 + if value is not None:

122 + rope_parameters[name] = value

123 +

124 + if rope_parameters.get("rope_type") == "original":

125 + rope_parameters = dict(rope_parameters)

126 + rope_parameters["rope_type"] = "default"

127 + return rope_parameters

128 +

129 +

130 +def _get_moe_renormalize(config) -> bool:

131 + explicit_value = getattr(

132 + config, "moe_router_renormalize", getattr(config, "moe_renormalize", None)

133 + )

134 + if explicit_value is not None:

135 + return bool(explicit_value)

136 + return not getattr(config, "residual_moe", False)

137 +

138 +

139 +class Grok1MLP(nn.Module):

140 + def __init__(

141 + self,

142 + hidden_size: int,

143 + intermediate_size: int,

144 + quant_config: QuantizationConfig | None = None,

145 + prefix: str = "",

146 + ) -> None:

147 + super().__init__()

148 + self.gate_up_proj = MergedColumnParallelLinear(

149 + input_size=hidden_size,

150 + output_sizes=[intermediate_size] * 2,

151 + bias=False,

152 + quant_config=quant_config,

153 + prefix=f"{prefix}.gate_up_proj",

154 + )

155 + self.down_proj = RowParallelLinear(

156 + input_size=intermediate_size,

157 + output_size=hidden_size,

158 + bias=False,

159 + quant_config=quant_config,

160 + prefix=f"{prefix}.down_proj",

161 + )

162 + self.act_fn = GeluAndMul()

163 +

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

165 + x, _ = self.gate_up_proj(x)

166 + x = self.act_fn(x)

167 + x, _ = self.down_proj(x)

168 + return x

169

170

171 class Grok1MoE(nn.Module):

@@ -85,9 +183,11 @@ class Grok1MoE(nn.Module):

183 top_k: int,

184 hidden_size: int,

185 intermediate_size: int,

186 + router_logit_soft_cap: float = 0.0,

187 params_dtype: torch.dtype | None = None,

188 quant_config: QuantizationConfig | None = None,

189 tp_size: int | None = None,

190 + renormalize: bool = False,

191 prefix: str = "",

192 ):

193 super().__init__()

@@ -110,12 +210,13 @@ class Grok1MoE(nn.Module):

210 intermediate_size=intermediate_size,

211 params_dtype=params_dtype,

212 reduce_results=True,

113 - renormalize=True,

213 + renormalize=renormalize,

214 quant_config=quant_config,

215 tp_size=tp_size,

216 activation="gelu",

217 prefix=f"{prefix}.experts",

218 )

219 + self.router_logit_soft_cap = router_logit_soft_cap

220

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

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

@@ -123,7 +224,10 @@ class Grok1MoE(nn.Module):

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

225 # router_logits: (num_tokens, n_experts)

226 router_logits, _ = self.gate(hidden_states)

126 - router_logits = 30.0 * F.tanh(router_logits / 30.0)

227 + if self.router_logit_soft_cap > 0:

228 + router_logits = self.router_logit_soft_cap * F.tanh(

229 + router_logits / self.router_logit_soft_cap

230 + )

231 final_hidden_states = self.experts(hidden_states, router_logits)

232 return final_hidden_states.view(orig_shape)

233

@@ -187,6 +291,15 @@ class Grok1Attention(nn.Module):

291 )

292

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

294 + attn_logit_softcapping_method = getattr(

295 + config, "attn_logit_softcapping_method", None

296 + )

297 + if attn_logit_softcapping_method not in (None, "tanh"):

298 + logger.warning_once(

299 + "Grok attention logit softcapping method '%s' is not "

300 + "supported; falling back to default behavior.",

301 + attn_logit_softcapping_method,

302 + )

303

304 self.attn = Attention(

305 self.num_heads,

@@ -238,30 +351,50 @@ class Grok1DecoderLayer(nn.Module):

351 num_heads=config.num_attention_heads,

352 max_position=config.max_position_embeddings,

353 num_kv_heads=config.num_key_value_heads,

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

354 + rope_parameters=_get_rope_parameters(config),

355 cache_config=cache_config,

356 quant_config=quant_config,

357 prefix=f"{prefix}.attn",

358 config=config,

359 ) # Pass config to Grok1Attention

360

248 - # Grok1 uses "num_experts" in its config

249 - num_experts = getattr(config, "num_experts", 8)

361 + num_experts = _get_num_experts(config)

362 num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)

363 + moe_intermediate_size = _get_moe_intermediate_size(config)

364 + moe_renormalize = _get_moe_renormalize(config)

365

366 self.moe_block = Grok1MoE(

367 num_experts=num_experts,

368 top_k=num_experts_per_tok,

369 hidden_size=config.hidden_size,

256 - intermediate_size=config.intermediate_size,

370 + intermediate_size=moe_intermediate_size,

371 + router_logit_soft_cap=max(

372 + getattr(

373 + config,

374 + "router_logit_softcapping",

375 + DEFAULT_ROUTER_LOGIT_SOFTCAP,

376 + ),

377 + 0.0,

378 + ),

379 quant_config=quant_config,

380 + renormalize=moe_renormalize,

381 prefix=f"{prefix}.moe_block",

382 )

383 + self.residual_moe = getattr(config, "residual_moe", False)

384 + self.residual_moe_scale = 1.0 / math.sqrt(2.0)

385

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

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

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

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

390 + self.mlp = None

391 + if self.residual_moe:

392 + self.mlp = Grok1MLP(

393 + hidden_size=config.hidden_size,

394 + intermediate_size=config.intermediate_size,

395 + quant_config=quant_config,

396 + prefix=f"{prefix}.mlp",

397 + )

398

399 def forward(

400 self,

@@ -286,7 +419,13 @@ class Grok1DecoderLayer(nn.Module):

419

420 # MoE block with normalization

421 hidden_states, residual = self.pre_moe_norm(hidden_states, residual)

289 - hidden_states = self.moe_block(hidden_states)

422 + if self.residual_moe:

423 + assert self.mlp is not None

424 + hidden_states = (

425 + self.moe_block(hidden_states) + self.mlp(hidden_states)

426 + ) * self.residual_moe_scale

427 + else:

428 + hidden_states = self.moe_block(hidden_states)

429 hidden_states = self.post_moe_norm(hidden_states)

430

431 return hidden_states, residual

@@ -294,7 +433,16 @@ class Grok1DecoderLayer(nn.Module):

433

434 @support_torch_compile

435 class Grok1Model(nn.Module):

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

436 + def __init__(

437 + self,

438 + *,

439 + vllm_config: VllmConfig,

440 + prefix: str = "",

441 + ckpt_gate_proj_name: str = "linear",

442 + ckpt_down_proj_name: str = "linear_1",

443 + ckpt_up_proj_name: str = "linear_v",

444 + weight_name_remapping: dict[str, str] | None = None,

445 + ):

446 super().__init__()

447

448 config = vllm_config.model_config.hf_config

@@ -305,6 +453,12 @@ class Grok1Model(nn.Module):

453 self.quant_config = quant_config

454 self.padding_idx = config.pad_token_id

455

456 + # Store expert naming for weight loading

457 + self.ckpt_gate_proj_name = ckpt_gate_proj_name

458 + self.ckpt_down_proj_name = ckpt_down_proj_name

459 + self.ckpt_up_proj_name = ckpt_up_proj_name

460 + self.weight_name_remapping = weight_name_remapping or {}

461 +

462 self.vocab_size = config.vocab_size

463

464 self.embedding_multiplier_scale = getattr(

@@ -365,14 +519,13 @@ class Grok1Model(nn.Module):

519 return hidden_states

520

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

368 - # Map Grok1's unique expert parameter names to standard names

369 - # Grok1 uses "num_experts" in its config

370 - num_experts = getattr(self.config, "num_experts", 8)

522 + # Map expert parameter names to standard names

523 + num_experts = _get_num_experts(self.config)

524 return FusedMoE.make_expert_params_mapping(

525 self,

373 - ckpt_gate_proj_name="linear", # Grok1 specific

374 - ckpt_down_proj_name="linear_1", # Grok1 specific

375 - ckpt_up_proj_name="linear_v", # Grok1 specific

526 + ckpt_gate_proj_name=self.ckpt_gate_proj_name,

527 + ckpt_down_proj_name=self.ckpt_down_proj_name,

528 + ckpt_up_proj_name=self.ckpt_up_proj_name,

529 num_experts=num_experts,

530 )

531

@@ -382,12 +535,18 @@ class Grok1Model(nn.Module):

535 ("qkv_proj", "q_proj", "q"),

536 ("qkv_proj", "k_proj", "k"),

537 ("qkv_proj", "v_proj", "v"),

538 + ("mlp.gate_up_proj", "mlp.gate_proj", 0),

539 + ("mlp.gate_up_proj", "mlp.up_proj", 1),

540 ]

541

542 params_dict = dict(self.named_parameters())

543 loaded_params: set[str] = set()

544 expert_params_mapping = self.get_expert_mapping()

545 for name, loaded_weight in weights:

546 + # Apply version-specific weight name remapping

547 + for old_pattern, new_pattern in self.weight_name_remapping.items():

548 + if old_pattern in name:

549 + name = name.replace(old_pattern, new_pattern)

550 if self.quant_config is not None and (

551 scale_name := self.quant_config.get_cache_scale(name)

552 ):

@@ -418,6 +577,8 @@ class Grok1Model(nn.Module):

577 name = maybe_remap_kv_scale_name(name, params_dict)

578 if name is None:

579 continue

580 + if name not in params_dict:

581 + continue

582 param = params_dict[name]

583 weight_loader = param.weight_loader

584 weight_loader(param, loaded_weight, shard_id)

@@ -464,6 +625,8 @@ class Grok1Model(nn.Module):

625 if "norm.scale" in name:

626 name = name.replace("scale", "weight")

627

628 + if name not in params_dict:

629 + continue

630 param = params_dict[name]

631 weight_loader = getattr(

632 param, "weight_loader", default_weight_loader

@@ -473,9 +636,12 @@ class Grok1Model(nn.Module):

636 return loaded_params

637

638

476 -class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

639 +class GrokBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

640 + """Base class for Grok models with shared logic."""

641 +

642 fall_back_to_pt_during_load = False

643

644 + # Subclasses should override these

645 packed_modules_mapping = {

646 "qkv_proj": [

647 "q_proj",

@@ -484,6 +650,15 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

650 ],

651 }

652

653 + # Expert weight naming - subclasses override these

654 + ckpt_gate_proj_name: str = "linear"

655 + ckpt_down_proj_name: str = "linear_1"

656 + ckpt_up_proj_name: str = "linear_v"

657 +

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

659 + """Return weight name remapping for this version. Override in subclasses."""

660 + return {}

661 +

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

663 super().__init__()

664

@@ -491,11 +666,15 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

666 quant_config = vllm_config.quant_config

667

668 self.config = config

494 -

669 self.quant_config = quant_config

670

671 self.model = Grok1Model(

498 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

672 + vllm_config=vllm_config,

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

674 + ckpt_gate_proj_name=self.ckpt_gate_proj_name,

675 + ckpt_down_proj_name=self.ckpt_down_proj_name,

676 + ckpt_up_proj_name=self.ckpt_up_proj_name,

677 + weight_name_remapping=self.get_weight_name_remapping(),

678 )

679

680 self.lm_head = ParallelLMHead(

@@ -512,7 +691,9 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

691 config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE

692 )

693 self.logits_processor = LogitsProcessor(

515 - config.vocab_size, scale=self.output_multiplier_scale

694 + config.vocab_size,

695 + scale=self.output_multiplier_scale,

696 + soft_cap=getattr(config, "final_logit_softcapping", None),

697 )

698

699 self.make_empty_intermediate_tensors = (

@@ -553,3 +734,70 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

734

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

736 return self.model.get_expert_mapping()

737 +

738 +

739 +class Grok1ForCausalLM(GrokBaseForCausalLM):

740 + """Grok1-specific implementation."""

741 +

742 + # Grok1 expert weight naming

743 + ckpt_gate_proj_name = "linear"

744 + ckpt_down_proj_name = "linear_1"

745 + ckpt_up_proj_name = "linear_v"

746 +

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

748 + # Grok1 uses standard naming, no remapping needed

749 + return {}

750 +

751 +

752 +class Grok2ForCausalLM(GrokBaseForCausalLM):

753 + """Grok2-specific implementation."""

754 +

755 + # Grok2 has additional packed modules for MLP

756 + packed_modules_mapping = {

757 + "qkv_proj": [

758 + "q_proj",

759 + "k_proj",

760 + "v_proj",

761 + ],

762 + "gate_up_proj": [

763 + "gate_proj",

764 + "up_proj",

765 + ],

766 + }

767 +

768 + # Grok2 expert weight naming

769 + ckpt_gate_proj_name = "w1"

770 + ckpt_down_proj_name = "w2"

771 + ckpt_up_proj_name = "w3"

772 +

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

774 + # Grok2 checkpoint uses different naming conventions

775 + return {

776 + ".self_attn.": ".attn.",

777 + ".block_sparse_moe.": ".moe_block.",

778 + }

779 +

780 +

781 +# Version dispatch mapping

782 +_GROK_VERSIONS: dict[str, type[GrokBaseForCausalLM]] = {

783 + "grok1": Grok1ForCausalLM,

784 + "grok2": Grok2ForCausalLM,

785 +}

786 +

787 +

788 +class GrokForCausalLM(GrokBaseForCausalLM):

789 + """Factory class that dispatches to version-specific implementation."""

790 +

791 + def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):

792 + config = vllm_config.model_config.hf_config

793 + version = _get_grok_version(config)

794 +

795 + instance_cls = _GROK_VERSIONS.get(version)

796 + if instance_cls is None:

797 + raise ValueError(f"Unsupported Grok version: {version}")

798 +

799 + # Merge class attributes for LoRA/quantization compatibility

800 + cls.packed_modules_mapping = dict(cls.packed_modules_mapping)

801 + cls.packed_modules_mapping.update(instance_cls.packed_modules_mapping)

802 +

803 + return instance_cls(vllm_config=vllm_config, prefix=prefix)

vllm/model_executor/models/registry.py

+2 -1

@@ -119,7 +119,8 @@ _TEXT_GENERATION_MODELS = {

119 "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), # noqa: E501

120 "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), # noqa: E501

121 "GritLM": ("gritlm", "GritLM"),

122 - "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),

122 + "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),

123 + "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),

124 "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),

125 "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),

126 "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),

vllm/tokenizers/grok2.py new

+443

@@ -0,0 +1,443 @@

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, Set

8 +from pathlib import Path

9 +from typing import Any, Literal, overload

10 +

11 +from huggingface_hub import hf_hub_download

12 +from huggingface_hub.utils import (

13 + EntryNotFoundError,

14 + HfHubHTTPError,

15 + RepositoryNotFoundError,

16 + RevisionNotFoundError,

17 +)

18 +from transformers import BatchEncoding

19 +from transformers.utils import chat_template_utils as hf_chat_utils

20 +

21 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

22 +from vllm.logger import init_logger

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_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_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 + def num_special_tokens_to_add(self) -> int:

281 + return 0

282 +

283 + @property

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

285 + return list(self._special_tokens.keys())

286 +

287 + @property

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

289 + return list(self._special_tokens.values())

290 +

291 + @property

292 + def bos_token_id(self) -> int:

293 + return self._bos_token_id

294 +

295 + @property

296 + def eos_token_id(self) -> int:

297 + return self._eos_token_id

298 +

299 + @property

300 + def pad_token_id(self) -> int:

301 + return self._pad_token_id

302 +

303 + @property

304 + def is_fast(self) -> bool:

305 + return False

306 +

307 + @property

308 + def vocab_size(self) -> int:

309 + return self._tokenizer.n_vocab

310 +

311 + @property

312 + def max_token_id(self) -> int:

313 + return self._tokenizer.n_vocab - 1

314 +

315 + @property

316 + def truncation_side(self) -> str:

317 + return self._truncation_side

318 +

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

320 + return dict(self._token_to_id)

321 +

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

323 + return dict(self._special_tokens)

324 +

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

326 + if max_length is None or len(tokens) <= max_length:

327 + return tokens

328 + if self.truncation_side == "left":

329 + return tokens[-max_length:]

330 + return tokens[:max_length]

331 +

332 + def encode(

333 + self,

334 + text: str,

335 + truncation: bool | None = None,

336 + max_length: int | None = None,

337 + add_special_tokens: bool = True,

338 + ) -> list[int]:

339 + del add_special_tokens

340 + tokens = self._tokenizer.encode(text)

341 + if truncation:

342 + tokens = self._maybe_truncate(tokens, max_length)

343 + return tokens

344 +

345 + def decode(self, ids: list[int] | int, skip_special_tokens: bool = False) -> str:

346 + if isinstance(ids, int):

347 + ids = [ids]

348 + if skip_special_tokens:

349 + ids = [

350 + token_id

351 + for token_id in ids

352 + if token_id not in self._special_tokens.values()

353 + ]

354 + return self._tokenizer.decode(ids)

355 +

356 + @overload

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

358 +

359 + @overload

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

361 +

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

363 + if isinstance(tokens, str):

364 + return self._token_to_id.get(tokens, self._unk_token_id)

365 + return [self._token_to_id.get(token, self._unk_token_id) for token in tokens]

366 +

367 + def convert_ids_to_tokens(

368 + self, ids: list[int], skip_special_tokens: bool = False

369 + ) -> list[str]:

370 + tokens = []

371 + for token_id in ids:

372 + if skip_special_tokens and token_id in self._special_tokens.values():

373 + continue

374 + tokens.append(self._id_to_token.get(token_id, "<|unk|>"))

375 + return tokens

376 +

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

378 + token_ids = self.convert_tokens_to_ids(tokens)

379 + return self.decode(token_ids, skip_special_tokens=False)

380 +

381 + def __call__(

382 + self,

383 + text: str | list[str],

384 + text_pair: str | None = None,

385 + add_special_tokens: bool = True,

386 + truncation: bool = False,

387 + max_length: int | None = None,

388 + ) -> BatchEncoding:

389 + if text_pair is not None:

390 + raise NotImplementedError("text_pair is not supported for Grok2Tokenizer.")

391 +

392 + if isinstance(text, list):

393 + input_ids_batch: list[list[int]] = [

394 + self.encode(

395 + item,

396 + truncation=truncation,

397 + max_length=max_length,

398 + add_special_tokens=add_special_tokens,

399 + )

400 + for item in text

401 + ]

402 + attention_mask_batch = [[1] * len(ids) for ids in input_ids_batch]

403 + return BatchEncoding(

404 + {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch}

405 + )

406 +

407 + input_ids = self.encode(

408 + text,

409 + truncation=truncation,

410 + max_length=max_length,

411 + add_special_tokens=add_special_tokens,

412 + )

413 + attention_mask = [1] * len(input_ids)

414 + return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})

415 +

416 + def get_chat_template(

417 + self, chat_template: str | None, tools: list[dict[str, Any]] | None = None

418 + ) -> str | None:

419 + del tools

420 + return chat_template or self._chat_template

421 +

422 + def apply_chat_template(

423 + self,

424 + messages: list[ChatCompletionMessageParam],

425 + tools: list[dict[str, Any]] | None = None,

426 + chat_template: str | None = None,

427 + tokenize: bool = False,

428 + **kwargs,

429 + ) -> str | list[int]:

430 + template = self.get_chat_template(chat_template, tools=tools)

431 + if template is None:

432 + raise ValueError(

433 + "No chat template available. Provide `chat_template` explicitly."

434 + )

435 + prompt = hf_chat_utils.apply_chat_template(

436 + conversation=messages,

437 + chat_template=template,

438 + tools=tools,

439 + **kwargs,

440 + )

441 + if tokenize:

442 + return self.encode(prompt, add_special_tokens=False)

443 + return prompt

vllm/tokenizers/registry.py

+12

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

32 _VLLM_TOKENIZERS = {

33 "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),

34 + "grok2": ("grok2", "Grok2Tokenizer"),

35 "hf": ("hf", "CachedHfTokenizer"),

36 "mistral": ("mistral", "MistralTokenizer"),

37 }

@@ -151,6 +152,17 @@ def resolve_tokenizer_args(

152 if len(files_list) > 0:

153 tokenizer_mode = "mistral"

154

155 + # Try to use Grok2 tiktoken tokenizer if possible

156 + if tokenizer_mode == "auto":

157 + allow_patterns = ["tokenizer.tok.json"]

158 + files_list = list_filtered_repo_files(

159 + model_name_or_path=str(tokenizer_name),

160 + allow_patterns=allow_patterns,

161 + revision=revision,

162 + )

163 + if len(files_list) > 0:

164 + tokenizer_mode = "grok2"

165 +

166 # Fallback to HF tokenizer

167 if tokenizer_mode == "auto":

168 tokenizer_mode = "hf"