Repository navigation
Fix D2H sync in Qwen VL vision encoder - #47703
Mahnoor-Zaffar wants to merge 3 commits into
Conversation
Move small control-flow tensors to CPU to avoid GPU-CPU synchronization that breaks CUDA stream overlap and prevents torch.compile graph capture. Key changes: - Compute cu_seqlens on CPU in get_vision_cu_seqlens, then move to GPU for flash attention (tiny H2D transfer ~1μs vs 250ms+ D2H sync) - Move grid_thw to CPU before .tolist() in get_vision_position_ids - Move cu_seqlens to CPU before .tolist() in VisionAttention non-flash path - Move control-flow tensors to CPU in get_rope_index, replace itertools.groupby with tensor boundary detection - Pre-allocate mrope_position_deltas tensor instead of building from GPU scalars - Move cu_seqlens to CPU in get_max_seqlen before .max().item() These changes eliminate D2H syncs that caused: - 250ms+ stalls per forward pass in vision encoder - CUDA graph capture failures - torch.compile graph breaks Applied to Qwen2VL, Qwen2.5VL, and Qwen3VL vision models.
|
Hey, investigated this a bit more and redacting my previous comment - |
|
Fair point. The core issue is where the GPU sync happens. Calling .item() on a GPU tensor in a Python loop forces the CPU to stop and wait for the GPU to finish every single iteration, which is what causes the 250ms stalls. Moving the control tensor to CPU upfront still forces one initial D2H transfer, but it eliminates all the repeated stream syncs inside the loops/layers. I'll pull together a quick benchmark with PyTorch Profiler traces to show the CUDA stream overlap and before/after numbers. |
Rework the D2H-sync fix so only tiny control-flow tensors (cu_seqlens, cu_window_seqlens, mm_token_type_ids, attention_mask, grid_thw) are moved to CPU once before the Python loops in the vision encoders and get_rope_index, instead of per-layer .cpu() calls and .cpu()/int() forcing of model activations that preserved the syncs. Revert vision_utils/get_max_seqlen to pre-PR behavior. Regenerate qwen2_5_vl/qwen3_vl and dependent models from modular sources. Add a GPU benchmark (benchmark/d2h_syncs.py) reporting forward latency and D2H sync counts.
|
The original PR added The actual problem:
Per-layer/per-batch syncs are gone; with the processor default (grid_thw on CPU) it's zero, and even with inputs forced to CUDA only a couple of small copies remain in the vision tower. For qwen2_5_vl/qwen3_vl I edited the modular files and regenerated via the converter, which cascaded to all dependent models (qwen3_vl_moe, glm*, cosmos3*, paddleocr_vl, etc.). Also added Passed: modular/copies consistency, ruff, compile. |
| llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) | ||
| if attention_mask is not None: | ||
| position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) | ||
| position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions |
There was a problem hiding this comment.
This is subtle but I think attention_mask[batch_idx].bool() as an index will cause a sync, for weird internal PyTorch reasons. When you use mixed integers and bools like this, Torch will convert the boolean mask to an integer index tensor using nonzero(), but nonzero() causes a sync because the size of the output is data-dependent. I think if you visualize kernel executions here you'll see the nonzero and a sync.
If I'm right, you might be able to avoid it via something like attention_mask_cpu[batch_idx].nonzero(as_tuple=True)[0] and then move that to the GPU with to(non_blocking=True). This definitely requires analysis and benchmarking, though, don't trust me blindly, since I didn't test this!
There was a problem hiding this comment.
okay, I will research it thoroughly
Resolve the only merge conflict, in modular_qwen3_vl.py: upstream removed accepts_precomputed_kwargs and maybe_autocast from the utils.generic import, so drop those while keeping is_flash_attention_requested added by this PR. Modular conversion and copies checks both pass with no diffs.
|
Thank you for your contribution 🤗! CI Security Gate — automatic approval blockedThis PR was not automatically approved for CI because the security gate failed. Possible reasons:
See the workflow run for the exact violations. A maintainer can review and manually approve CI if a finding is a false positive. |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: cosmos3_edge, cosmos3_omni, exaone4_5, glm46v, glm4v, glm4v_moe, glm_ocr, minicpmv4_6, paddleocr_vl, qwen2_5_vl, qwen2_vl, qwen3_5, qwen3_5_moe, qwen3_omni_moe, qwen3_vl, qwen3_vl_moe |
CI recapDashboard: View test results in Grafana |
|
hmmm, tbh these are like processor stuff since it is single time input preparation. I am not sure how vllm calls it, but we defi have a plan to allow qwen's mrope to be called outisde model via a standlaone fn. It is being done for export and the like in #47974 |
Was profiling Qwen3VL inference and noticed ~250ms stalls in the vision encoder. Turned out we were doing D2H transfers on every attention layer when splitting sequences.
The issue is that
cu_seqlensandgrid_thware tiny control-flow tensors that were living on GPU, forcing syncs whenever we called.tolist()or.item()on them. This breaks CUDA stream overlap and prevents torch.compile from capturing the whole thing as one graph.Fixed by moving these small tensors to CPU before the Python loops. The H2D transfer for flash attention is negligible (<100 elements). Also replaced the
itertools.groupbywith tensor boundary detection inget_rope_indexto avoid more syncs.Changes:
get_vision_cu_seqlens: compute on CPU, move to GPU for flashVisionAttention: movecu_seqlensto CPU in non-flash pathget_vision_position_ids: useint()on CPU tensor instead of.item()get_rope_index: move control tensors to CPU, pre-allocate deltasget_max_seqlen: move to CPU before.max().item()Tested the algorithmic changes (tensor groupby logic, cu_seqlens computation, etc.) - all match the old behavior. Full model tests need PyTorch environment (Python 3.14 doesn't have wheels yet).
Should eliminate the stalls and make the vision encoder torch.compile-friendly.
Rebased on latest main
Merged latest
main(360 commits) into this branch. The only conflict was the...utils.genericimport inmodular_qwen3_vl.py: upstream removedaccepts_precomputed_kwargsandmaybe_autocast, so those were dropped while keeping the newly addedis_flash_attention_requested.Verified after the merge:
utils/check_modular_conversion.pypasses (no diffs;--fix_and_overwriteproduces zero changes).utils/check_copies.pypasses (dependent models: qwen3_vl_moe, glm*, cosmos3*, paddleocr_vl, exaone4_5, etc.).ruff checkclean on all touched files.GPU verification still pending
The authoring environment has no GPU and no PyTorch wheels for the local Python version, so the following still need to be run on a GPU runner (CPU CI above covers the algorithmic parity checks):
torch.profilerbefore/after D2H counts onQwen/Qwen3-VL-2B(--scope vision/--scope fullviabenchmark/d2h_syncs.py).tests/models/qwen2_5_vl,tests/models/qwen3_vl, ...) in a PyTorch environment.