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 2, 07:15 UTC
[Feature] Add support for JetBrains' Mellum v2 code generation model (#43992)
Signed-off-by: Madeesh Kannan <madeeswaran.kannan@jetbrains.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Madeesh Kannan committed
Jun 1, 2026 at 14:11 UTC
023808c23d234387298732ebc942fff5939dbd8b
7 files changed
+266
docs/models/supported_models.md
+1
@@ -437,6 +437,7 @@ th {
437
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
438
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
439
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
440
+| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
441
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
442
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
443
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
tests/models/registry.py
+1
@@ -522,6 +522,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
522
"Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
523
"Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
524
"Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
525
+ "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
526
"Qwen3NextForCausalLM": _HfExamplesInfo(
527
"Qwen/Qwen3-Next-80B-A3B-Instruct",
528
extras={"tiny-random": "tiny-random/qwen3-next-moe"},
vllm/model_executor/models/mellum.py
new
+253
@@ -0,0 +1,253 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Any
5
+
6
+from torch import nn
7
+
8
+from vllm.compilation.decorators import support_torch_compile
9
+from vllm.config import VllmConfig
10
+from vllm.distributed import get_tensor_model_parallel_world_size
11
+from vllm.model_executor.layers.attention import Attention
12
+from vllm.model_executor.layers.layernorm import RMSNorm
13
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
14
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
+from vllm.model_executor.layers.rotary_embedding import get_rope
16
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
17
+
18
+from .qwen3_moe import (
19
+ Qwen3MoeAttention,
20
+ Qwen3MoeDecoderLayer,
21
+ Qwen3MoeForCausalLM,
22
+ Qwen3MoeMLP,
23
+ Qwen3MoeModel,
24
+ Qwen3MoeSparseMoeBlock,
25
+)
26
+from .utils import PPMissingLayer, extract_layer_index, maybe_prefix
27
+
28
+
29
+class MellumAttention(Qwen3MoeAttention):
30
+ """
31
+ Differences from `Qwen3MoeAttention`:
32
+ - Supports `per_layer_sliding_window` for `Attention`.
33
+ """
34
+
35
+ def __init__(
36
+ self,
37
+ hidden_size: int,
38
+ num_heads: int,
39
+ num_kv_heads: int,
40
+ rope_parameters: dict[str, Any],
41
+ max_position_embeddings: int = 8192,
42
+ head_dim: int | None = None,
43
+ rms_norm_eps: float = 1e-06,
44
+ qkv_bias: bool = False,
45
+ cache_config: Any | None = None,
46
+ quant_config: Any | None = None,
47
+ prefix: str = "",
48
+ dual_chunk_attention_config: dict[str, Any] | None = None,
49
+ per_layer_sliding_window: int | None = None,
50
+ ) -> None:
51
+ nn.Module.__init__(self)
52
+
53
+ self.hidden_size = hidden_size
54
+ tp_size = get_tensor_model_parallel_world_size()
55
+ self.total_num_heads = num_heads
56
+ assert self.total_num_heads % tp_size == 0
57
+ self.num_heads = self.total_num_heads // tp_size
58
+ self.total_num_kv_heads = num_kv_heads
59
+ if self.total_num_kv_heads >= tp_size:
60
+ assert self.total_num_kv_heads % tp_size == 0
61
+ else:
62
+ assert tp_size % self.total_num_kv_heads == 0
63
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
64
+ self.head_dim = head_dim or (hidden_size // self.total_num_heads)
65
+ self.q_size = self.num_heads * self.head_dim
66
+ self.kv_size = self.num_kv_heads * self.head_dim
67
+ self.scaling = self.head_dim**-0.5
68
+ self.max_position_embeddings = max_position_embeddings
69
+ self.dual_chunk_attention_config = dual_chunk_attention_config
70
+
71
+ self.qkv_proj = QKVParallelLinear(
72
+ hidden_size,
73
+ self.head_dim,
74
+ self.total_num_heads,
75
+ self.total_num_kv_heads,
76
+ bias=qkv_bias,
77
+ quant_config=quant_config,
78
+ prefix=f"{prefix}.qkv_proj",
79
+ )
80
+
81
+ self.o_proj = RowParallelLinear(
82
+ self.total_num_heads * self.head_dim,
83
+ hidden_size,
84
+ bias=False,
85
+ quant_config=quant_config,
86
+ prefix=f"{prefix}.o_proj",
87
+ )
88
+
89
+ self.rotary_emb = get_rope(
90
+ self.head_dim,
91
+ max_position=max_position_embeddings,
92
+ rope_parameters=rope_parameters,
93
+ dual_chunk_attention_config=dual_chunk_attention_config,
94
+ )
95
+ self.attn = Attention(
96
+ self.num_heads,
97
+ self.head_dim,
98
+ self.scaling,
99
+ num_kv_heads=self.num_kv_heads,
100
+ cache_config=cache_config,
101
+ quant_config=quant_config,
102
+ per_layer_sliding_window=per_layer_sliding_window,
103
+ prefix=f"{prefix}.attn",
104
+ **(
105
+ {
106
+ "layer_idx": extract_layer_index(prefix),
107
+ "dual_chunk_attention_config": dual_chunk_attention_config,
108
+ }
109
+ if dual_chunk_attention_config
110
+ else {}
111
+ ),
112
+ )
113
+
114
+ self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
115
+ self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
116
+
117
+
118
+class MellumDecoderLayer(Qwen3MoeDecoderLayer):
119
+ """
120
+ Differences from `Qwen3MoeDecoderLayer`:
121
+ - Supports interleaved SWA and per-layer RoPE scaling.
122
+ """
123
+
124
+ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
125
+ nn.Module.__init__(self)
126
+
127
+ config = vllm_config.model_config.hf_text_config
128
+ cache_config = vllm_config.cache_config
129
+ quant_config = vllm_config.quant_config
130
+
131
+ self.hidden_size = config.hidden_size
132
+ max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
133
+ dual_chunk_attention_config = getattr(
134
+ config, "dual_chunk_attention_config", None
135
+ )
136
+
137
+ layer_idx = extract_layer_index(prefix)
138
+ layer_type = config.layer_types[layer_idx]
139
+ if layer_type == "sliding_attention":
140
+ sliding_window = getattr(config, "sliding_window", None)
141
+ else:
142
+ sliding_window = None
143
+ rope_parameters = config.rope_parameters[layer_type]
144
+
145
+ self.self_attn = MellumAttention(
146
+ hidden_size=self.hidden_size,
147
+ num_heads=config.num_attention_heads,
148
+ num_kv_heads=config.num_key_value_heads,
149
+ rope_parameters=rope_parameters,
150
+ max_position_embeddings=max_position_embeddings,
151
+ rms_norm_eps=config.rms_norm_eps,
152
+ qkv_bias=getattr(config, "attention_bias", False),
153
+ head_dim=getattr(config, "head_dim", None),
154
+ cache_config=cache_config,
155
+ quant_config=quant_config,
156
+ prefix=f"{prefix}.self_attn",
157
+ dual_chunk_attention_config=dual_chunk_attention_config,
158
+ per_layer_sliding_window=sliding_window,
159
+ )
160
+
161
+ if config.mlp_layer_types[layer_idx] == "sparse":
162
+ self.mlp = Qwen3MoeSparseMoeBlock(
163
+ vllm_config=vllm_config, prefix=f"{prefix}.mlp"
164
+ )
165
+ else:
166
+ self.mlp = Qwen3MoeMLP(
167
+ hidden_size=config.hidden_size,
168
+ intermediate_size=config.intermediate_size,
169
+ hidden_act=config.hidden_act,
170
+ quant_config=quant_config,
171
+ prefix=f"{prefix}.mlp",
172
+ )
173
+
174
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
175
+ self.post_attention_layernorm = RMSNorm(
176
+ config.hidden_size, eps=config.rms_norm_eps
177
+ )
178
+
179
+
180
+@support_torch_compile
181
+class MellumModel(Qwen3MoeModel):
182
+ """
183
+ Differences from `Qwen3MoeModel`:
184
+ - Uses `MellumDecoderLayer`.
185
+ """
186
+
187
+ def __init__(
188
+ self,
189
+ *,
190
+ vllm_config: VllmConfig,
191
+ prefix: str = "",
192
+ ):
193
+ super().__init__(
194
+ vllm_config=vllm_config,
195
+ prefix=prefix,
196
+ decoder_layer_type=MellumDecoderLayer,
197
+ )
198
+
199
+
200
+class MellumForCausalLM(Qwen3MoeForCausalLM):
201
+ """
202
+ Differences from `Qwen3MoeForCausalLM`:
203
+ - Uses `MellumModel`.
204
+ """
205
+
206
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
207
+ nn.Module.__init__(self)
208
+ config = vllm_config.model_config.hf_text_config
209
+ quant_config = vllm_config.quant_config
210
+ self.config = config
211
+ self.quant_config = quant_config
212
+ if "dense" in getattr(config, "mlp_layer_types", []):
213
+ self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
214
+ self.model = MellumModel(
215
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
216
+ )
217
+ self.lm_head = ParallelLMHead(
218
+ config.vocab_size,
219
+ config.hidden_size,
220
+ quant_config=quant_config,
221
+ prefix=maybe_prefix(prefix, "lm_head"),
222
+ )
223
+ if self.config.tie_word_embeddings:
224
+ self.lm_head.weight = self.model.embed_tokens.weight
225
+ self.logits_processor = LogitsProcessor(config.vocab_size)
226
+ self.make_empty_intermediate_tensors = (
227
+ self.model.make_empty_intermediate_tensors
228
+ )
229
+
230
+ self.expert_weights = []
231
+
232
+ self.moe_layers = []
233
+ example_layer = None
234
+ for layer in self.model.layers:
235
+ if isinstance(layer, PPMissingLayer):
236
+ continue
237
+
238
+ assert isinstance(layer, Qwen3MoeDecoderLayer)
239
+ if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
240
+ example_layer = layer.mlp
241
+ self.moe_layers.append(layer.mlp.experts)
242
+
243
+ if example_layer is None:
244
+ raise RuntimeError("No MoE layer found in the model.layers.")
245
+
246
+ self.num_moe_layers = len(self.moe_layers)
247
+ self.num_expert_groups = 1
248
+ self.num_shared_experts = 0
249
+ self.num_logical_experts = example_layer.n_logical_experts
250
+ self.num_physical_experts = example_layer.n_physical_experts
251
+ self.num_local_physical_experts = example_layer.n_local_physical_experts
252
+ self.num_routed_experts = example_layer.n_routed_experts
253
+ self.num_redundant_experts = example_layer.n_redundant_experts
vllm/model_executor/models/registry.py
+1
@@ -159,6 +159,7 @@ _TEXT_GENERATION_MODELS = {
159
"LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
160
"MambaForCausalLM": ("mamba", "MambaForCausalLM"),
161
"Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
162
+ "MellumForCausalLM": ("mellum", "MellumForCausalLM"),
163
"MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
164
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
165
"MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
vllm/transformers_utils/config.py
+1
@@ -116,6 +116,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
116
RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct)
117
mlp_speculator="MLPSpeculatorConfig",
118
medusa="MedusaConfig",
119
+ mellum="MellumConfig",
120
midashenglm="MiDashengLMConfig",
121
moondream3="Moondream3Config",
122
eagle="EAGLEConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -49,6 +49,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
49
"LagunaConfig": "vllm.transformers_utils.configs.laguna",
50
"Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",
51
"MedusaConfig": "vllm.transformers_utils.configs.medusa",
52
+ "MellumConfig": "vllm.transformers_utils.configs.mellum",
53
"MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
54
"MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",
55
"Moondream3Config": "vllm.transformers_utils.configs.moondream3",
@@ -117,6 +118,7 @@ __all__ = [
118
"LagunaConfig",
119
"Lfm2MoeConfig",
120
"MedusaConfig",
121
+ "MellumConfig",
122
"MiDashengLMConfig",
123
"MLPSpeculatorConfig",
124
"Moondream3Config",
vllm/transformers_utils/configs/mellum.py
new
+7
@@ -0,0 +1,7 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from transformers import Qwen3MoeConfig
4
+
5
+
6
+class MellumConfig(Qwen3MoeConfig):
7
+ model_type = "mellum"