@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6d729c43fb

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 load & Enable sliding window attention for bidirectional attention. (#22637)

Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Max de Bayser <mbayser@br.ibm.com>

wang.yuqi committed Aug 12, 2025 at 15:23 UTC 6d729c43fbaf63d534e71c0b8aa61f0a82dd2018

4 files changed +101 -59

tests/models/language/pooling/test_gte.py

+18 -3

@@ -4,10 +4,11 @@ from typing import Any

5 import pytest

7 -from ...utils import (CLSPoolingEmbedModelInfo, EmbedModelInfo,

8 - LASTPoolingEmbedModelInfo, check_transformers_version)

7 +from ...utils import (CLSPoolingEmbedModelInfo, CLSPoolingRerankModelInfo,

8 + EmbedModelInfo, LASTPoolingEmbedModelInfo,

9 + RerankModelInfo, check_transformers_version)

10 from .embed_utils import correctness_test_embed_models

10 -from .mteb_utils import mteb_test_embed_models

11 +from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models

13 MODELS = [

14 ########## BertModel

@@ -58,6 +59,14 @@ MODELS = [

59 enable_test=False),

60 ]

62 +RERANK_MODELS = [

63 + # classifier_pooling: mean

64 + CLSPoolingRerankModelInfo(

65 + "Alibaba-NLP/gte-reranker-modernbert-base",

66 + architecture="ModernBertForSequenceClassification",

67 + enable_test=True),

68 +]

69 +

71 @pytest.mark.parametrize("model_info", MODELS)

72 def test_embed_models_mteb(hf_runner, vllm_runner,

@@ -88,3 +97,9 @@ def test_embed_models_correctness(hf_runner, vllm_runner,

98 correctness_test_embed_models(hf_runner, vllm_runner, model_info,

99 example_prompts, vllm_extra_kwargs)

100 +

101 +

102 +@pytest.mark.parametrize("model_info", RERANK_MODELS)

103 +def test_rerank_models_mteb(hf_runner, vllm_runner,

104 + model_info: RerankModelInfo) -> None:

105 + mteb_test_rerank_models(hf_runner, vllm_runner, model_info)

vllm/model_executor/models/modernbert.py

+15 -16

@@ -26,8 +26,7 @@ from vllm.model_executor.pooling_metadata import PoolingMetadata

26 from vllm.sequence import IntermediateTensors

27 from vllm.tasks import PoolingTask

29 -from .interfaces import (SupportsCrossEncoding, SupportsV0Only,

30 - default_pooling_type)

29 +from .interfaces import SupportsCrossEncoding, default_pooling_type

30 from .utils import WeightsMapper, maybe_prefix

@@ -93,16 +92,14 @@ class ModernBertAttention(nn.Module):

92 bias=config.attention_bias,

93 )

95 + sliding_window = None

96 if layer_id % config.global_attn_every_n_layers != 0:

97 - self.local_attention = (config.local_attention // 2,

98 - config.local_attention // 2)

97 + sliding_window = config.local_attention // 2

98 + rope_theta = config.local_rope_theta if config.local_rope_theta \

99 + is not None else config.global_rope_theta

100 else:

100 - self.local_attention = (-1, -1)

101 + rope_theta = config.global_rope_theta

102

102 - rope_theta = config.global_rope_theta

103 - if self.local_attention != (

104 - -1, -1) and config.local_rope_theta is not None:

105 - rope_theta = config.local_rope_theta

103 self.rotary_emb = ModernBertRotaryEmbedding(config=config,

104 head_size=self.head_dim,

105 dim=self.head_dim,

@@ -111,7 +108,8 @@ class ModernBertAttention(nn.Module):

108 self.head_dim,

109 self.scaling,

110 prefix=f"{layer_id}.attn",

114 - attn_type=AttentionType.ENCODER_ONLY)

111 + attn_type=AttentionType.ENCODER_ONLY,

112 + per_layer_sliding_window=sliding_window)

113 self.Wo = RowParallelLinear(config.hidden_size,

114 config.hidden_size,

115 bias=config.attention_bias)

@@ -278,6 +276,7 @@ class ModernBertPooler(Pooler):

276 return self.pooling.get_pooling_updates(task)

277

278 def _head(self, pooled_output: torch.Tensor):

279 + pooled_output = pooled_output.to(self.dense.weight.dtype)

280 return self.norm(self.act(self.dense(pooled_output)))

281

282 def forward(

@@ -296,8 +295,7 @@ class ModernBertPooler(Pooler):

295

296

297 @default_pooling_type("CLS")

299 -class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,

300 - SupportsCrossEncoding):

298 +class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

299

300 is_pooling_model = True

301

@@ -308,6 +306,7 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,

306 self.model = ModernBertModel(vllm_config=vllm_config,

307 prefix=maybe_prefix(prefix, "modernbert"))

308 self.classifier = nn.Linear(config.hidden_size, config.num_labels)

309 + self.pooling = ModernBertPooler(config)

310

311 pooler_config = vllm_config.model_config.pooler_config

312 assert pooler_config is not None

@@ -317,14 +316,14 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,

316 Pooler.for_encode(pooler_config),

317 "classify":

318 ClassifierPooler(

320 - pooling=ModernBertPooler(config),

319 + pooling=self.pooling,

320 classifier=self.classifier,

321 act_fn=ClassifierPooler.act_fn_for_seq_cls(

322 vllm_config.model_config),

323 ),

324 "score":

325 ClassifierPooler(

327 - pooling=ModernBertPooler(config),

326 + pooling=self.pooling,

327 classifier=self.classifier,

328 act_fn=ClassifierPooler.act_fn_for_cross_encoder(

329 vllm_config.model_config),

@@ -353,7 +352,7 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,

352 default_weight_loader)

353 weight_loader(param, loaded_weight)

354 if name.startswith("head"):

356 - param = params_dict["_pooler.pooler." + name[len("head") + 1:]]

355 + param = params_dict["pooling." + name[len("head") + 1:]]

356 weight_loader = getattr(param, "weight_loader",

357 default_weight_loader)

358 weight_loader(param, loaded_weight)

@@ -368,5 +367,5 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,

367 return self.model(

368 input_ids=input_ids,

369 inputs_embeds=inputs_embeds,

371 - position_ids=positions,

370 + positions=positions,

371 )

vllm/v1/attention/backends/flash_attn.py

@@ -384,6 +384,8 @@ class FlashAttentionImpl(AttentionImpl):

384 self.alibi_slopes = alibi_slopes

385 if sliding_window is None:

386 self.sliding_window = (-1, -1)

387 + elif attn_type == AttentionType.ENCODER_ONLY:

388 + self.sliding_window = (sliding_window - 1, sliding_window - 1)

389 else:

390 self.sliding_window = (sliding_window - 1, 0)

391 self.kv_cache_dtype = kv_cache_dtype

vllm/v1/worker/gpu_model_runner.py

+66 -40

@@ -826,7 +826,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

826 # Prepare encoder attention metadata separately

827 # (encoder layers are not in KV cache groups)

828 if self.is_encoder_only_model:

829 - common_attn_metadata, encoder_attn_metadata = \

829 +

830 + per_layer_metadata = \

831 self._build_encoder_only_attn_metadata(

832 scheduler_output)

833

@@ -835,6 +836,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

836 self.vllm_config, Attention)

837 for layer_name, attn_module in attention_layers.items():

838 if attn_module.attn_type == AttentionType.ENCODER_ONLY:

839 + common_attn_metadata, encoder_attn_metadata =\

840 + per_layer_metadata[layer_name]

841 attn_metadata[layer_name] = encoder_attn_metadata

842

843 # Prepare the attention metadata for each KV cache group and make layers

@@ -2683,30 +2686,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2686 # Check if model is encoder-only

2687 block_size = self.vllm_config.cache_config.block_size

2688 use_mla = self.vllm_config.model_config.use_mla

2686 - attn_specs = list[AttentionSpec]()

2687 - for attn_module in attn_layers.values():

2689 + attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list)

2690 + for layer_name, attn_module in attn_layers.items():

2691

2692 if attn_module.attn_type == AttentionType.ENCODER_ONLY:

2690 - assert attn_module.sliding_window is None, "Sliding "

2691 - "window attention is not supported for encoder-only models"

2692 -

2693 - attn_specs.append(

2694 - FullAttentionSpec(block_size=block_size,

2695 - num_kv_heads=attn_module.num_kv_heads,

2696 - head_size=attn_module.head_size,

2697 - dtype=self.kv_cache_dtype,

2698 - use_mla=use_mla))

2693 + if attn_module.sliding_window is None:

2694 + attn_spec: AttentionSpec = FullAttentionSpec(

2695 + block_size=block_size,

2696 + num_kv_heads=attn_module.num_kv_heads,

2697 + head_size=attn_module.head_size,

2698 + dtype=self.kv_cache_dtype,

2699 + use_mla=use_mla)

2700 + else:

2701 + attn_spec = SlidingWindowSpec(

2702 + block_size=block_size,

2703 + num_kv_heads=attn_module.num_kv_heads,

2704 + head_size=attn_module.head_size,

2705 + dtype=self.kv_cache_dtype,

2706 + sliding_window=attn_module.sliding_window,

2707 + use_mla=use_mla)

2708 + attn_specs[attn_spec].append(layer_name)

2709 +

2710 else:

2711 raise ValueError("Expected only encoder-only layers")

2712

2713 if len(attn_specs) > 0:

2703 - assert len(attn_specs) == len(attn_layers), \

2704 - "All or none of the layers are expected to be encoder-only"

2714 + total_layers = 0

2715 + for attn_spec, layer_names in attn_specs.items():

2716

2706 - attn_backends = get_attn_backends_for_layers(attn_layers.keys())

2717 + attn_backends = get_attn_backends_for_layers(layer_names)

2718 + total_layers += len(layer_names)

2719

2708 - self.attn_groups.append(

2709 - create_attn_groups(attn_backends, attn_specs[0]))

2720 + self.attn_groups.append(

2721 + create_attn_groups(attn_backends, attn_spec))

2722 + assert total_layers == len(attn_layers), \

2723 + "All or none of the layers are expected to be encoder-only"

2724 self.is_encoder_only_model = True

2725

2726 def calculate_reorder_batch_threshold(self) -> None:

@@ -3071,7 +3085,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3085

3086 def _build_encoder_only_attn_metadata(

3087 self, scheduler_output: "SchedulerOutput") -> \

3074 - tuple[CommonAttentionMetadata, Any]:

3088 + dict[str, tuple[CommonAttentionMetadata, Any]]:

3089 """Prepare encoder attention metadata for encoder-only models.

3090

3091 Args:

@@ -3088,10 +3102,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3102 tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]

3103 max_num_scheduled_tokens = max(tokens)

3104

3091 - # Use the first attention metadata builder

3092 - # to create encoder attention metadata

3093 - builder = self.attn_groups[0][0].metadata_builder

3094 -

3105 dummy_block_table = torch.zeros((num_reqs, 1),

3106 dtype=torch.int32,

3107 device=self.device)

@@ -3099,22 +3109,38 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3109 dtype=torch.int32,

3110 device=self.device)

3111

3102 - common_metadata = CommonAttentionMetadata(

3103 - query_start_loc=self.query_start_loc[:num_reqs + 1],

3104 - query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],

3105 - seq_lens=self.seq_lens[:num_reqs],

3106 - seq_lens_cpu=self.seq_lens_cpu[:num_reqs],

3107 - num_computed_tokens_cpu=self.input_batch.

3108 - num_computed_tokens_cpu_tensor[:num_reqs],

3109 - num_reqs=num_reqs,

3110 - num_actual_tokens=total_num_scheduled_tokens,

3111 - max_query_len=max_num_scheduled_tokens,

3112 - block_table_tensor=dummy_block_table,

3113 - slot_mapping=dummy_slot_mapping,

3114 - causal=False,

3115 - )

3112 + group_metadata = dict[str, tuple[CommonAttentionMetadata, Any]]()

3113

3117 - return common_metadata, builder.build(

3118 - common_prefix_len=0, # No cascade for encoder

3119 - common_attn_metadata=common_metadata,

3120 - )

3114 + for attn_group_list in self.attn_groups:

3115 +

3116 + assert len(attn_group_list) == 1

3117 + attn_group = attn_group_list[0]

3118 +

3119 + # Use the first attention metadata builder

3120 + # to create encoder attention metadata

3121 + builder = attn_group.metadata_builder

3122 +

3123 + common_metadata = CommonAttentionMetadata(

3124 + query_start_loc=self.query_start_loc[:num_reqs + 1],

3125 + query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],

3126 + seq_lens=self.seq_lens[:num_reqs],

3127 + seq_lens_cpu=self.seq_lens_cpu[:num_reqs],

3128 + num_computed_tokens_cpu=self.input_batch.

3129 + num_computed_tokens_cpu_tensor[:num_reqs],

3130 + num_reqs=num_reqs,

3131 + num_actual_tokens=total_num_scheduled_tokens,

3132 + max_query_len=max_num_scheduled_tokens,

3133 + block_table_tensor=dummy_block_table,

3134 + slot_mapping=dummy_slot_mapping,

3135 + causal=False,

3136 + )

3137 +

3138 + metadata = builder.build(

3139 + common_prefix_len=0, # No cascade for encoder

3140 + common_attn_metadata=common_metadata,

3141 + )

3142 +

3143 + for layer_name in attn_group.layer_names:

3144 + group_metadata[layer_name] = (common_metadata, metadata)

3145 +

3146 + return group_metadata