Add VLM support to the dynamic-batching inference server - #6260
Open
RPrenger wants to merge 36 commits into
Open
Add VLM support to the dynamic-batching inference server#6260RPrenger wants to merge 36 commits into
RPrenger wants to merge 36 commits into
Conversation
svcnvidia-nemo-ci
marked this pull request as draft
August 4, 2026 20:24
Contributor
|
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:
See the contribution guide for more details. |
RPrenger
force-pushed
the
vlm-inference
branch
3 times, most recently
from
August 11, 2026 04:00
5d668b2 to
83afeb8
Compare
RPrenger
marked this pull request as ready for review
August 11, 2026 04:03
Contributor
Author
|
/ok to test 83afeb8 |
RPrenger
force-pushed
the
vlm-inference
branch
from
August 11, 2026 04:27
83afeb8 to
9e36a7c
Compare
Contributor
Author
|
/ok to test 9e36a7c |
Signed-off-by: Cory Ye <cye@nvidia.com>
Signed-off-by: Cory Ye <cye@nvidia.com>
Pure formatting pass over the files this PR touches, so the CI linting job stays green and Claude review doesn't get distracted by whitespace nits. No semantic changes. Signed-off-by: rprenger <rprenger@nvidia.com>
_preprocess_data now returns final_position_ids=None on the dynamic-
inference path (decoder_input is supplied directly and the LM ignores
position_ids), but the >language_max_sequence_length truncation branch
still subscripted it unconditionally. When a combined VLM sequence
exceeded the language max, that raised
TypeError: 'NoneType' object is not subscriptable
inside _preprocess_data instead of the actual bookkeeping.
Skip the slice when final_position_ids is None; the None value is
propagated to the caller and matches how the language model consumes it.
Addresses PR NVIDIA#6260 review comment on llava_model.py:902.
Signed-off-by: rprenger <rprenger@nvidia.com>
…sition_ids
_preprocess_data returns combined_position_ids as a plain arange over the
combined [text + image-expanded] sequence, which does not encode the per-
image token layout the surrounding code computes. Forwarding that to the
language model as position_ids was inert for plain RoPE with decoder_input
supplied, but for position_embedding_type in {learned_absolute, mrope} the
LM reads position_ids directly and would consume the wrong positions
silently.
Pass position_ids=None instead. The LM either computes positions itself
from decoder_input (RoPE path) or fails clearly (learned_absolute, mrope)
rather than silently degrading. Deriving an image-aware combined position
sequence is a follow-up; a TODO in _preprocess_data at final_position_ids
tracks it.
Addresses PR NVIDIA#6260 review comment on llava_model.py:1336.
Signed-off-by: rprenger <rprenger@nvidia.com>
pixel_shuffle used to accept optional h/w kwargs; when both were passed
the function short-circuited to a fast reshape path for non-square patch
grids. Dropping those kwargs broke:
* examples/mimo tests (test_radio_model.py) that call pixel_shuffle
with h/w for the non-square case,
* the RADIO dynamic-resolution code path in llava_model that supplies
per-image h/w.
Restore h/w as optional kwargs and the fast-path reshape they enable.
When both are None the function keeps the square derivation from
x.shape[1] we had, so square-tile callers see no change.
Addresses PR NVIDIA#6260 review comment on llava_model.py:1471.
Signed-off-by: rprenger <rprenger@nvidia.com>
… silently ignoring The forward signature still accepts sound_clips / sound_length / sound_timestamps / num_sound_clips / num_frames to keep source compatibility with upstream LLaVAModel, but this PR reduced the multimodal tree to the LLaVA + vision path this inference engine uses, so the audio branch of _preprocess_data and the temporal (>1 frame) video handling are gone. Previously the kwargs were accepted and silently ignored, giving a user with an audio- or video-capable checkpoint a plausibly-correct completion for the wrong modality. Raise NotImplementedError with a message naming the missing paths when any of these kwargs is set to a non-default value. num_frames == 1 still passes because the training tree uses it as a signal for the static-image path and dropping it would break existing image-only callers. Restoring the full audio/video path is a follow-up. Addresses PR NVIDIA#6260 review comment on llava_model.py:1091. Signed-off-by: rprenger <rprenger@nvidia.com>
The image-index guard in _forward_dynamic was written as ``assert max_idx < image_embeddings_flat.shape[0]``, which python -O drops. Under an optimized runtime a real mismatch between ``expand_image_tokens`` and ``_forward_vision_encoder`` (class-token handling divergence, pixel-shuffle rounding, etc.) then silently indexes out of bounds and either raises a confusing kernel-side error or writes adjacent memory. Replace the assert with an explicit ``raise RuntimeError`` so the check survives ``-O`` and continues to point at the two producers that disagreed. Addresses PR NVIDIA#6260 review comment on vlm_inference_wrapper.py:405. Signed-off-by: rprenger <rprenger@nvidia.com>
Two related issues in the RADIO branch of LLaVAModel.__init__: * The constructor honored only the ``class_token_len`` constructor argument as an override for the per-radio-variant defaults, and ignored ``vision_transformer_config.class_token_len`` -- which is the field the encoder registry writes for RADIO checkpoints whose effective class-token length disagrees with the hardcoded default (radio: registry says 10, code assumed 8). * After the RADIO branch resolved a ``radio_class_token_len``, ``class_token_len`` -- the outer variable that ``self._class_token_len`` is stored from at the bottom of the constructor -- was left at the placeholder value 1 set for the CLIP-family fallback path. Consumers of ``self._class_token_len`` therefore saw ``1`` for every RADIO model instead of the actual value fed to the ViT. Extend the override chain to consult ``vision_transformer_config.class_token_len`` between the constructor arg and the hardcoded default, and write the resolved value back to ``class_token_len`` so ``self._class_token_len`` matches the ViT's class-token width. Addresses PR NVIDIA#6260 review comments on llava_model.py:314 and :497. Signed-off-by: rprenger <rprenger@nvidia.com>
The clamped-cu_seqlens rebuild path computed ``(cu[1:] - cu[:-1]).max().item()`` twice, once for ``max_seqlen_q`` and once for ``max_seqlen_kv``. Each invocation triggers a device-to-host sync, and the subtract+max is duplicated. Both fields carry the same value on this path. Compute the max once, sync a single scalar, and reuse it. Net effect on the training forward path: two D2H syncs collapse to one, and one element-wise pass drops. Addresses PR NVIDIA#6260 review comment on llava_model.py:1323. Signed-off-by: rprenger <rprenger@nvidia.com>
… MegatronAsyncLLM The base ``_MegatronLLMBase`` already defaults ``inference_wrapper_cls`` to ``GPTInferenceWrapper``. The subclass wrappers ``MegatronLLM`` and ``MegatronAsyncLLM`` declared the parameter as ``Optional[...] = None`` and passed the ``None`` through to the base, which overrode the base default. The base then called ``inference_wrapper_cls(model, context)`` with ``None`` and raised ``TypeError: 'NoneType' object is not callable`` -- a hard break for every existing caller that constructed ``MegatronLLM`` / ``MegatronAsyncLLM`` without explicitly passing a wrapper class. Give both subclasses the same ``GPTInferenceWrapper`` default the base already uses so the value passed through is a valid class and existing callers keep working. VLM callers still pass ``VLMInferenceWrapper`` explicitly. Addresses PR NVIDIA#6260 review comment on llm.py:45. Signed-off-by: rprenger <rprenger@nvidia.com>
The base ``_generate_impl`` gained ``multi_modal_data_list`` as a **required positional** parameter with no default, breaking every pre-VLM caller and subclass override (Sync ``MegatronLLM.generate``, Async ``MegatronAsyncLLM.generate``, plus any external subclass) that was already forwarding ``prompts`` and ``sp`` alone. Default the parameter to ``None`` and interpret ``None`` as "no multi-modal data attached to any prompt". Callers with images still pass a list the same length as ``prompts``; text-only callers keep working with no code change. Addresses PR NVIDIA#6260 review comment on _llm_base.py:481. Signed-off-by: rprenger <rprenger@nvidia.com>
…ll branch The decode short-circuit returned a 1D ``(padded_active_token_count,)`` tensor, while the prefill branch returned a 2D ``[1, padded_active_token_count]`` (from ``mask.unsqueeze(0)`` at the tail). Callers advanced-index a batch-first ``[b, seq, h]`` embedding tensor with the mask, and the 2D form is the one that broadcasts correctly across the batch axis. On a decode step that reached the ``image_token_mask is not None`` branch, the 1D mask would either error on shape mismatch or silently select along dim 0. Unsqueeze the decode return so both branches produce ``[1, N]``. Addresses PR NVIDIA#6260 review comment on dynamic_context.py:4746. Signed-off-by: rprenger <rprenger@nvidia.com>
``current_image_embeddings`` walked the active request slice, ran a ``.tolist()`` (D2H sync), gathered per-request embeddings, and returned a concatenated tensor even on pure decode steps -- where the returned tensor is not consumed (the decode wrapper path uses only the mask existence check). The sync fired every decode step, defeating the overlap the dynamic engine is built around, especially at high concurrency. Add the same ``is_decode_only()`` short-circuit already used by ``current_image_token_mask`` so decode steps return ``None`` without paying the sync or the cat. Addresses PR NVIDIA#6260 review comment on dynamic_context.py:4838. Signed-off-by: rprenger <rprenger@nvidia.com>
``TextGenerationController._dynamic_step_forward_logits`` called ``current_image_token_mask`` and ``current_image_embeddings`` on every step regardless of whether any request had image data attached. Both helpers already short-circuit on empty per-request dicts, but the lookup + method-call overhead still fires per decode step on the critical path. Expose a ``has_vlm_data`` property on ``DynamicInferenceContext`` (true iff any active request has attached image data) and gate the helper calls on it. Text-only workloads skip both calls entirely; VLM workloads see no change. Addresses PR NVIDIA#6260 review comment on text_generation_controller.py:826. Signed-off-by: rprenger <rprenger@nvidia.com>
…_tokens The dynamic-resolution branch of ``expand_image_tokens`` looped over ``imgs_sizes`` in Python and called ``imgs_sizes[i][0].item(), imgs_sizes[i][1].item()`` per iteration, firing two D2H syncs per image on the admission path. For an N-image request that's 2N blocking syncs before the vision encoder can even run. Do one ``imgs_sizes.tolist()`` up front and iterate over the resulting Python nested list. Same math, single sync. Addresses PR NVIDIA#6260 review comment on vlm_inference_wrapper.py:179. Signed-off-by: rprenger <rprenger@nvidia.com>
_build_vlm_request always computed ``int(num_tiles.sum().item())`` for every image-bearing admission, even when the request came in on the dynamic-resolution path (imgs + imgs_sizes) where ``num_tiles`` was either absent or unused: that path derives its per-image embedding count from ``imgs_sizes`` inside ``expand_image_tokens``. The sync was therefore paying a D2H roundtrip for a value that was never read on the dynamic-res path. Guard the ``num_tiles.sum().item()`` behind the static-tiling branch so dynamic-res admissions no longer stall on it; static-tiling admissions keep the exact same behavior. Companion to the vlm_inference_wrapper .tolist() batching. Combined they remove two of the three per-admission syncs Claude flagged; the remaining ``tokens.tolist()`` (needed to feed ``expand_image_tokens`` which takes ``List[List[int]]``) would need a wider on-device expansion refactor to eliminate -- tracking in a follow-up. Addresses PR NVIDIA#6260 review comment on dynamic_engine.py:1385. Signed-off-by: rprenger <rprenger@nvidia.com>
The SSRF allowlist check ran against the URL's original hostname, but ``urlopen`` then followed 3xx responses to wherever they pointed -- including private addresses -- with no further check. A public URL that returned ``302 http://169.254.169.254/...`` slipped straight past the guard. Install a small ``HTTPRedirectHandler`` subclass that raises on every 3xx and route the fetch through a dedicated opener. Everything else on the fetch path stays as it was (allowlist, timeout, size cap, user-agent). DNS rebinding between the allowlist check and the socket connect is a known residual attack surface that would need socket-level control to fully close; deployments that expose this endpoint to untrusted networks should also run it behind an egress policy. Addresses PR NVIDIA#6260 review comment on chat_completions.py:274. Signed-off-by: rprenger <rprenger@nvidia.com>
Every admission for a request with n > 1 completions independently re-preprocesses ``image_bytes_list`` and re-runs the vision encoder, so the same images are processed n times when the client only asked for n text completions of one prompt. The proper fix is to preprocess (and optionally encode) the images once at admission time and share the result across the n requests, which needs the HTTP layer to see the engine's ``ImageProcessingConfig`` (not currently threaded up here) and for embeddings to be shippable on the wire. That plumbing is a larger, orthogonal change and not scoped to this PR. Adding a TODO here so a follow-up has a landing point. Addresses PR NVIDIA#6260 review comment on chat_completions.py:784. Signed-off-by: rprenger <rprenger@nvidia.com>
…mples/ import
``preprocess_image_bytes_tiled`` inverted the dependency direction by
importing ``examples.multimodal.image_processing.ImageTransform`` from
inside ``megatron/core``. No supported encoder in the registry is on
the tiling path (``use_tiling`` is unset everywhere in the registry),
and no in-tree caller submits raw bytes with tiling enabled today, so
we drop the helper rather than move its ~170-line dependency into
core.
Wire clients that need static tiling should preprocess bytes into a
tensor dict themselves and submit
``multi_modal_data['image'] = {'imgs': ..., 'num_tiles': ...,
'num_img_embeddings_per_tile': ...}`` on the wire; the engine already
accepts that shape without any examples/ dependency. The
dynamic-resolution path used by every current encoder is unaffected.
Addresses PR NVIDIA#6260 review comment on image_preprocessing.py:229.
Signed-off-by: rprenger <rprenger@nvidia.com>
``vlm_dynamic_inference`` appended ``examples/multimodal/`` to ``sys.path`` at import time so a later ``from model import model_provider`` inside ``get_model`` would resolve. A library module in ``megatron/core`` mutating a global on import is a bad pattern -- it silently changes lookup order for every module in the process, whether or not the caller ever exercises the VLM path. The caller that actually invokes ``get_model`` (``tools/run_dynamic_text_generation_server.py``) already sets both the repo root and ``examples/multimodal/`` on ``sys.path``, so the mutation here was redundant. Drop it and document in a NOTE that ``get_model``'s bare ``from model import`` requires the caller to have set that up (which every in-tree entry point already does). Also removes the now-unused ``os`` and ``sys`` imports from this module. Addresses PR NVIDIA#6260 review comment on vlm_dynamic_inference.py:33. Signed-off-by: rprenger <rprenger@nvidia.com>
…spec_te The old comment claimed the moe_layer branch would fail loudly when ``config`` didn't carry the required MoE fields, but no assert was present. Passing a config with ``num_moe_experts=None`` (either accidentally, or because the caller only wanted a non-MoE hybrid) fell through to ``get_moe_module_spec(num_experts=None, ...)`` and either built a spec whose moe_layer trip would fail deep inside the MoE spec, or produced a checkpoint-mismatched architecture silently. Add the assert the comment already promised, and replace the free-form comment with a proper docstring so callers can see the ``config`` contract without reading source. Signature is unchanged: ``config`` and ``padding`` stay in the same positions, so no callers move. Addresses PR NVIDIA#6260 review comment on layer_specs.py:131. Signed-off-by: rprenger <rprenger@nvidia.com>
The previous order was:
1. compute pos = arange over patch tokens only
2. shift pos by class_token_len (to "make room" for CLS)
3. add position_embeddings to the patch-only tensor
4. prepend CLS
That produced two off-by-class_token_len effects:
* The CLS slot never received a position embedding (it was concatenated
onto x after the +position_embeddings step and started life as pure
class_token content, no positional signal).
* Patch tokens ended up at positions class_token_len..class_token_len+N,
which is what a patch would want AFTER the CLS was in place — but the
prepending happened later, so the effective positions the transformer
saw for patches drifted relative to the pre-shift arange.
Prepend CLS first, then compute ``pos = arange(x.shape[1])`` over the
CLS + patch tensor and add ``position_embeddings(pos)``. CLS now gets
positions [0, class_token_len) and patches sit at [class_token_len, N +
class_token_len). RoPE branch is unaffected — rope is passed to the
transformer as a separate tensor covering patch tokens only.
Addresses PR NVIDIA#6260 review comment on vit_model.py:320.
Signed-off-by: rprenger <rprenger@nvidia.com>
The fast path (h_patches / w_patches match the stored resolution) returned ``self.weight.reshape(...)`` without moving the weight tensor to the caller-supplied ``device``. The slow path (bicubic interp) does ``w = self.weight.to(device=device)`` first, so downstream operations land on the right device. If a caller passed a ``device`` that differed from where the module's weight lived (fresh instantiation before ``.to(device)``, or an explicit device override at call time), the fast path returned tensor would be on the wrong device and the addition in ``ViTModel.forward`` would either error on device-mismatch or trigger an implicit copy. Hoist the ``.to(device=device)`` above the fast-path branch so both paths behave identically. Addresses PR NVIDIA#6260 review comment on vit_model.py:801. Signed-off-by: rprenger <rprenger@nvidia.com>
…urn shape Newer ``transformers`` releases default ``apply_chat_template`` to ``return_dict=True``, which returns a ``BatchEncoding`` instead of the list/tensor the ``[0]`` subscript below (and the ``len()`` on the per-turn helper) expects. Restore ``return_dict=False`` on both call sites so behavior stays consistent across ``transformers`` versions. Addresses PR NVIDIA#6260 review comment on multimodal_tokenizer.py:272. Signed-off-by: rprenger <rprenger@nvidia.com>
``image_preprocessing.py`` uses ``PIL.Image`` to decode client-supplied image bytes on the VLM inference path, but Pillow was not declared as a ``megatron-core`` install dependency: on a fresh ``pip install megatron-core[dev]`` without Pillow separately available, the VLM inference request path would raise ``ImportError`` on the first image request, not at install time. Add ``Pillow`` next to the other multimedia deps (``av``, energon's audio/video decoders) in the ``dev`` optional-dependency block. This keeps the megatron-core install self-contained for users who install the multimodal path without cloning the repo. ``torchvision`` is also used by the same file but is already declared via ``override-dependencies`` in the same way as ``torch``, treating it as provided by the NGC PyTorch base image; the convention is unchanged. Requires a follow-up ``uv lock`` regeneration (``UV_PYTHON=3.12 uvx uv@0.7.2 lock``) so ``uv.lock`` picks up the Pillow entry. Signed-off-by: rprenger <rprenger@nvidia.com>
…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>
Contributor
Author
|
/ok to test a9464c7 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
What does this PR do?
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
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"
.github/CODEOWNERS.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, theFinal Reviewlabel 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
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.