Skip to content

Fix VLM inference engine #2418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meatybobby merged 5 commits into from
Feb 20, 2026
Merged

Fix VLM inference engine #2418

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3cad2ec
Fix VLM inference engine
meatybobby Feb 17, 2026
aa1b9c3
format
meatybobby Feb 17, 2026
5d0bf0e
Fix test
meatybobby Feb 18, 2026
c1995c6
Remove batch size in generate
meatybobby Feb 20, 2026
d48c6d7
Add max batch size in init
meatybobby Feb 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions examples/inference/vlm/vlm_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from typing import Optional

import torch
from megatron.core.inference.common_inference_params import CommonInferenceParams
from megatron.core.inference.sampling_params import SamplingParams
from qwen_vl_utils import process_vision_info

from megatron.bridge.inference.vlm.base import generate, setup_model_and_tokenizer
Expand Down Expand Up @@ -94,7 +94,7 @@ def main(args) -> None:
text, image_inputs, video_inputs = process_image_inputs(processor, args.image_path, prompt)

# Setup inference parameters
inference_params = CommonInferenceParams(
inference_params = SamplingParams(
temperature=args.temperature,
top_p=args.top_p,
top_k=args.top_k,
Expand All @@ -111,7 +111,7 @@ def main(args) -> None:
processor=processor,
max_batch_size=1,
random_seed=0,
inference_params=inference_params,
sampling_params=inference_params,
)

# Print results
Expand Down
10 changes: 9 additions & 1 deletion src/megatron/bridge/inference/vlm/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import torch
import torch.distributed
from megatron.core.inference.contexts import StaticInferenceContext
from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
AbstractModelInferenceWrapper,
)
Expand Down Expand Up @@ -120,6 +121,7 @@ def _expose_decoder_from_language_model(model):
if hasattr(current, "language_model"):
language_model = current.language_model
current.decoder = language_model.decoder
current.vocab_size = language_model.vocab_size


def setup_inference_wrapper(
Expand All @@ -128,6 +130,7 @@ def setup_inference_wrapper(
params_dtype: torch.dtype = torch.bfloat16,
inference_batch_times_seqlen_threshold: int = 1000,
inference_max_seq_length: int = 8192,
inference_max_batch_size: int = 4,

@coderabbitai coderabbitai Bot Feb 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

inference_max_batch_size is not surfaced through setup_model_and_tokenizer.

setup_inference_wrapper gains the new inference_max_batch_size parameter (default 4), but setup_model_and_tokenizer — the public entry point — never accepts or forwards it (lines 104–110). Any caller who needs generate(max_batch_size=N) with N > 4 will hit a StaticInferenceContext KV-cache overflow at runtime, with no way to prevent it through the public API.

🔧 Suggested fix: surface the parameter in the public function 
 def setup_model_and_tokenizer(
     megatron_model_path: str,
     tp: int = 1,
     pp: int = 1,
     params_dtype: torch.dtype = torch.bfloat16,
     inference_batch_times_seqlen_threshold: int = 1000,
     inference_max_seq_length: int = 8192,
+    inference_max_batch_size: int = 4,
 ):
     ...
     inference_wrapped_model = setup_inference_wrapper(
         model[0],
         processor.tokenizer,
         params_dtype=torch.bfloat16,
         inference_batch_times_seqlen_threshold=1000,
         inference_max_seq_length=inference_max_seq_length,
+        inference_max_batch_size=inference_max_batch_size,
     )
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/inference/vlm/base.py` at line 133, The public function
setup_model_and_tokenizer currently does not accept or forward the new
inference_max_batch_size parameter, causing callers to be unable to set
generate(max_batch_size=N) and risking StaticInferenceContext KV-cache
overflows; update setup_model_and_tokenizer to add an inference_max_batch_size:
int = 4 parameter and pass it through to setup_inference_wrapper (which already
accepts inference_max_batch_size) so downstream calls to generate() can use the
configured max_batch_size and avoid KV-cache overflow.

@athitten athitten Feb 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@meatybobby could you please check if this is relevant ?

@coderabbitai coderabbitai Bot Feb 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!

@meatybobby meatybobby Feb 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MCore still has backward compatibility for this, but I just removed it to avoid confusing.

):
"""Set up inference wrapper for the model"""
config = model.config
Expand All @@ -148,7 +151,12 @@ def setup_inference_wrapper(
else:
raise ValueError(f"Unknown model config: {config}")

inference_wrapped_model = wrapper_cls(mcore_model)
inference_wrapped_model = wrapper_cls(mcore_model,
inference_context=StaticInferenceContext(
max_batch_size=inference_max_batch_size,
max_sequence_length=inference_max_seq_length,
),
)

return inference_wrapped_model

Expand Down
4 changes: 2 additions & 2 deletions src/megatron/bridge/inference/vlm/qwenvl_inference_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ class QwenVLInferenceWrapper(AbstractModelInferenceWrapper):
model (Qwen2VLModel): The Qwen2VL model
"""

def __init__(self, model):
super().__init__(model)
def __init__(self, model, inference_context=None):
super().__init__(model, inference_context=inference_context)
Comment thread
meatybobby marked this conversation as resolved.

def prep_inference_input(
self,
Expand Down
Loading