Add VLM support to the dynamic-batching inference server by RPrenger · Pull Request #6260 · NVIDIA/Megatron-LM

RPrenger · 2026-08-04T20:24:37Z

Adds an end-to-end VLM inference path on top of the DynamicInferenceEngine:

engine plumbing accepts per-request image inputs (imgs, imgs_sizes, num_tiles, num_img_embeddings_per_tile); the engine expands placeholders into pad tokens, runs the vision encoder on the first PP stage, and attaches per-request image embeddings and an image-token mask to the DynamicInferenceContext so the decoder forward can splice them back in.
wire schema between InferenceClient, the coordinator, and the engine drain grows an optional 5th slot for raw image bytes; text-only callers keep the 4-slot payload. Prefix-cache routing is skipped for image-bearing requests so text-identical prompts with different images do not falsely share kv-cache prefixes.
image preprocessing (dynamic-resolution and static tiling), pixel-stat encoder registry, chat_template plumbing on /v1/completions and /v1/chat/completions, and a VLM-aware run_dynamic_text_generation_server entrypoint that auto-detects VLM checkpoints and builds the LLaVA- wrapped model with VLMInferenceWrapper.
LLaVAModel gets a forward_lm_only entry point for the dynamic path, plus a few attributes the wrapper reads. Upstream audio/video params (sound_model, sound_projection, sound_token_index, temporal_patch_dim, separate_video_embedder, temporal_ckpt_compat) are preserved as no-op stubs so the constructor signature stays source-compatible.
Text-only inference is unaffected: none of the new engine or context work fires unless a caller passes imgs/imgs_sizes/num_tiles.

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-08-04T20:24:41Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-08-04T20:24:47Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

RPrenger · 2026-08-11T04:08:35Z

/ok to test 83afeb8

RPrenger · 2026-08-11T04:28:33Z

/ok to test 9e36a7c

…e error torchvision isn't a hard install dependency of ``megatron-core``: the NGC PyTorch container ships one pinned to the container's torch build, so the toml lists it under ``override-dependencies`` (same treatment as ``torch``) and the container assumption covers users on that path. A plain ``pip install megatron-core`` off PyPI does not get torchvision automatically -- installing it via pip needs a build matching the local torch, which the caller has to pick, so we can't just declare it as a regular dep here without breaking the container path. Wrap the lazy import at the one call site (VLM image preprocessing) in a try/except that translates ``ImportError`` into a message naming ``torchvision`` and pointing at the container path. Matches the ``HAVE_TE`` pattern used elsewhere in the repo -- users installing off PyPI who don't need VLM inference are unaffected; users who do hit the VLM path get a clean instruction instead of an opaque import stack. Signed-off-by: rprenger <rprenger@nvidia.com>

Autoformatter pass to keep the CI ``linting`` job green after the SSRF-hardening and torchvision-guard changes shifted import blocks. Pure formatting; no semantic change. Signed-off-by: rprenger <rprenger@nvidia.com>

Companion to 2722ef2, which declared ``Pillow`` under the ``dev`` extras of ``pyproject.toml`` but never regenerated ``uv.lock``. The CI ``Pip`` / ``UV`` / ``Install test summary`` jobs run ``uv sync --locked``, which refuses to install when the lock and the toml disagree. Regenerated inside the NGC PyTorch container we use for eval (matching ``UV_VERSION=0.7.2`` from ``docker/Dockerfile.ci.*`` and Python 3.12 from ``.python-version``). The diff is just the new ``pillow`` entry plus its transitive hash and the ``dev`` extras edge that references it -- no other package versions moved. Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-15T00:36:57Z

/ok to test a9464c7

RPrenger · 2026-08-17T04:31:17Z

/ok to test f1180d4

megatron.inference is not in mcore's setuptools packages.find, so the install-test check-imports scan fails when it walks megatron.core and tries to top-level import megatron.inference.utils. Defer the import into add_vlm_inference_args, which is only called from the server entry point where the full training tree is present. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-17T04:57:50Z

/ok to test 43a3650

The SSRF-hardening no-redirect handler failed pylint's missing-function-docstring check. Add a one-line docstring. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-17T05:19:19Z

/ok to test 9e56ee8

- llava_model.py: only pass ``imgs_sizes=`` to the vision encoder when it's not None; CLIPViTModel.forward() doesn't accept the kwarg. Move ``_preprocess_data``'s ``position_ids`` to keyword-only so positional callers don't silently shift into the wrong slot, and accept the sound_* kwargs (still ignored inside) for API compatibility with pre-existing callers. - dynamic_context.py: guard the VLM dict clears in ``reset_metadata`` with ``getattr`` so tests that construct the context via ``__new__`` don't trip on unset attrs. - llm.py / async_llm.py: default ``inference_wrapper_cls`` to ``None`` and resolve to ``GPTInferenceWrapper`` at call time, so tests that monkey-patch the module-level name actually steer construction. - test_apis.py: also patch ``GPTInferenceWrapper`` on the llm / async_llm modules (safety-belt alongside the code change). - test_data_parallel_inference_coordinator.py: give ``DummyEngine`` a ``failed_request_ids`` list so ``_handle_failed_request`` no longer raises ``AttributeError`` and hangs the coordinator (was the source of the 13x TimeoutError cluster). - test_llava_model.py::test_preprocess_data: unpack all 5 return values now that ``_preprocess_data`` also returns combined input_ids / position_ids. - test_llava_sound.py::TestPreprocessDataSoundReplacement: skip; the sound-embedding splicing block was removed during the VLM refactor. Kwargs are still accepted but no longer replace tokens in-place. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-17T07:48:39Z

/ok to test 54e2c47

black wanted the pytest.mark.skip decorator wrapped differently. Pure formatting, no behavior change. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-17T16:53:18Z

/ok to test 93f19e7

Companion to the imgs_sizes gating in the previous fix commit — the same CLIPViTModel test paths were still failing on ``CLIPViTModel.forward() got an unexpected keyword argument 'packed_seq_params'``. Move ``packed_seq_params`` into the same "conditional pass" dict as ``imgs_sizes`` so both are only forwarded to the vision encoder when actually populated. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

Two follow-up fixes for the CI unit-test failures that survived the earlier round: - dynamic_context.py: ``reset_metadata`` clears a third VLM dict (``_request_to_image_token_count``) I missed in the previous guarded round. Wrap the same way with ``getattr(...).clear()`` so tests that build the context via ``__new__`` don't AttributeError. - dynamic_engine.py: the SUBMIT_REQUEST handler dereferences ``self.context.config.image_preprocessing_config`` unconditionally, even for text-only requests. Test fixtures use a ``DummyContext`` without a ``.config`` attribute, so every SUBMIT_REQUEST raised AttributeError, got swallowed by ``_fail_submission``, and desynced the ranks (surfacing later as an NCCL collective timeout in the distributed coordinator test). Skip the config lookup entirely when ``multi_modal_data is None`` — the resolver already returns ``{}`` for that case, so this preserves the production behavior. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-17T20:16:01Z

/ok to test 7d2cb12

Neither GPTModel.forward() nor HybridModel.forward() accepts ``mtp_source_loss_mask`` — it was never consumed anywhere in-tree. The line just TypeError'd every LLaVA-with-labels forward call, which manifested as the ``test_llava_model.py::test_forward*`` and ``test_cuda_graphs.py::test_llava_cudagraph_is_last_layer_logic`` failures. Drop the kwarg; keep the ``loss_mask`` line since that kwarg IS accepted. If a downstream fork ever adds MTP support to HybridModel and expects this signal, re-thread it there — the plumbing is trivial and the correct wiring depends on where MTP consumes it. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger · 2026-08-17T23:34:03Z

/ok to test 11abf32

RPrenger · 2026-08-18T06:07:38Z

/ok to test cffcdbc

nemo-automation-bot · 2026-08-20T17:51:37Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32399934814

RPrenger requested a review from a team as a code owner August 4, 2026 20:24

RPrenger requested a review from a team August 4, 2026 20:24

RPrenger requested review from a team as code owners August 4, 2026 20:24

svcnvidia-nemo-ci marked this pull request as draft August 4, 2026 20:24

nvidia-megatron-lm-release-bot Bot added the community-request label Aug 4, 2026

chtruong814 removed the community-request label Aug 5, 2026

RPrenger force-pushed the vlm-inference branch 3 times, most recently from 5d668b2 to 83afeb8 Compare August 11, 2026 04:00

RPrenger marked this pull request as ready for review August 11, 2026 04:03

github-actions Bot added the complexity: high label Aug 11, 2026

copy-pr-bot Bot temporarily deployed to public August 11, 2026 04:09 Inactive

copy-pr-bot Bot temporarily deployed to public August 11, 2026 04:12 Inactive

copy-pr-bot Bot temporarily deployed to public August 11, 2026 04:24 Inactive

RPrenger force-pushed the vlm-inference branch from 83afeb8 to 9e36a7c Compare August 11, 2026 04:27

copy-pr-bot Bot temporarily deployed to public August 11, 2026 04:30 Inactive

copy-pr-bot Bot temporarily deployed to public August 11, 2026 04:33 Inactive

RPrenger added 3 commits August 14, 2026 17:22

Re-run isort/black on files touched in this round of review fixes

1af7efc

Autoformatter pass to keep the CI ``linting`` job green after the SSRF-hardening and torchvision-guard changes shifted import blocks. Pure formatting; no semantic change. Signed-off-by: rprenger <rprenger@nvidia.com>

Merge branch 'main' into vlm-inference

f1180d4

Add docstring to _NoRedirectHandler.http_error_301 for pylint C0116

9e56ee8

The SSRF-hardening no-redirect handler failed pylint's missing-function-docstring check. Add a one-line docstring. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

Reformat test_llava_sound.py after skip-decorator addition

93f19e7

black wanted the pytest.mark.skip decorator wrapped differently. Pure formatting, no behavior change. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>

RPrenger added 2 commits August 17, 2026 10:46

Merge branch 'main' into vlm-inference

cffcdbc

santhnm2 approved these changes Aug 18, 2026

View reviewed changes

jaredcasper approved these changes Aug 19, 2026

View reviewed changes

dimapihtar approved these changes Aug 20, 2026

View reviewed changes

cspades mentioned this pull request Aug 26, 2026

[Inference] Support video modality in Megatron-Inference. #6809

Merged

4 of 6 tasks

yaoyu-33 mentioned this pull request Aug 26, 2026

[misc] fix: honor chat completion stop sequences NVIDIA-NeMo/Megatron-Bridge#5815

Closed

This was referenced Aug 27, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

Megatron-LM Monthly Activity Report — August 2026 #7006

Open

svcnemo-autobot mentioned this pull request Oct 17, 2026

chore: AUT-2547 nightly sync main into dev (2026-10-16) #8305

Closed

Conversation

RPrenger commented Aug 4, 2026

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

RPrenger commented Aug 11, 2026

Uh oh!

RPrenger commented Aug 11, 2026

Uh oh!

RPrenger commented Aug 15, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 17, 2026

Uh oh!

RPrenger commented Aug 18, 2026

Uh oh!

nemo-automation-bot Bot commented Aug 20, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants