megatron_generate drops VLM vision inputs during no-cache decoding

[cuichenx]

Describe the bug

modelopt.torch.utils.plugins.megatron_generate.megatron_generate drops all VLM vision inputs after generation step 0 even when KV-cache decoding is disabled.

With enable_kv_cache=False, the function recomputes the entire growing token prefix on every step. However, pixel_values, image_grid_thw, and image_sizes are each gated only by step == 0. On steps 1+, the full prefix is therefore recomputed without replacing image-placeholder tokens with vision embeddings. Generation can complete successfully while becoming visually ungrounded.

This also affects callers that request KV cache when model.config.sequence_parallel is true, because megatron_generate automatically falls back to no-cache decoding.

The behavior is present in ModelOpt 0.46.0rc1 and current main at ddca53b4bf99a1a370a8f3c8fe562710695a3448.

Steps/Code to reproduce bug

1. Load a Megatron-Core Qwen3-VL model with a real image and prompt.
2. Call:

megatron_generate(
    model=model,
    input_ids=input_ids,
    pixel_values=pixel_values,
    image_grid_thw=image_grid_thw,
    osl=32,
    enable_kv_cache=False,
)

3. Instrument the model forward calls. Step 0 receives the visual tensors, but all later calls receive the complete growing prefix without visual tensors.

A matched semantic A/B was run with Qwen3-VL-8B-Instruct, an FP8 ModelOpt checkpoint, TP=2, the public Qwen demo beach image, and prompt Describe this image.:

• Stock helper: described the image as an unrelated “visual pun or meme” about “irritable.”
• Same loaded model/checkpoint/input, changing only the three visual-input gates to retain the tensors whenever no inference context exists: correctly described a woman and her dog on a beach at sunset.

The effective A/B condition was:

replay_vision_inputs = inference_context is None or step == 0

and that condition was applied to pixel_values, image_grid_thw, and image_sizes.

This is a semantic correctness failure rather than an exception: the stock call exits successfully and reports successful generation.

Expected behavior

• With KV-cache decoding, visual tensors should be passed during prefill only.
• Without KV-cache decoding, visual tensors should be passed on every full-prefix recomputation.
• A focused regression test should verify the forward-call arguments for both modes.

Who can help?

ModelOpt Megatron/VLM generation owners.

System information

• OS: Linux, x86_64
• GPU: 2x H100 80GB
• Python: 3.12
• ModelOpt: 0.46.0rc1 (982d72eafcba34bb5715703e919c5180c80d1b5a)
• Also confirmed in current main: ddca53b4bf99a1a370a8f3c8fe562710695a3448
• Megatron Bridge: 0.6.0 (7b5057e03224f08b122b132a7f8b1d361b039a7d)
• Megatron Core: 0.19.0 (16ad357ee7973af32916fc1ca39d71065e5f03d4)
• Transformer Engine: 2.17.1
• PyTorch: 2.13.0a0+8145d630e8.nv26.6.54250401
• CUDA: 13.3
• Transformers: 5.12.1

No product source or checkpoint modification was used in the stock reproduction. The corrected arm patched only the generator function in memory to isolate the visual-input gating behavior.