Skip to content

Fix VLM runtime stability - #30865

Closed
mickqian wants to merge 1 commit into
sgl-project:mainfrom
mickqian:codex/vlm-runtime-stability
Closed

mickqian wants to merge 1 commit into
sgl-project:mainfrom
mickqian:codex/vlm-runtime-stability

Conversation

@mickqian

@mickqian mickqian commented Jul 11, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

  • make mRoPE and forward-batch metadata dynamic for piecewise CUDA graphs
  • avoid capture-stream misuse during runtime recompilation and preserve Kimi deepstack warmup inputs
  • fix Kimi MoonViT DP encoder behavior for TP=1 and avoid DP/TP CUDA-graph capture interference
  • remove TileLang import side effects from the Qwen/DeepSeek-V4 runtime path

Root cause

VLM-specific metadata and optional deepstack inputs were specialized by CUDA-graph warmup. Kimi's TP=1 DP vision path also returned a list of image tensors where the projector requires a tensor.

Validation

  • H100 CUDA: Kimi encoder tests (6 passed)
  • H100 CUDA: broader VLM/MoE regression set (67 passed)
  • ruff format --check, ruff check (with repository legacy import-style exclusions), git diff --check, and project pre-commit hooks

Impact

Eliminates runtime PCG recompiles for tested 128/256/512 token VLM shapes and restores Kimi-VL TP=1 serving.


CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.

@gemini-code-assist gemini-code-assist Bot left a comment

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.

Code Review

This pull request introduces several optimizations and bug fixes for multimodal models (VLMs), particularly focusing on CUDA graph support, parallel execution (DP/TP) for MoonViT (Kimi-VL), and profiling. Key changes include updating dynamic dimension inference for mRoPE models, caching sequence metadata to avoid FA3 synchronization, deferring MHC kernel imports in DeepSeek-V4 to prevent side effects, and warming up the deepstack branch in Qwen3-VL. Feedback from the reviewer highlights critical issues such as potential NameErrors due to a missing math import in mm_utils.py and a potentially missing inspect import in prefill_cuda_graph_runner.py. Additionally, the reviewer suggests fixing a library-matching bug in cuda_wrapper.py and adding a size limit to the position embedding cache in kimi_vl_moonvit.py to prevent memory exhaustion.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

continue
path = line[line.index("/") :].strip()
filename = os.path.basename(path)
if filename.rpartition(".so")[0].startswith(lib_name):

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.

high

Using filename.rpartition(".so")[0] will return an empty string "" if the filename is exactly libcudart.so (without any version suffix), because .rpartition splits at the last occurrence and the left part is empty. This causes .startswith(lib_name) to return False, silently ignoring the valid library. Using .partition(".so")[0] instead of .rpartition(".so")[0] correctly handles both versioned (e.g., libcudart.so.11.0) and unversioned (e.g., libcudart.so) filenames.

Suggested change
if filename.rpartition(".so")[0].startswith(lib_name):
if filename.partition(".so")[0].startswith(lib_name):

image_embeds = vision_model(
pixel_values,
grid_hw=grid_thw,
max_seqlen=max(math.prod(grid) for grid in grid_thw_list),

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.

high

math is not imported in this file, which will cause a NameError when calling math.prod(grid). Since rope_type == "rope_2d" guarantees a 2D grid, we can directly compute the product as grid[0] * grid[1] to avoid any external dependency or NameError.

Suggested change
max_seqlen=max(math.prod(grid) for grid in grid_thw_list),
max_seqlen=max(grid[0] * grid[1] for grid in grid_thw_list),

pixel_values_local, torch.tensor(local_grid_thw_list)
pixel_values_local,
grid_hw=local_grid_thw,
max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list),

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.

high

math is not imported in this file, which will cause a NameError when calling math.prod(grid). Since rope_type == "rope_2d" guarantees a 2D grid, we can directly compute the product as grid[0] * grid[1] to avoid any external dependency or NameError.

Suggested change
max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list),
max_seqlen=max(grid[0] * grid[1] for grid in local_grid_thw_list),

Comment on lines +253 to +254
if not self.training:
self._interpolated_pos_emb_cache[cache_key] = pos_emb

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.

medium

The _interpolated_pos_emb_cache dictionary can grow indefinitely if clients send images with highly diverse or arbitrary resolutions during long-running serving. To prevent potential memory exhaustion (OOM), we should add a simple size limit to clear the cache when it exceeds a threshold.

Suggested change
if not self.training:
self._interpolated_pos_emb_cache[cache_key] = pos_emb
if not self.training:
if len(self._interpolated_pos_emb_cache) >= 1024:
self._interpolated_pos_emb_cache.clear()
self._interpolated_pos_emb_cache[cache_key] = pos_emb

Comment on lines +511 to +514
if (
"input_deepstack_embeds"
not in inspect.signature(language_model.forward).parameters
):

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.

medium

inspect is used here but might not be imported at the module level in this file, potentially causing a NameError. Importing inspect inline makes this function self-contained and robust.

Suggested change
if (
"input_deepstack_embeds"
not in inspect.signature(language_model.forward).parameters
):
import inspect
if (
"input_deepstack_embeds"
not in inspect.signature(language_model.forward).parameters
):

@mickqian

Copy link
Copy Markdown
Collaborator Author

Superseded by smaller independently-reviewable PRs:

Each replacement PR contains its applicable measured performance or before/after behavior data.

@mickqian mickqian closed this Jul 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant