@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 429e4e2d42

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 14, 05:30 UTC

[Bugfix] Fix ModernBert cuda graph capturing in v1 (#21901)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <2037008807@qq.com>

Isotr0py committed Aug 9, 2025 at 13:17 UTC 429e4e2d420f7c648d37b7d90430f5df6a7dc61f

5 files changed +39 -42

tests/models/language/pooling/mteb_utils.py

+3 -2

@@ -162,7 +162,8 @@ def mteb_test_embed_models(hf_runner,

162 vllm_runner,

163 model_info: EmbedModelInfo,

164 vllm_extra_kwargs=None,

165 - hf_model_callback=None):

165 + hf_model_callback=None,

166 + atol=MTEB_RERANK_TOL):

167 if not model_info.enable_test:

168 # A model family has many models with the same architecture,

169 # and we don't need to test each one.

@@ -198,7 +199,7 @@ def mteb_test_embed_models(hf_runner,

199 print("SentenceTransformers:", st_dtype, st_main_score)

200 print("Difference:", st_main_score - vllm_main_score)

201

201 - assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_EMBED_TOL)

202 + assert st_main_score == pytest.approx(vllm_main_score, abs=atol)

203

204

205 def run_mteb_rerank(cross_encoder, tasks, languages):

vllm/model_executor/models/bert.py

+1 -1

@@ -466,7 +466,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):

466

467 def forward(

468 self,

469 - input_ids: Optional[torch.Tensor],

469 + input_ids: torch.Tensor,

470 positions: torch.Tensor,

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

472 intermediate_tensors: Optional[IntermediateTensors] = None,

vllm/model_executor/models/bert_with_rope.py

+21 -25

@@ -8,13 +8,15 @@ from torch import nn

8 from transformers import PretrainedConfig

10 from vllm.attention import Attention, AttentionType

11 +from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import CacheConfig, VllmConfig

13 from vllm.distributed import (divide, get_tensor_model_parallel_rank,

14 get_tensor_model_parallel_world_size,

15 tensor_model_parallel_all_reduce)

16 from vllm.model_executor.layers.activation import (get_act_and_mul_fn,

17 get_act_fn)

17 -from vllm.model_executor.layers.fused_moe import fused_moe

18 +from vllm.model_executor.layers.fused_moe.fused_moe import (

19 + fused_topk, torch_vllm_outplace_fused_experts)

20 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

21 MergedColumnParallelLinear,

22 QKVParallelLinear,

@@ -284,15 +286,22 @@ class NomicMoE(nn.Module):

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

287 # router_logits: (num_tokens, n_experts)

288 router_logits, _ = self.router(hidden_states)

287 - final_hidden_states = fused_moe(hidden_states,

288 - self.w1,

289 - self.w2,

290 - router_logits,

291 - self.top_k,

292 - renormalize=False,

293 - inplace=False,

294 - activation=self.hidden_act,

295 - is_act_and_mul=False)

289 + # FIXME(Isotr0py): This implementation is too tricky,

290 + # we should use FusedMoE instead in the future

291 + # after supporting ungated activation for it.

292 + topk_weights, topk_ids, _ = fused_topk(hidden_states,

293 + router_logits,

294 + self.top_k,

295 + renormalize=False)

296 + final_hidden_states = torch_vllm_outplace_fused_experts(

297 + hidden_states=hidden_states,

298 + w1=self.w1,

299 + w2=self.w2,

300 + topk_weights=topk_weights,

301 + topk_ids=topk_ids,

302 + activation=self.hidden_act,

303 + is_act_and_mul=False,

304 + )

305

306 if self.tp_size > 1:

307 final_hidden_states = tensor_model_parallel_all_reduce(

@@ -391,6 +400,7 @@ class BertWithRopeEncoder(nn.Module):

400 return hidden_states

401

402

403 +@support_torch_compile

404 class BertWithRope(nn.Module, SupportsQuant):

405 hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})

406

@@ -407,7 +417,7 @@ class BertWithRope(nn.Module, SupportsQuant):

417

418 def forward(

419 self,

410 - input_ids: Optional[torch.Tensor],

420 + input_ids: torch.Tensor,

421 positions: torch.Tensor,

422 intermediate_tensors: Optional[IntermediateTensors] = None,

423 inputs_embeds: Optional[torch.Tensor] = None,

@@ -554,20 +564,6 @@ class JinaRobertaModel(BertWithRope):

564 "norm2": "mlp_ln",

565 })

566

557 - def forward(

558 - self,

559 - input_ids: torch.Tensor,

560 - position_ids: torch.Tensor,

561 - intermediate_tensors: Optional[IntermediateTensors] = None,

562 - inputs_embeds: Optional[torch.Tensor] = None,

563 - token_type_ids: Optional[torch.Tensor] = None,

564 - ) -> torch.Tensor:

565 - return super().forward(input_ids=input_ids,

566 - positions=position_ids,

567 - intermediate_tensors=intermediate_tensors,

568 - inputs_embeds=inputs_embeds,

569 - token_type_ids=token_type_ids)

570 -

567 @torch.inference_mode()

568 def jina_merge_lora_weights(self, weights: Iterable[tuple[str,

569 torch.Tensor]]):

vllm/model_executor/models/modernbert.py

+11 -11

@@ -8,6 +8,7 @@ from torch import nn

8 from transformers import ModernBertConfig

10 from vllm.attention import Attention, AttentionType

11 +from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import VllmConfig

13 from vllm.distributed import get_tensor_model_parallel_world_size

14 from vllm.model_executor.layers.linear import (QKVParallelLinear,

@@ -46,7 +47,7 @@ class ModernBertEmbeddings(nn.Module):

47 input_ids: torch.Tensor,

48 inputs_embeds: Optional[torch.Tensor] = None,

49 ) -> torch.Tensor:

49 - if inputs_embeds:

50 + if inputs_embeds is not None:

51 return self.norm(inputs_embeds)

52 else:

53 inputs_embeds = self.tok_embeddings(input_ids)

@@ -117,7 +118,7 @@ class ModernBertAttention(nn.Module):

118 def forward(

119 self,

120 hidden_states: torch.Tensor,

120 - position_ids: Optional[torch.LongTensor] = None,

121 + position_ids: torch.Tensor,

122 ) -> torch.Tensor:

123 qkv, _ = self.Wqkv(hidden_states)

124 q, k, v = qkv.split([self.all_head_size] * 3, dim=-1)

@@ -169,9 +170,9 @@ class ModernBertLayer(nn.Module):

170 def forward(

171 self,

172 hidden_states: torch.Tensor,

172 - position_ids: Optional[torch.LongTensor] = None,

173 - ):

174 - attn_outputs = self.attn(self.attn_norm(hidden_states),

173 + position_ids: torch.Tensor,

174 + ) -> torch.Tensor:

175 + attn_outputs = self.attn(hidden_states=self.attn_norm(hidden_states),

176 position_ids=position_ids)

177 hidden_states = hidden_states + attn_outputs

178 mlp_output = self.mlp(self.mlp_norm(hidden_states))

@@ -192,13 +193,14 @@ class ModernBertEncoderLayer(nn.Module):

193 def forward(

194 self,

195 hidden_states: torch.Tensor,

195 - position_ids: Optional[torch.LongTensor] = None,

196 + position_ids: torch.Tensor,

197 ) -> torch.Tensor:

198 for i, layer in enumerate(self.layers):

199 hidden_states = layer(hidden_states, position_ids)

200 return hidden_states

201

202

203 +@support_torch_compile

204 class ModernBertModel(nn.Module):

205 hf_to_vllm_mapper = WeightsMapper(

206 orig_to_new_prefix={"layers.": "encoder_layer.layers."})

@@ -234,13 +236,11 @@ class ModernBertModel(nn.Module):

236

237 def forward(

238 self,

237 - input_ids: Optional[torch.LongTensor] = None,

238 - positions: Optional[torch.Tensor] = None,

239 + input_ids: torch.Tensor,

240 + positions: torch.Tensor,

241 intermediate_tensors: Optional[IntermediateTensors] = None,

242 inputs_embeds: Optional[torch.Tensor] = None,

241 - position_ids: Optional[torch.LongTensor] = None,

243 ) -> torch.Tensor:

243 - position_ids = positions if positions is not None else position_ids

244 if inputs_embeds is not None:

245 hidden_states = inputs_embeds

246 else:

@@ -249,7 +249,7 @@ class ModernBertModel(nn.Module):

249

250 outputs = self.encoder_layer(

251 hidden_states=hidden_states,

252 - position_ids=position_ids,

252 + position_ids=positions,

253 )

254 norm_outputs = self.final_norm(outputs)

255 return norm_outputs

vllm/model_executor/models/roberta.py

+3 -3

@@ -105,7 +105,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

105

106 def forward(

107 self,

108 - input_ids: Optional[torch.Tensor],

108 + input_ids: torch.Tensor,

109 positions: torch.Tensor,

110 token_type_ids: Optional[torch.Tensor] = None,

111 intermediate_tensors: Optional[IntermediateTensors] = None,

@@ -119,8 +119,8 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

119 position_ids=positions,

120 padding_idx=self.padding_idx)

121

122 - return self.model(input_ids=input_ids,

123 - position_ids=positions,

122 + return self.model(input_ids,

123 + positions,

124 token_type_ids=token_type_ids,

125 inputs_embeds=inputs_embeds,

126 intermediate_tensors=intermediate_tensors)