Stabilize local vLLM DeepGEMM warmup startup

nemo_retriever/src/nemo_retriever/models/inference/vllm.py

@@ -16,12 +16,23 @@
 from __future__ import annotations
 
 import logging
+import os
 from typing import Any, List, Optional
 
 logger = logging.getLogger(__name__)
 
 VLLM_DTYPE = "bfloat16"
 VLLM_ATTENTION_BACKEND = "FLASH_ATTN"
+VLLM_DEEP_GEMM_WARMUP_DEFAULT = "skip"
+
+
+def apply_vllm_startup_defaults() -> None:
+    """Apply conservative vLLM startup defaults without overriding users."""
+
+    # DeepGEMM can still be used by vLLM at runtime. This only skips the
+    # ahead-of-time warmup path, which may fail before local inference starts
+    # when the optional DeepGEMM/CUDA-toolkit stack is not discoverable.
+    os.environ.setdefault("VLLM_DEEP_GEMM_WARMUP", VLLM_DEEP_GEMM_WARMUP_DEFAULT)
 
 
 def create_vllm_llm(
@@ -44,6 +55,7 @@ def create_vllm_llm(
     Uses bfloat16 and FLASH_ATTN backend (fixed for this module).
 
     """
+    apply_vllm_startup_defaults()
     try:
         from vllm import LLM
     except ImportError as e:
@@ -185,4 +197,9 @@ def embed_multimodal_with_vllm_llm(
     return all_embeddings
 
 
-__all__ = ["create_vllm_llm", "embed_with_vllm_llm", "embed_multimodal_with_vllm_llm"]
+__all__ = [
+    "apply_vllm_startup_defaults",
+    "create_vllm_llm",
+    "embed_with_vllm_llm",
+    "embed_multimodal_with_vllm_llm",
+]

nemo_retriever/src/nemo_retriever/models/local/nemotron_parse_v1_2.py

@@ -81,6 +81,9 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults
+
+        apply_vllm_startup_defaults()
         try:
             from vllm import LLM, SamplingParams  # noqa: F401
         except ImportError as e:

nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py

@@ -74,6 +74,9 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults
+
+        apply_vllm_startup_defaults()
         try:
             from vllm import LLM
         except ImportError as e:

nemo_retriever/src/nemo_retriever/models/local/nemotron_vlm_captioner.py

@@ -84,6 +84,9 @@ def __init__(
         profile = get_caption_model_profile(model_path, target="local")
         model_path = profile.local_model_id
 
+        from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults
+
+        apply_vllm_startup_defaults()
         try:
             from vllm import LLM, SamplingParams  # noqa: F401
         except ImportError as e:

nemo_retriever/tests/test_vllm_embed.py

@@ -6,6 +6,7 @@
 
 import base64
 import io
+import os
 import sys
 from types import ModuleType, SimpleNamespace
 from unittest.mock import MagicMock, patch
@@ -14,7 +15,11 @@
 
 torch = pytest.importorskip("torch")
 
-from nemo_retriever.models.inference.vllm import embed_multimodal_with_vllm_llm, embed_with_vllm_llm
+from nemo_retriever.models.inference.vllm import (
+    apply_vllm_startup_defaults,
+    embed_multimodal_with_vllm_llm,
+    embed_with_vllm_llm,
+)
 from nemo_retriever.models.local.llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2Embedder
 
 
@@ -192,6 +197,32 @@ def test_limit_mm_per_prompt_forwarded_when_provided(self):
         _, kwargs = mock_llm_cls.call_args
         assert kwargs.get("limit_mm_per_prompt") == {"image": 1}
 
+    def test_applies_vllm_startup_defaults_before_constructing_llm(self, monkeypatch):
+        monkeypatch.delenv("VLLM_DEEP_GEMM_WARMUP", raising=False)
+        with patch("vllm.LLM") as mock_llm_cls:
+            mock_llm_cls.return_value = MagicMock()
+            from nemo_retriever.models.inference.vllm import create_vllm_llm
+
+            create_vllm_llm("some-model")
+
+        assert os.environ["VLLM_DEEP_GEMM_WARMUP"] == "skip"
+
+
+class TestVllmStartupDefaults:
+    def test_deep_gemm_warmup_defaults_to_skip(self, monkeypatch):
+        monkeypatch.delenv("VLLM_DEEP_GEMM_WARMUP", raising=False)
+
+        apply_vllm_startup_defaults()
+
+        assert os.environ["VLLM_DEEP_GEMM_WARMUP"] == "skip"
+
+    def test_deep_gemm_warmup_respects_user_override(self, monkeypatch):
+        monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "full")
+
+        apply_vllm_startup_defaults()
+
+        assert os.environ["VLLM_DEEP_GEMM_WARMUP"] == "full"
+
 
 class TestVLLMEmbedderImages:
     def setup_method(self):