Stabilize local vLLM DeepGEMM warmup startup

[jioffe502]

Summary

• Default NRL local vLLM startup to VLLM_DEEP_GEMM_WARMUP=skip via os.environ.setdefault(...).
• Apply that default before every local NRL vllm.LLM constructor: embedding, VL rerank, captioning, and Nemotron Parse.
• Do not set VLLM_USE_DEEP_GEMM=0 and do not hard-code CUDA_HOME; users can still opt into another vLLM warmup mode explicitly.

Why

JP20 local harness runs have been failing during ingest before any rows are written. The failing artifact points to local embedding/vLLM startup, not caption or rerank:

• results.json: exit_code: 10, failed_phase: ingest
• ingest_plan.json: caption: null, local nvidia/llama-nemotron-embed-1b-v2, local_ingest_embed_backend: "vllm"
• query_plan.json: rerank: false

Original trace:

RuntimeError: DeepGEMM backend is not available or outdated. Please install or update the `deep_gemm` to a newer version to enable FP8 kernels.
RuntimeError: Engine core initialization failed. See root cause above. Failed core proc(s): {}

File ".../llama_nemotron_embed_1b_v2_embedder.py", line 69, in _ensure_loaded
    self._llm = create_vllm_llm(
File ".../models/inference/vllm.py", line 87, in create_vllm_llm
    return LLM(**kwargs)

Downstream symptom:

lancedb_write summary: total=3147 accepted=0 dropped_bad_length=3147 expected_dim=2048

When DeepGEMM warmup can help

This change skips DeepGEMM's ahead-of-time kernel warmup; it does not disable runtime DeepGEMM. With warmup skipped, an eligible kernel is JIT-compiled when first encountered during model execution.

Warmup is useful when all of the following are true:

• The workload runs on Hopper or Blackwell with a compatible DeepGEMM installation.
• The model uses FP8 layers that vLLM actually routes through DeepGEMM.
• First-request latency or stable benchmark latency matters enough to move JIT compilation into startup.

NRL's default local embedding, VL rerank, Nemotron Parse, and caption paths use BF16, so they are not expected to benefit. The realistic NRL candidates are the optional local FP8 caption profiles, especially the FP8 Nemotron 3 Omni MoE model. For those deployments, operators can opt into VLLM_DEEP_GEMM_WARMUP=relax first, or full when every possible GEMM shape must be compiled before accepting work.

Review question

Is skipping optional DeepGEMM warmup the right default for NRL local startup reliability, while letting Hopper/Blackwell performance owners opt into VLLM_DEEP_GEMM_WARMUP=full or another upstream-supported mode?

Validation

• uv run --project nemo_retriever pytest nemo_retriever/tests/test_vllm_embed.py nemo_retriever/tests/test_nemotron_rerank_vl_v2.py nemo_retriever/tests/test_caption_model_profiles.py nemo_retriever/tests/test_nemotron_parse_v1_2.py -q
  • 100 passed, 2 warnings
• Rebased onto current main with Typer 0.27 help support; combined CLI-contract and DeepGEMM-focused suite: 105 passed, 2 warnings
• python -m compileall on changed Python files
• git diff --check
• Live local embedding smoke with CUDA/DeepGEMM env vars unset:

warmup skip
shape (1, 2048)
dtype torch.float32

[jioffe502]

To discuss with Julio and Jeremy

[jdye64]

I have no problem with this.

[greptile-apps]

Greptile Summary

This PR fixes a reproducible JP20 local-harness failure where RuntimeError: DeepGEMM backend is not available or outdated aborted vLLM startup during ingest before any rows were written. The fix adds apply_vllm_startup_defaults() — a single function that calls os.environ.setdefault("VLLM_DEEP_GEMM_WARMUP", "skip") — and invokes it before every local vllm.LLM constructor.

• vllm.py: apply_vllm_startup_defaults() is added and called inside create_vllm_llm, covering the embedding path. It is also exported in __all__ for the three model classes that construct LLM directly.
• NemotronParseV12 / NemotronRerankVLV2VLLM / NemotronVLMCaptioner: each __init__ now calls apply_vllm_startup_defaults() before the from vllm import LLM import, ensuring the env var is visible to vLLM's process-startup code. setdefault semantics preserve any user-supplied value so Hopper/Blackwell operators can opt into full or relax warmup explicitly.
• All four constructor paths now have matching unit tests (including the new test_nemotron_parse_v1_2.py file) that verify the env var is set at the moment LLM() is invoked.

Confidence Score: 5/5

Safe to merge. The change is a targeted, additive env-var default that only fires when the user has not already set VLLM_DEEP_GEMM_WARMUP, and all four local vLLM constructor paths now have matching tests verifying the env var is visible at LLM construction time.

The implementation uses os.environ.setdefault, which is a no-op when the user has already configured the variable — no existing behavior is broken. Each of the four affected code paths (create_vllm_llm for the embedder, and the three direct-constructor classes) is covered by a test that captures the env var at the exact moment LLM() runs. The new test file carries the required SPDX header. There are no logic errors, no hardcoded secrets, no circular imports, and no bare excepts introduced.

No files require special attention.

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/models/inference/vllm.py	Adds apply_vllm_startup_defaults() using os.environ.setdefault (preserves user override), calls it in create_vllm_llm before LLM construction, and exports it via all. Logic is correct and minimal.
nemo_retriever/src/nemo_retriever/models/local/nemotron_parse_v1_2.py	Adds apply_vllm_startup_defaults() call before the vllm.LLM import/construction in NemotronParseV12.init. Placement is correct — env var is set before vLLM's process startup code runs.
nemo_retriever/src/nemo_retriever/models/local/nemotron_rerank_vl_v2.py	Adds apply_vllm_startup_defaults() call before vllm.LLM import/construction in NemotronRerankVLV2VLLM.init. Change is minimal and correctly placed.
nemo_retriever/src/nemo_retriever/models/local/nemotron_vlm_captioner.py	Adds apply_vllm_startup_defaults() call after profile resolution but before vllm.LLM import/construction in NemotronVLMCaptioner.init. Correct placement.
nemo_retriever/tests/test_vllm_embed.py	Adds TestVllmStartupDefaults covering setdefault semantics and user-override preservation, plus a test in TestCreateVllmLlm verifying the env var is set after create_vllm_llm runs.
nemo_retriever/tests/test_nemotron_parse_v1_2.py	New test file with SPDX header. Uses side_effect to verify VLLM_DEEP_GEMM_WARMUP is set at the exact moment LLM() is called during NemotronParseV12 construction.
nemo_retriever/tests/test_nemotron_rerank_vl_v2.py	Adds test_applies_vllm_startup_defaults_before_constructing_llm using side_effect pattern to verify env var ordering at LLM construction time.
nemo_retriever/tests/test_caption_model_profiles.py	Adds deep_gemm_warmup capture to FakeLLM.init and a new test that verifies VLLM_DEEP_GEMM_WARMUP=="skip" is visible at FakeLLM construction time during NemotronVLMCaptioner init.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User
    participant NRLModel as NRL Model __init__
    participant ApplyDefaults as apply_vllm_startup_defaults()
    participant Env as os.environ
    participant VLLM as vllm.LLM()

    Note over User,VLLM: Embedding path (via create_vllm_llm)
    User->>NRLModel: LlamaNemotronEmbed1BV2Embedder.__init__
    NRLModel->>ApplyDefaults: create_vllm_llm() → apply_vllm_startup_defaults()
    ApplyDefaults->>Env: setdefault("VLLM_DEEP_GEMM_WARMUP", "skip")
    ApplyDefaults-->>NRLModel: (returns)
    NRLModel->>VLLM: "LLM(**kwargs) — warmup=skip visible"

    Note over User,VLLM: Direct-constructor paths (new call sites)
    User->>NRLModel: NemotronParseV12 / NemotronRerankVLV2VLLM / NemotronVLMCaptioner __init__
    NRLModel->>ApplyDefaults: apply_vllm_startup_defaults()
    ApplyDefaults->>Env: setdefault("VLLM_DEEP_GEMM_WARMUP", "skip")
    ApplyDefaults-->>NRLModel: (returns)
    NRLModel->>VLLM: "LLM(**kwargs) — warmup=skip visible"

    Note over Env: User override preserved: if VLLM_DEEP_GEMM_WARMUP already set, setdefault is a no-op

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User
    participant NRLModel as NRL Model __init__
    participant ApplyDefaults as apply_vllm_startup_defaults()
    participant Env as os.environ
    participant VLLM as vllm.LLM()

    Note over User,VLLM: Embedding path (via create_vllm_llm)
    User->>NRLModel: LlamaNemotronEmbed1BV2Embedder.__init__
    NRLModel->>ApplyDefaults: create_vllm_llm() → apply_vllm_startup_defaults()
    ApplyDefaults->>Env: setdefault("VLLM_DEEP_GEMM_WARMUP", "skip")
    ApplyDefaults-->>NRLModel: (returns)
    NRLModel->>VLLM: "LLM(**kwargs) — warmup=skip visible"

    Note over User,VLLM: Direct-constructor paths (new call sites)
    User->>NRLModel: NemotronParseV12 / NemotronRerankVLV2VLLM / NemotronVLMCaptioner __init__
    NRLModel->>ApplyDefaults: apply_vllm_startup_defaults()
    ApplyDefaults->>Env: setdefault("VLLM_DEEP_GEMM_WARMUP", "skip")
    ApplyDefaults-->>NRLModel: (returns)
    NRLModel->>VLLM: "LLM(**kwargs) — warmup=skip visible"

    Note over Env: User override preserved: if VLLM_DEEP_GEMM_WARMUP already set, setdefault is a no-op

Loading

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into codex/vllm-deep..." | Re-trigger Greptile}