Repository navigation
Conversation
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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.
| 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), |
There was a problem hiding this comment.
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.
| 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), |
There was a problem hiding this comment.
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.
| 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), |
| if not self.training: | ||
| self._interpolated_pos_emb_cache[cache_key] = pos_emb |
There was a problem hiding this comment.
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.
| 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 |
| if ( | ||
| "input_deepstack_embeds" | ||
| not in inspect.signature(language_model.forward).parameters | ||
| ): |
There was a problem hiding this comment.
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.
| 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 | |
| ): |
|
Superseded by smaller independently-reviewable PRs:
Each replacement PR contains its applicable measured performance or before/after behavior data. |
Summary
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
ruff format --check,ruff check(with repository legacy import-style exclusions),git diff --check, and project pre-commit hooksImpact
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-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.