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 15, 08:30 UTC
[CI] Enable test_initialization to run on V1 (#16736)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 23, 2025 at 18:09 UTC
0ddf88e16e6ef4d985716f5bdec60fd053a260fa
5 files changed
+54
-45
.buildkite/test-pipeline.yaml
+1
-4
@@ -472,10 +472,7 @@ steps:
472
- pytest -v -s models/test_registry.py
473
- pytest -v -s models/test_utils.py
474
- pytest -v -s models/test_vision.py
475
- # V1 Test: https://github.com/vllm-project/vllm/issues/14531
476
- - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'
477
- - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'llama4'
478
- - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'plamo2'
475
+ - pytest -v -s models/test_initialization.py
476
477
- label: Language Models Test (Standard)
478
mirror_hardwares: [amdexperimental]
tests/models/registry.py
+24
-8
@@ -55,9 +55,18 @@ class _HfExamplesInfo:
55
trust_remote_code: bool = False
56
"""The ``trust_remote_code`` level required to load the model."""
57
58
+ v0_only: bool = False
59
+ """The model is only available with the vLLM V0 engine."""
60
+
61
hf_overrides: dict[str, Any] = field(default_factory=dict)
62
"""The ``hf_overrides`` required to load the model."""
63
64
+ max_model_len: Optional[int] = None
65
+ """
66
+ The maximum model length to use for this model. Some models default to a
67
+ length that is too large to fit into memory in CI.
68
+ """
69
+
70
def check_transformers_version(
71
self,
72
*,
@@ -215,10 +224,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
224
"OrionForCausalLM": _HfExamplesInfo("OrionStarAI/Orion-14B-Chat",
225
trust_remote_code=True),
226
"PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),
218
- "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
227
+ "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2", v0_only=True),
228
"Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),
229
"Phi3SmallForCausalLM": _HfExamplesInfo("microsoft/Phi-3-small-8k-instruct",
221
- trust_remote_code=True),
230
+ trust_remote_code=True,
231
+ v0_only=True),
232
"PhiMoEForCausalLM": _HfExamplesInfo("microsoft/Phi-3.5-MoE-instruct",
233
trust_remote_code=True),
234
"Plamo2ForCausalLM": _HfExamplesInfo("pfnet/plamo-2-1b",
@@ -234,7 +244,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
244
is_available_online=False),
245
"StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b", # noqa: E501
246
is_available_online=False),
237
- "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
247
+ "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t",
248
+ v0_only=True),
249
"Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),
250
"SolarForCausalLM": _HfExamplesInfo("upstage/solar-pro-preview-instruct"),
251
"TeleChat2ForCausalLM": _HfExamplesInfo("Tele-AI/TeleChat2-3B",
@@ -303,7 +314,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
314
"AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),
315
"AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereForAI/aya-vision-8b"), # noqa: E501
316
"Blip2ForConditionalGeneration": _HfExamplesInfo("Salesforce/blip2-opt-2.7b", # noqa: E501
306
- extras={"6b": "Salesforce/blip2-opt-6.7b"}), # noqa: E501
317
+ extras={"6b": "Salesforce/blip2-opt-6.7b"}, # noqa: E501
318
+ v0_only=True),
319
"ChameleonForConditionalGeneration": _HfExamplesInfo("facebook/chameleon-7b"), # noqa: E501
320
"DeepseekVLV2ForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-vl2-tiny", # noqa: E501
321
extras={"fork": "Isotr0py/deepseek-vl2-tiny"}, # noqa: E501
@@ -328,9 +340,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
340
{"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"}), # noqa: E501
341
"KimiVLForConditionalGeneration": _HfExamplesInfo("moonshotai/Kimi-VL-A3B-Instruct", # noqa: E501
342
extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, # noqa: E501
331
- trust_remote_code=True),
343
+ trust_remote_code=True,
344
+ v0_only=True),
345
"Llama4ForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct", # noqa: E501
333
- min_transformers_version="4.51"),
346
+ min_transformers_version="4.51",
347
+ max_model_len=10240),
348
"LlavaForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-1.5-7b-hf",
349
extras={"mistral": "mistral-community/pixtral-12b", # noqa: E501
350
"mistral-fp8": "nm-testing/pixtral-12b-FP8-dynamic"}), # noqa: E501
@@ -349,7 +363,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
363
extras={"2.6": "openbmb/MiniCPM-V-2_6"}, # noqa: E501
364
trust_remote_code=True),
365
"MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo("MiniMaxAI/MiniMax-VL-01", # noqa: E501
352
- trust_remote_code=True),
366
+ trust_remote_code=True,
367
+ v0_only=True),
368
"Mistral3ForConditionalGeneration": _HfExamplesInfo("mistralai/Mistral-Small-3.1-24B-Instruct-2503", # noqa: E501
369
extras={"fp8": "nm-testing/Mistral-Small-3.1-24B-Instruct-2503-FP8-dynamic"}), # noqa: E501
370
"MolmoForCausalLM": _HfExamplesInfo("allenai/Molmo-7B-D-0924",
@@ -372,7 +387,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
387
"Phi4MMForCausalLM": _HfExamplesInfo("microsoft/Phi-4-multimodal-instruct",
388
trust_remote_code=True),
389
"PixtralForConditionalGeneration": _HfExamplesInfo("mistralai/Pixtral-12B-2409", # noqa: E501
375
- tokenizer_mode="mistral"),
390
+ tokenizer_mode="mistral",
391
+ v0_only=True),
392
"QwenVLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen-VL",
393
extras={"chat": "Qwen/Qwen-VL-Chat"}, # noqa: E501
394
trust_remote_code=True,
tests/models/test_initialization.py
+13
-4
@@ -15,12 +15,12 @@ from .registry import HF_EXAMPLE_MODELS
15
16
17
@pytest.mark.parametrize("model_arch", HF_EXAMPLE_MODELS.get_supported_archs())
18
-def test_can_initialize(model_arch):
18
+def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):
19
model_info = HF_EXAMPLE_MODELS.get_hf_info(model_arch)
20
model_info.check_available_online(on_fail="skip")
21
model_info.check_transformers_version(on_fail="skip")
22
23
- # Avoid OOM
23
+ # Avoid OOM and reduce initialization time by only using 1 layer
24
def hf_overrides(hf_config: PretrainedConfig) -> PretrainedConfig:
25
hf_config.update(model_info.hf_overrides)
26
@@ -34,6 +34,12 @@ def test_can_initialize(model_arch):
34
"num_local_experts": 2,
35
})
36
37
+ if hasattr(hf_config, "vision_config"):
38
+ hf_config.vision_config.update({
39
+ "num_layers": 1,
40
+ "num_hidden_layers": 1,
41
+ })
42
+
43
return hf_config
44
45
# Avoid calling model.forward()
@@ -46,7 +52,7 @@ def test_can_initialize(model_arch):
52
scheduler_kv_cache_config = get_kv_cache_config(
53
vllm_config,
54
kv_cache_specs[0],
49
- 20 * GiB_bytes,
55
+ 10 * GiB_bytes,
56
)
57
58
# gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
@@ -55,7 +61,9 @@ def test_can_initialize(model_arch):
61
with (patch.object(V0LLMEngine, "_initialize_kv_caches",
62
_initialize_kv_caches_v0),
63
patch.object(V1EngineCore, "_initialize_kv_caches",
58
- _initialize_kv_caches_v1)):
64
+ _initialize_kv_caches_v1), monkeypatch.context() as m):
65
+ if model_info.v0_only:
66
+ m.setenv("VLLM_USE_V1", "0")
67
LLM(
68
model_info.default,
69
tokenizer=model_info.tokenizer,
@@ -65,6 +73,7 @@ def test_can_initialize(model_arch):
73
"num_speculative_tokens": 1,
74
} if model_info.speculative_model else None,
75
trust_remote_code=model_info.trust_remote_code,
76
+ max_model_len=model_info.max_model_len,
77
load_format="dummy",
78
hf_overrides=hf_overrides,
79
)
vllm/model_executor/models/grok1.py
+7
-23
@@ -28,7 +28,7 @@ import torch
28
import torch.nn.functional as F
29
from torch import nn
30
31
-from vllm.attention import Attention, AttentionMetadata
31
+from vllm.attention import Attention
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, VllmConfig
34
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -182,25 +182,20 @@ class Grok1Attention(nn.Module):
182
quant_config=quant_config,
183
logits_soft_cap=attn_logits_soft_cap,
184
prefix=f"{prefix}.attn")
185
+ self.attn_multiplier = getattr(self.config, "attn_output_multiplier",
186
+ 1.0) if self.config else 1.0
187
188
def forward(
189
self,
190
positions: torch.Tensor,
191
hidden_states: torch.Tensor,
190
- kv_cache: torch.Tensor,
191
- attn_metadata: AttentionMetadata,
192
) -> torch.Tensor:
193
qkv, _ = self.qkv_proj(hidden_states)
194
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
195
q, k = self.rotary_emb(positions, q, k)
196
- attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
196
+ attn_output = self.attn(q, k, v)
197
output, _ = self.o_proj(attn_output)
198
-
199
- # Apply attention output multiplier if specified in config
200
- attn_multiplier = getattr(self.config, "attn_output_multiplier",
201
- None) if self.config else None
202
- if attn_multiplier is not None:
203
- output = output * attn_multiplier
198
+ output *= self.attn_multiplier
199
return output
200
201
@@ -261,8 +256,6 @@ class Grok1DecoderLayer(nn.Module):
256
self,
257
positions: torch.Tensor,
258
hidden_states: torch.Tensor,
264
- kv_cache: torch.Tensor,
265
- attn_metadata: AttentionMetadata,
259
residual: Optional[torch.Tensor],
260
) -> tuple[torch.Tensor, torch.Tensor]:
261
# Self Attention
@@ -276,8 +269,6 @@ class Grok1DecoderLayer(nn.Module):
269
hidden_states = self.attn(
270
positions=positions,
271
hidden_states=hidden_states,
279
- kv_cache=kv_cache,
280
- attn_metadata=attn_metadata,
272
)
273
274
# Post attention normalization
@@ -341,8 +332,6 @@ class Grok1Model(nn.Module):
332
self,
333
input_ids: torch.Tensor,
334
positions: torch.Tensor,
344
- kv_caches: list[torch.Tensor],
345
- attn_metadata: AttentionMetadata,
335
intermediate_tensors: Optional[IntermediateTensors],
336
inputs_embeds: Optional[torch.Tensor] = None,
337
) -> Union[torch.Tensor, IntermediateTensors]:
@@ -359,9 +348,7 @@ class Grok1Model(nn.Module):
348
349
for i in range(self.start_layer, self.end_layer):
350
layer = self.layers[i]
362
- hidden_states, residual = layer(positions, hidden_states,
363
- kv_caches[i - self.start_layer],
364
- attn_metadata, residual)
351
+ hidden_states, residual = layer(positions, hidden_states, residual)
352
353
if not get_pp_group().is_last_rank:
354
return IntermediateTensors({
@@ -529,13 +516,10 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
516
self,
517
input_ids: torch.Tensor,
518
positions: torch.Tensor,
532
- kv_caches: list[torch.Tensor],
533
- attn_metadata: AttentionMetadata,
519
intermediate_tensors: Optional[IntermediateTensors] = None,
520
inputs_embeds: Optional[torch.Tensor] = None,
521
) -> Union[torch.Tensor, IntermediateTensors]:
537
- hidden_states = self.model(input_ids, positions, kv_caches,
538
- attn_metadata, intermediate_tensors,
522
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
523
inputs_embeds)
524
return hidden_states
525
vllm/utils.py
+9
-6
@@ -2794,14 +2794,17 @@ def cprofile(save_file: Optional[str] = None, enabled: bool = True):
2794
2795
# Only relevant for models using ALiBi (e.g, MPT)
2796
def check_use_alibi(model_config: ModelConfig) -> bool:
2797
- return (getattr(model_config.hf_text_config, "alibi", False) # Falcon
2797
+ cfg = model_config.hf_text_config
2798
+ return (getattr(cfg, "alibi", False) # Falcon
2799
or ("BloomForCausalLM" in getattr(model_config.hf_config,
2800
"architectures", [])) # Bloom
2800
- or getattr(model_config.hf_text_config, "position_encoding_type",
2801
- "") == "alibi" # codellm_1b_alibi
2802
- or
2803
- (hasattr(model_config.hf_text_config, "attn_config") # MPT
2804
- and model_config.hf_text_config.attn_config.get("alibi", False)))
2801
+ or getattr(cfg, "position_encoding_type", "") ==
2802
+ "alibi" # codellm_1b_alibi
2803
+ or (hasattr(cfg, "attn_config") # MPT
2804
+ and ((isinstance(cfg.attn_config, dict)
2805
+ and cfg.attn_config.get("alibi", False)) or
2806
+ (not isinstance(cfg.attn_config, dict)
2807
+ and getattr(cfg.attn_config, "alibi", False)))))
2808
2809
2810
def sha256(input) -> int: