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 7, 05:15 UTC
Order `config.py` in Lexicographical order (#35866)
Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>
Andrii Skliar committed
Mar 5, 2026 at 05:56 UTC
0a12cea25f4a0c2a2ce1c145677a7f54545d8d7d
1 file changed
+370
-370
vllm/model_executor/models/config.py
+370
-370
@@ -28,6 +28,36 @@ class VerifyAndUpdateConfig:
28
return
29
30
31
+class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
32
+ @classmethod
33
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
34
+ """
35
+ Updated fp8 cache to custom "fp8_ds_mla" format for DeepSeekV32
36
+ """
37
+ hf_config = vllm_config.model_config.hf_config
38
+
39
+ # Mirror the check in vllm/model_executor/models/deepseek_v2.py
40
+ is_v32 = hasattr(hf_config, "index_topk")
41
+ assert is_v32
42
+
43
+ # For DeepSeekV3.2, a custom fp8 format is used when fp8 kv-cache is enabled.
44
+ cache_config = vllm_config.cache_config
45
+ if cache_config.cache_dtype.startswith("fp8"):
46
+ cache_config.cache_dtype = "fp8_ds_mla"
47
+ logger.info("Using custom fp8 kv-cache format for DeepSeekV3.2")
48
+ if cache_config.cache_dtype == "bfloat16":
49
+ cache_config.cache_dtype = "auto"
50
+ logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
51
+
52
+
53
+class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig):
54
+ @staticmethod
55
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
56
+ # Ernie4.5-VL conditionally executes text/vision MoE branches, so
57
+ # fast_moe_cold_start can silently produce incorrect execution order.
58
+ vllm_config.compilation_config.fast_moe_cold_start = False
59
+
60
+
61
class Gemma3TextModelConfig(VerifyAndUpdateConfig):
62
@staticmethod
63
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -35,6 +65,29 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):
65
hf_config.is_causal = not hf_config.use_bidirectional_attention
66
67
68
+class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
69
+ @staticmethod
70
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
71
+ structured_outputs_config = vllm_config.structured_outputs_config
72
+ if structured_outputs_config.reasoning_parser == "":
73
+ structured_outputs_config.reasoning_parser = "openai_gptoss"
74
+
75
+ # Increase the max capture size from 512 to 1024 for performance.
76
+ # NOTE(woosuk): This will increase the number of CUDA graphs
77
+ # from 67 to 83.
78
+ compilation_config = vllm_config.compilation_config
79
+ # Only override when the user has not set either of
80
+ # cudagraph_capture_sizes or max_cudagraph_capture_size.
81
+ if (
82
+ compilation_config.cudagraph_capture_sizes is None
83
+ and compilation_config.max_cudagraph_capture_size is None
84
+ ):
85
+ compilation_config.max_cudagraph_capture_size = 1024
86
+ logger.info(
87
+ "Overriding max cuda graph capture size to %d for performance.", 1024
88
+ )
89
+
90
+
91
class GteNewModelConfig(VerifyAndUpdateConfig):
92
@staticmethod
93
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -55,6 +108,166 @@ class GteNewModelConfig(VerifyAndUpdateConfig):
108
}
109
110
111
+class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
112
+ @classmethod
113
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
114
+ """
115
+ Ensure that page size of attention layers is greater than or
116
+ equal to the mamba layers. If not, automatically set the attention
117
+ block size to ensure that it is. If the attention page size is
118
+ strictly greater than the mamba page size, we pad the mamba page size
119
+ to make them equal.
120
+
121
+ Args:
122
+ vllm_config: vLLM Config
123
+ """
124
+ # Save the user input before it gets modified by MambaModelConfig
125
+ mamba_block_size = vllm_config.cache_config.mamba_block_size
126
+ # Enable FULL_AND_PIECEWISE by default
127
+ MambaModelConfig.verify_and_update_config(vllm_config)
128
+
129
+ attention_config = vllm_config.attention_config
130
+ cache_config = vllm_config.cache_config
131
+ model_config = vllm_config.model_config
132
+ parallel_config = vllm_config.parallel_config
133
+
134
+ if cache_config.cache_dtype == "auto":
135
+ kv_cache_dtype = model_config.dtype
136
+ else:
137
+ kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
138
+
139
+ # get attention page size (for 1 token)
140
+ # Attention backend constraints:
141
+ # - FlashAttention (FA) requires block size to be multiple of 16
142
+ # - MLA (Multi-head Latent Attention) requires larger alignment:
143
+ # * CUTLASS_MLA backend: kernel_block_size 128 alignment
144
+ # * Other MLA backends: kernel_block_size 64 alignment
145
+ if model_config.use_mla:
146
+ use_cutlass_mla = (
147
+ attention_config.backend == AttentionBackendEnum.CUTLASS_MLA
148
+ )
149
+ kernel_block_alignment_size = 128 if use_cutlass_mla else 64
150
+ attn_page_size_1_token = MLAAttentionSpec(
151
+ block_size=1,
152
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
153
+ head_size=model_config.get_head_size(),
154
+ dtype=kv_cache_dtype,
155
+ ).page_size_bytes
156
+ else:
157
+ kernel_block_alignment_size = 16
158
+ if (
159
+ current_platform.is_device_capability_family(100)
160
+ and model_config.get_head_size() == 256
161
+ and (
162
+ attention_config.backend is None
163
+ or attention_config.backend == AttentionBackendEnum.FLASHINFER
164
+ )
165
+ ):
166
+ # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that`
167
+ # head size 256 and block size 16 is not supported on blackwell.
168
+ kernel_block_alignment_size = 32
169
+ attn_page_size_1_token = FullAttentionSpec(
170
+ block_size=1,
171
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
172
+ head_size=model_config.get_head_size(),
173
+ dtype=kv_cache_dtype,
174
+ ).page_size_bytes
175
+
176
+ model_cls, _ = ModelRegistry.resolve_model_cls(
177
+ model_config.architecture,
178
+ model_config=model_config,
179
+ )
180
+
181
+ # get mamba page size
182
+ mamba_page_size = MambaSpec(
183
+ shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
184
+ dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
185
+ block_size=-1, # block_size doesn't matter for mamba page size
186
+ ).page_size_bytes
187
+
188
+ # Model may be marked as is_hybrid
189
+ # but mamba is skipped via config,
190
+ # return directly
191
+ if mamba_page_size == 0:
192
+ return
193
+
194
+ if cache_config.mamba_cache_mode == "all":
195
+ # With prefix caching, select attention block size to
196
+ # optimize for mamba kernel performance
197
+
198
+ # Mamba2 SSD kernel uses a chunk_size, e.g. 256
199
+ # Align the block to the kernel: use lowest multiple of chunk_size
200
+ # of attention tokens that would fit mamba_page_size:
201
+ # e.g. for mamba page size = 788kB
202
+ # attn_1_token = 2kB -> fits ~394 tokens
203
+ # then round up to a multiple of 256 -> 512 tokens
204
+ # End result:
205
+ # attn_block_size = 512
206
+ # mamba_block_size = 512 (aligned to a multiple of chunk_size)
207
+ # TODO(tdoublep): this constraint can be relaxed fairly
208
+ # easily by changing the way we layout chunks in the
209
+ # mamba2 kernels.
210
+
211
+ base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
212
+ attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
213
+ chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
214
+ attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
215
+ cache_config.mamba_block_size = attn_block_size
216
+ else:
217
+ # Without prefix caching, select minimum valid attention block size
218
+ # to minimize mamba state padding
219
+
220
+ # Calculate minimum attention block size that satisfies both:
221
+ # 1. Backend alignment requirements (kernel_block_alignment_size)
222
+ # 2. Mamba page size compatibility (attn_page_size >= mamba_page_size)
223
+ attn_block_size = kernel_block_alignment_size * cdiv(
224
+ mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token
225
+ )
226
+
227
+ # override attention block size if either (a) the
228
+ # user has not set it or (b) the user has set it
229
+ # too small.
230
+ if cache_config.block_size is None or cache_config.block_size < attn_block_size:
231
+ cache_config.block_size = attn_block_size
232
+ logger.info(
233
+ "Setting attention block size to %d tokens "
234
+ "to ensure that attention page size is >= mamba page size.",
235
+ attn_block_size,
236
+ )
237
+
238
+ # By default, mamba block size will be set to max_model_len.
239
+ # When enabling prefix caching and using align mamba cache
240
+ # mode, we align mamba block size to the block size as the
241
+ # basic granularity for prefix caching.
242
+ if cache_config.mamba_cache_mode == "align":
243
+ cache_config.mamba_block_size = cache_config.block_size
244
+
245
+ # compute new attention page size
246
+ attn_page_size = cache_config.block_size * attn_page_size_1_token
247
+
248
+ assert attn_page_size >= mamba_page_size
249
+
250
+ if attn_page_size == mamba_page_size:
251
+ # don't need to pad mamba page size
252
+ return
253
+
254
+ # pad mamba page size to exactly match attention
255
+ if (
256
+ cache_config.mamba_page_size_padded is None
257
+ or cache_config.mamba_page_size_padded != attn_page_size
258
+ ):
259
+ cache_config.mamba_page_size_padded = attn_page_size
260
+ mamba_padding_pct = (
261
+ 100 * (attn_page_size - mamba_page_size) / mamba_page_size
262
+ )
263
+ logger.info(
264
+ "Padding mamba page size by %.2f%% to ensure "
265
+ "that mamba page size and attention page size are "
266
+ "exactly equal.",
267
+ mamba_padding_pct,
268
+ )
269
+
270
+
271
class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
272
@staticmethod
273
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -91,6 +304,16 @@ class JinaRobertaModelConfig(VerifyAndUpdateConfig):
304
}
305
306
307
+class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):
308
+ @staticmethod
309
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
310
+ config = model_config.hf_config
311
+ config.num_labels = 1
312
+ pooler_config = model_config.pooler_config
313
+ if pooler_config.logit_bias is None:
314
+ pooler_config.logit_bias = 2.65
315
+
316
+
317
class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
318
@staticmethod
319
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -148,30 +371,119 @@ class LlamaNemotronVLConfig(VerifyAndUpdateConfig):
371
model_config.pooler_config.seq_pooling_type = pooling_type
372
373
151
-class NomicBertModelConfig(VerifyAndUpdateConfig):
152
- @staticmethod
153
- def verify_and_update_model_config(model_config: "ModelConfig") -> None:
154
- config = model_config.hf_config
155
-
156
- assert config.__class__.__name__ == "NomicBertConfig"
157
- assert config.activation_function in ["swiglu", "gelu"]
158
- config.position_embedding_type = getattr(
159
- config, "position_embedding_type", "rope"
160
- )
161
-
162
- if config.activation_function == "swiglu":
163
- config.hidden_act = "silu"
164
- else:
165
- config.hidden_act = config.activation_function
166
-
167
- assert config.mlp_fc1_bias == config.mlp_fc2_bias == config.qkv_proj_bias
168
- config.bias = config.qkv_proj_bias
374
+class MambaModelConfig(VerifyAndUpdateConfig):
375
+ @classmethod
376
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
377
+ """
378
+ Enable FULL_AND_PIECEWISE cuda graph mode by default (required
379
+ to get good performance for mamba layers in V1).
380
170
- assert config.rotary_emb_scale_base is None
171
- assert not config.rotary_emb_interleaved
381
+ Args:
382
+ vllm_config: vLLM Config
383
+ """
384
+ model_config = vllm_config.model_config
385
+ cache_config = vllm_config.cache_config
386
173
- config.layer_norm_eps = config.layer_norm_epsilon
174
- config.intermediate_size = config.n_inner
387
+ if cache_config.enable_prefix_caching:
388
+ if cache_config.mamba_cache_mode == "none":
389
+ cache_config.mamba_cache_mode = (
390
+ "all" if model_config.supports_mamba_prefix_caching else "align"
391
+ )
392
+ logger.warning(
393
+ "Mamba cache mode is set to '%s' for %s by default "
394
+ "when prefix caching is enabled",
395
+ cache_config.mamba_cache_mode,
396
+ model_config.architecture,
397
+ )
398
+ if (
399
+ cache_config.mamba_cache_mode == "all"
400
+ and not model_config.supports_mamba_prefix_caching
401
+ ):
402
+ cache_config.mamba_cache_mode = "align"
403
+ logger.warning(
404
+ "Hybrid or mamba-based model detected without support "
405
+ "for prefix caching with Mamba cache 'all' mode: "
406
+ "falling back to 'align' mode."
407
+ )
408
+ if cache_config.mamba_cache_mode == "align":
409
+ assert vllm_config.scheduler_config.enable_chunked_prefill, (
410
+ "Chunked prefill is required for mamba cache mode 'align'."
411
+ )
412
+ logger.info(
413
+ "Warning: Prefix caching in Mamba cache '%s' "
414
+ "mode is currently enabled. "
415
+ "Its support for Mamba layers is experimental. "
416
+ "Please report any issues you may observe.",
417
+ cache_config.mamba_cache_mode,
418
+ )
419
+ # By default, mamba block size will be set to max_model_len (see
420
+ # below). When enabling prefix caching, we align mamba block size
421
+ # to the block size as the basic granularity for prefix caching.
422
+ if cache_config.mamba_block_size is None:
423
+ cache_config.mamba_block_size = cache_config.block_size
424
+ else:
425
+ if cache_config.mamba_cache_mode != "none":
426
+ cache_config.mamba_cache_mode = "none"
427
+ logger.warning(
428
+ "Mamba cache mode is set to 'none' when prefix caching is disabled"
429
+ )
430
+ if cache_config.mamba_block_size is None:
431
+ cache_config.mamba_block_size = model_config.max_model_len
432
+
433
+
434
+class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
435
+ @staticmethod
436
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
437
+ """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto'
438
+ (or not explicitly set), to the value specified in the HF config, or to
439
+ float16 if not specified.
440
+ """
441
+ cache_config = vllm_config.cache_config
442
+ if cache_config.mamba_ssm_cache_dtype == "auto":
443
+ hf_config = vllm_config.model_config.hf_config
444
+ mamba_ssm_cache_dtype = getattr(
445
+ hf_config, "mamba_ssm_cache_dtype", "float16"
446
+ )
447
+ logger.info(
448
+ "Updating mamba_ssm_cache_dtype to '%s' for NemotronH model",
449
+ mamba_ssm_cache_dtype,
450
+ )
451
+ cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype
452
+
453
+
454
+class NemotronHNanoVLV2Config(VerifyAndUpdateConfig):
455
+ @staticmethod
456
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
457
+ mm_config = model_config.multimodal_config
458
+ if mm_config is not None:
459
+ video_kwargs = mm_config.media_io_kwargs.setdefault("video", {})
460
+ video_kwargs.setdefault("video_backend", "nemotron_vl")
461
+
462
+
463
+class NomicBertModelConfig(VerifyAndUpdateConfig):
464
+ @staticmethod
465
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
466
+ config = model_config.hf_config
467
+
468
+ assert config.__class__.__name__ == "NomicBertConfig"
469
+ assert config.activation_function in ["swiglu", "gelu"]
470
+ config.position_embedding_type = getattr(
471
+ config, "position_embedding_type", "rope"
472
+ )
473
+
474
+ if config.activation_function == "swiglu":
475
+ config.hidden_act = "silu"
476
+ else:
477
+ config.hidden_act = config.activation_function
478
+
479
+ assert config.mlp_fc1_bias == config.mlp_fc2_bias == config.qkv_proj_bias
480
+ config.bias = config.qkv_proj_bias
481
+
482
+ assert config.rotary_emb_scale_base is None
483
+ assert not config.rotary_emb_interleaved
484
+
485
+ config.layer_norm_eps = config.layer_norm_epsilon
486
+ config.intermediate_size = config.n_inner
487
config.hidden_size = config.n_embd
488
config.num_hidden_layers = config.n_layer
489
model_config.model_arch_config.hidden_size = config.hidden_size
@@ -299,338 +611,6 @@ class Qwen3VLForSequenceClassificationConfig(Qwen3ForSequenceClassificationConfi
611
pass
612
613
302
-class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):
303
- @staticmethod
304
- def verify_and_update_model_config(model_config: "ModelConfig") -> None:
305
- config = model_config.hf_config
306
- config.num_labels = 1
307
- pooler_config = model_config.pooler_config
308
- if pooler_config.logit_bias is None:
309
- pooler_config.logit_bias = 2.65
310
-
311
-
312
-class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
313
- @staticmethod
314
- def verify_and_update_model_config(model_config: "ModelConfig") -> None:
315
- config = model_config.hf_config
316
-
317
- assert config.__class__.__name__ == "GteConfig"
318
- assert config.hidden_act == "gelu"
319
-
320
- config.hidden_act = "geglu"
321
-
322
- head_dim = config.hidden_size // config.num_attention_heads
323
- rotary_dim = getattr(config, "rotary_emb_dim", head_dim)
324
- config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim
325
- config.rotary_kwargs = {
326
- "head_size": head_dim,
327
- "max_position": config.max_position_embeddings,
328
- "rope_parameters": config.rope_parameters,
329
- }
330
-
331
-
332
-class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig):
333
- @staticmethod
334
- def verify_and_update_config(vllm_config: "VllmConfig") -> None:
335
- # Ernie4.5-VL conditionally executes text/vision MoE branches, so
336
- # fast_moe_cold_start can silently produce incorrect execution order.
337
- vllm_config.compilation_config.fast_moe_cold_start = False
338
-
339
-
340
-class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
341
- @staticmethod
342
- def verify_and_update_config(vllm_config: "VllmConfig") -> None:
343
- structured_outputs_config = vllm_config.structured_outputs_config
344
- if structured_outputs_config.reasoning_parser == "":
345
- structured_outputs_config.reasoning_parser = "openai_gptoss"
346
-
347
- # Increase the max capture size from 512 to 1024 for performance.
348
- # NOTE(woosuk): This will increase the number of CUDA graphs
349
- # from 67 to 83.
350
- compilation_config = vllm_config.compilation_config
351
- # Only override when the user has not set either of
352
- # cudagraph_capture_sizes or max_cudagraph_capture_size.
353
- if (
354
- compilation_config.cudagraph_capture_sizes is None
355
- and compilation_config.max_cudagraph_capture_size is None
356
- ):
357
- compilation_config.max_cudagraph_capture_size = 1024
358
- logger.info(
359
- "Overriding max cuda graph capture size to %d for performance.", 1024
360
- )
361
-
362
-
363
-class MambaModelConfig(VerifyAndUpdateConfig):
364
- @classmethod
365
- def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
366
- """
367
- Enable FULL_AND_PIECEWISE cuda graph mode by default (required
368
- to get good performance for mamba layers in V1).
369
-
370
- Args:
371
- vllm_config: vLLM Config
372
- """
373
- model_config = vllm_config.model_config
374
- cache_config = vllm_config.cache_config
375
-
376
- if cache_config.enable_prefix_caching:
377
- if cache_config.mamba_cache_mode == "none":
378
- cache_config.mamba_cache_mode = (
379
- "all" if model_config.supports_mamba_prefix_caching else "align"
380
- )
381
- logger.warning(
382
- "Mamba cache mode is set to '%s' for %s by default "
383
- "when prefix caching is enabled",
384
- cache_config.mamba_cache_mode,
385
- model_config.architecture,
386
- )
387
- if (
388
- cache_config.mamba_cache_mode == "all"
389
- and not model_config.supports_mamba_prefix_caching
390
- ):
391
- cache_config.mamba_cache_mode = "align"
392
- logger.warning(
393
- "Hybrid or mamba-based model detected without support "
394
- "for prefix caching with Mamba cache 'all' mode: "
395
- "falling back to 'align' mode."
396
- )
397
- if cache_config.mamba_cache_mode == "align":
398
- assert vllm_config.scheduler_config.enable_chunked_prefill, (
399
- "Chunked prefill is required for mamba cache mode 'align'."
400
- )
401
- logger.info(
402
- "Warning: Prefix caching in Mamba cache '%s' "
403
- "mode is currently enabled. "
404
- "Its support for Mamba layers is experimental. "
405
- "Please report any issues you may observe.",
406
- cache_config.mamba_cache_mode,
407
- )
408
- # By default, mamba block size will be set to max_model_len (see
409
- # below). When enabling prefix caching, we align mamba block size
410
- # to the block size as the basic granularity for prefix caching.
411
- if cache_config.mamba_block_size is None:
412
- cache_config.mamba_block_size = cache_config.block_size
413
- else:
414
- if cache_config.mamba_cache_mode != "none":
415
- cache_config.mamba_cache_mode = "none"
416
- logger.warning(
417
- "Mamba cache mode is set to 'none' when prefix caching is disabled"
418
- )
419
- if cache_config.mamba_block_size is None:
420
- cache_config.mamba_block_size = model_config.max_model_len
421
-
422
-
423
-class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
424
- @classmethod
425
- def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
426
- """
427
- Ensure that page size of attention layers is greater than or
428
- equal to the mamba layers. If not, automatically set the attention
429
- block size to ensure that it is. If the attention page size is
430
- strictly greater than the mamba page size, we pad the mamba page size
431
- to make them equal.
432
-
433
- Args:
434
- vllm_config: vLLM Config
435
- """
436
- # Save the user input before it gets modified by MambaModelConfig
437
- mamba_block_size = vllm_config.cache_config.mamba_block_size
438
- # Enable FULL_AND_PIECEWISE by default
439
- MambaModelConfig.verify_and_update_config(vllm_config)
440
-
441
- attention_config = vllm_config.attention_config
442
- cache_config = vllm_config.cache_config
443
- model_config = vllm_config.model_config
444
- parallel_config = vllm_config.parallel_config
445
-
446
- if cache_config.cache_dtype == "auto":
447
- kv_cache_dtype = model_config.dtype
448
- else:
449
- kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
450
-
451
- # get attention page size (for 1 token)
452
- # Attention backend constraints:
453
- # - FlashAttention (FA) requires block size to be multiple of 16
454
- # - MLA (Multi-head Latent Attention) requires larger alignment:
455
- # * CUTLASS_MLA backend: kernel_block_size 128 alignment
456
- # * Other MLA backends: kernel_block_size 64 alignment
457
- if model_config.use_mla:
458
- use_cutlass_mla = (
459
- attention_config.backend == AttentionBackendEnum.CUTLASS_MLA
460
- )
461
- kernel_block_alignment_size = 128 if use_cutlass_mla else 64
462
- attn_page_size_1_token = MLAAttentionSpec(
463
- block_size=1,
464
- num_kv_heads=model_config.get_num_kv_heads(parallel_config),
465
- head_size=model_config.get_head_size(),
466
- dtype=kv_cache_dtype,
467
- ).page_size_bytes
468
- else:
469
- kernel_block_alignment_size = 16
470
- if (
471
- current_platform.is_device_capability_family(100)
472
- and model_config.get_head_size() == 256
473
- and (
474
- attention_config.backend is None
475
- or attention_config.backend == AttentionBackendEnum.FLASHINFER
476
- )
477
- ):
478
- # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that`
479
- # head size 256 and block size 16 is not supported on blackwell.
480
- kernel_block_alignment_size = 32
481
- attn_page_size_1_token = FullAttentionSpec(
482
- block_size=1,
483
- num_kv_heads=model_config.get_num_kv_heads(parallel_config),
484
- head_size=model_config.get_head_size(),
485
- dtype=kv_cache_dtype,
486
- ).page_size_bytes
487
-
488
- model_cls, _ = ModelRegistry.resolve_model_cls(
489
- model_config.architecture,
490
- model_config=model_config,
491
- )
492
-
493
- # get mamba page size
494
- mamba_page_size = MambaSpec(
495
- shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
496
- dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
497
- block_size=-1, # block_size doesn't matter for mamba page size
498
- ).page_size_bytes
499
-
500
- # Model may be marked as is_hybrid
501
- # but mamba is skipped via config,
502
- # return directly
503
- if mamba_page_size == 0:
504
- return
505
-
506
- if cache_config.mamba_cache_mode == "all":
507
- # With prefix caching, select attention block size to
508
- # optimize for mamba kernel performance
509
-
510
- # Mamba2 SSD kernel uses a chunk_size, e.g. 256
511
- # Align the block to the kernel: use lowest multiple of chunk_size
512
- # of attention tokens that would fit mamba_page_size:
513
- # e.g. for mamba page size = 788kB
514
- # attn_1_token = 2kB -> fits ~394 tokens
515
- # then round up to a multiple of 256 -> 512 tokens
516
- # End result:
517
- # attn_block_size = 512
518
- # mamba_block_size = 512 (aligned to a multiple of chunk_size)
519
- # TODO(tdoublep): this constraint can be relaxed fairly
520
- # easily by changing the way we layout chunks in the
521
- # mamba2 kernels.
522
-
523
- base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
524
- attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
525
- chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
526
- attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
527
- cache_config.mamba_block_size = attn_block_size
528
- else:
529
- # Without prefix caching, select minimum valid attention block size
530
- # to minimize mamba state padding
531
-
532
- # Calculate minimum attention block size that satisfies both:
533
- # 1. Backend alignment requirements (kernel_block_alignment_size)
534
- # 2. Mamba page size compatibility (attn_page_size >= mamba_page_size)
535
- attn_block_size = kernel_block_alignment_size * cdiv(
536
- mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token
537
- )
538
-
539
- # override attention block size if either (a) the
540
- # user has not set it or (b) the user has set it
541
- # too small.
542
- if cache_config.block_size is None or cache_config.block_size < attn_block_size:
543
- cache_config.block_size = attn_block_size
544
- logger.info(
545
- "Setting attention block size to %d tokens "
546
- "to ensure that attention page size is >= mamba page size.",
547
- attn_block_size,
548
- )
549
-
550
- # By default, mamba block size will be set to max_model_len.
551
- # When enabling prefix caching and using align mamba cache
552
- # mode, we align mamba block size to the block size as the
553
- # basic granularity for prefix caching.
554
- if cache_config.mamba_cache_mode == "align":
555
- cache_config.mamba_block_size = cache_config.block_size
556
-
557
- # compute new attention page size
558
- attn_page_size = cache_config.block_size * attn_page_size_1_token
559
-
560
- assert attn_page_size >= mamba_page_size
561
-
562
- if attn_page_size == mamba_page_size:
563
- # don't need to pad mamba page size
564
- return
565
-
566
- # pad mamba page size to exactly match attention
567
- if (
568
- cache_config.mamba_page_size_padded is None
569
- or cache_config.mamba_page_size_padded != attn_page_size
570
- ):
571
- cache_config.mamba_page_size_padded = attn_page_size
572
- mamba_padding_pct = (
573
- 100 * (attn_page_size - mamba_page_size) / mamba_page_size
574
- )
575
- logger.info(
576
- "Padding mamba page size by %.2f%% to ensure "
577
- "that mamba page size and attention page size are "
578
- "exactly equal.",
579
- mamba_padding_pct,
580
- )
581
-
582
-
583
-class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
584
- @classmethod
585
- def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
586
- """
587
- Updated fp8 cache to custom "fp8_ds_mla" format for DeepSeekV32
588
- """
589
- hf_config = vllm_config.model_config.hf_config
590
-
591
- # Mirror the check in vllm/model_executor/models/deepseek_v2.py
592
- is_v32 = hasattr(hf_config, "index_topk")
593
- assert is_v32
594
-
595
- # For DeepSeekV3.2, a custom fp8 format is used when fp8 kv-cache is enabled.
596
- cache_config = vllm_config.cache_config
597
- if cache_config.cache_dtype.startswith("fp8"):
598
- cache_config.cache_dtype = "fp8_ds_mla"
599
- logger.info("Using custom fp8 kv-cache format for DeepSeekV3.2")
600
- if cache_config.cache_dtype == "bfloat16":
601
- cache_config.cache_dtype = "auto"
602
- logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
603
-
604
-
605
-class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
606
- @staticmethod
607
- def verify_and_update_config(vllm_config: "VllmConfig") -> None:
608
- """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto'
609
- (or not explicitly set), to the value specified in the HF config, or to
610
- float16 if not specified.
611
- """
612
- cache_config = vllm_config.cache_config
613
- if cache_config.mamba_ssm_cache_dtype == "auto":
614
- hf_config = vllm_config.model_config.hf_config
615
- mamba_ssm_cache_dtype = getattr(
616
- hf_config, "mamba_ssm_cache_dtype", "float16"
617
- )
618
- logger.info(
619
- "Updating mamba_ssm_cache_dtype to '%s' for NemotronH model",
620
- mamba_ssm_cache_dtype,
621
- )
622
- cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype
623
-
624
-
625
-class NemotronHNanoVLV2Config(VerifyAndUpdateConfig):
626
- @staticmethod
627
- def verify_and_update_model_config(model_config: "ModelConfig") -> None:
628
- mm_config = model_config.multimodal_config
629
- if mm_config is not None:
630
- video_kwargs = mm_config.media_io_kwargs.setdefault("video", {})
631
- video_kwargs.setdefault("video_backend", "nemotron_vl")
632
-
633
-
614
class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
615
@staticmethod
616
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -658,6 +638,26 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
638
)
639
640
641
+class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
642
+ @staticmethod
643
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
644
+ config = model_config.hf_config
645
+
646
+ assert config.__class__.__name__ == "GteConfig"
647
+ assert config.hidden_act == "gelu"
648
+
649
+ config.hidden_act = "geglu"
650
+
651
+ head_dim = config.hidden_size // config.num_attention_heads
652
+ rotary_dim = getattr(config, "rotary_emb_dim", head_dim)
653
+ config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim
654
+ config.rotary_kwargs = {
655
+ "head_size": head_dim,
656
+ "max_position": config.max_position_embeddings,
657
+ "rope_parameters": config.rope_parameters,
658
+ }
659
+
660
+
661
class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
662
@staticmethod
663
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -666,33 +666,33 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
666
667
668
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
669
+ "ColBERTJinaRobertaModel": JinaRobertaModelConfig,
670
+ "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
671
+ "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
672
+ "FalconMambaForCausalLM": MambaModelConfig,
673
+ "Gemma3TextModel": Gemma3TextModelConfig,
674
+ "GptOssForCausalLM": GptOssForCausalLMConfig,
675
"GteModel": SnowflakeGteNewModelConfig,
670
- "GteNewModel": GteNewModelConfig,
676
"GteNewForSequenceClassification": GteNewModelConfig,
672
- "Gemma3TextModel": Gemma3TextModelConfig,
673
- "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config,
677
+ "GteNewModel": GteNewModelConfig,
678
+ "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
679
+ "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
680
"LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
681
"LlamaBidirectionalModel": LlamaBidirectionalConfig,
676
- "LlamaNemotronVLModel": LlamaNemotronVLConfig,
682
"LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig,
683
+ "LlamaNemotronVLModel": LlamaNemotronVLConfig,
684
+ "Mamba2ForCausalLM": MambaModelConfig,
685
+ "MambaForCausalLM": MambaModelConfig,
686
+ "NemotronHForCausalLM": NemotronHForCausalLMConfig,
687
+ "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
688
+ "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config,
689
"NomicBertModel": NomicBertModelConfig,
690
"Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
691
"Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
692
"Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,
693
"Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,
683
- "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
684
- "XLMRobertaModel": JinaRobertaModelConfig,
685
- "ColBERTJinaRobertaModel": JinaRobertaModelConfig,
686
- "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
687
- "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
688
- "GptOssForCausalLM": GptOssForCausalLMConfig,
689
- "MambaForCausalLM": MambaModelConfig,
690
- "Mamba2ForCausalLM": MambaModelConfig,
691
- "FalconMambaForCausalLM": MambaModelConfig,
692
- "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
693
- "NemotronHForCausalLM": NemotronHForCausalLMConfig,
694
- "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
694
"Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
695
"Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
696
"VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
697
+ "XLMRobertaModel": JinaRobertaModelConfig,
698
}