[Dev] Add Triton fused mRoPE for Qwen3.5-VL - #5962
Conversation
Signed-off-by: qizhang <qizhang@nvidia.com>
|
/claude strict-review |
Strict Review Summary — PR #5962: Triton fused mRoPE for Qwen3.5-VLReviewed the fused Triton mRoPE kernels (BSHD + THD), the RoPE dispatcher rework in Findings by severity — CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 1 IMPORTANT
SUGGESTION
Overall risk: Low-ModerateKernel correctness is solid and well-tested (70 passed / distributed CP verified). The main risks are the silent default-behavior flip for existing mrope users and the CPU-GPU sync overhead under CP, both addressable without changing the kernel. No correctness defects found in the rotary math, dtype handling, or CP partitioning. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| ) | ||
| if cp_size > 1: | ||
| # Guard: each packed sub-sequence length must satisfy seqlen % cp_size == 0. | ||
| seq_bounds = cu_seqlens.tolist() |
There was a problem hiding this comment.
[IMPORTANT Performance] cu_seqlens.tolist() forces a device→host copy and a full CUDA sync every time this guard runs. get_fused_mrope_thd_unavailable_reason is called on the per-layer attention hot path (via _try_fused_mrope_thd in rope_utils.apply_rotary_pos_emb), and it is then invoked a second time inside _launch_fused_mrope_thd's assert unavailable_reason is None — so with cp_size > 1 you pay two host syncs per attention layer per microbatch, serializing the CPU against the GPU.
Why it matters: this is precisely the kind of Python-side inspection of tensor values in the training loop that kills throughput and defeats overlap; the fused kernel was added to reduce overhead. Note the MLA yarn kernel deliberately does the analogous per-sub-sequence divisibility check inside the Triton kernel to avoid exactly this.
Suggestions: (1) skip the re-validation in _launch_fused_mrope_thd when the caller has already checked (or gate the assert behind a debug flag); and (2) move the per-sub-sequence divisibility guard into the kernel, or compute it with a device-side op (((cu_seqlens[1:] - cu_seqlens[:-1]) % cp_size != 0).any()) whose result you only materialize once, rather than .tolist() iterating on host.
| def _get_thd_raw_mrope_freqs_on_this_cp_rank( | ||
| cp_rank: int, cp_size: int, x: Tensor, freqs: Tensor, offset: int = 0 | ||
| ) -> Tensor: | ||
| """Get raw mRoPE frequency slices for this CP rank in THD layout.""" | ||
| if cp_size > 1: | ||
| first_cp_seg = (x.size(0) + 1) // 2 | ||
| second_cp_seg = x.size(0) // 2 | ||
| full_seqlen = cp_size * x.size(0) | ||
| return torch.cat( | ||
| [ | ||
| freqs[ | ||
| :, :, offset + cp_rank * first_cp_seg : offset + (cp_rank + 1) * first_cp_seg | ||
| ], | ||
| freqs[ | ||
| :, | ||
| :, | ||
| offset | ||
| + full_seqlen | ||
| - (cp_rank + 1) * second_cp_seg : offset | ||
| + full_seqlen | ||
| - cp_rank * second_cp_seg, | ||
| ], | ||
| ], | ||
| dim=2, | ||
| ) | ||
| else: | ||
| return freqs[:, :, offset : offset + x.size(0)] | ||
|
|
||
|
|
||
| def _get_thd_cp_splits(cu_seqlens: Tensor, cp_size: int) -> tuple[list[int], list[int]]: | ||
| """Return global sequence offsets and per-rank sequence lengths for THD CP fallback.""" | ||
| cu_seqlens_list = cu_seqlens.tolist() | ||
| local_seqlens = [] | ||
| for seq_start, seq_end in zip(cu_seqlens_list[:-1], cu_seqlens_list[1:]): | ||
| seq_len = seq_end - seq_start | ||
| if cp_size > 1 and seq_len % cp_size != 0: | ||
| raise ValueError( | ||
| "THD sequence lengths must be divisible by context parallel size, " | ||
| f"got sequence length {seq_len}, cp_size={cp_size}" | ||
| ) | ||
| local_seqlens.append(seq_len // cp_size) | ||
| return cu_seqlens_list, local_seqlens | ||
|
|
||
|
|
||
| def _pack_thd_raw_mrope_freqs( | ||
| t: Tensor, | ||
| cu_seqlens: Tensor, | ||
| freqs: Tensor, | ||
| cp_group: torch.distributed.ProcessGroup, | ||
| total_seqlen: Optional[int] = None, | ||
| ) -> Tensor: | ||
| """Pack raw mRoPE freqs into the same local token order as THD tensor ``t``.""" | ||
| cp_size = cp_group.size() | ||
| cp_rank = cp_group.rank() | ||
| cu_seqlens_list, seqlens = _get_thd_cp_splits(cu_seqlens, cp_size) | ||
| sequence_splits = torch.split(t, seqlens) | ||
| if total_seqlen is None: | ||
| total_seqlen = cu_seqlens_list[-1] | ||
| assert freqs.size(2) == total_seqlen, ( | ||
| f"raw mRoPE THD freqs sequence length {freqs.size(2)} must match " | ||
| f"cu_seqlens[-1] = {total_seqlen}" | ||
| ) | ||
|
|
||
| freq_slices = [] | ||
| for i, x in enumerate(sequence_splits): | ||
| seq_start_offset = cu_seqlens_list[i] | ||
| freq_slices.append( | ||
| _get_thd_raw_mrope_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) | ||
| ) | ||
|
|
||
| packed_freqs = torch.cat(freq_slices, dim=2) | ||
| assert packed_freqs.shape[2] == t.shape[0], ( | ||
| f"packed raw mRoPE freqs sequence length {packed_freqs.shape[2]} " | ||
| f"does not match THD tensor length {t.shape[0]}" | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION Simplification] _pack_thd_raw_mrope_freqs, together with its helpers _get_thd_cp_splits and _get_thd_raw_mrope_freqs_on_this_cp_rank, has no production call site. The only references are in tests/unit_tests/fusions/test_fused_mrope.py, where _pack_thd_raw_mrope_freqs is monkeypatched to an unexpected_pack that asserts it is never called ("raw THD mRoPE fusion should not materialize packed freqs").
Why it matters: the THD fused path passes the full unpacked freqs straight to the kernel, which does the CP position mapping internally, so this ~45-line packing utility (plus two helpers) is dead code. Carrying it adds maintenance surface and reads as an unfinished migration shim.
Suggestion: either remove all three functions, or if they are intended as the documented unfused-fallback packer for a future path, add a TODO explaining when they will be wired in. Per the review policy on newly-added unused identifiers, deferred usage should be flagged with an explicit TODO + migration note.
| if args.use_rotary_position_embeddings: | ||
| args.position_embedding_type = 'rope' | ||
| if args.position_embedding_type != 'rope': | ||
| if args.position_embedding_type not in ('rope', 'mrope'): |
There was a problem hiding this comment.
[IMPORTANT Compatibility] apply_rope_fusion defaults to True (it is only cleared via --no-rope-fusion). Previously, any position_embedding_type != 'rope' — including 'mrope' — force-disabled fusion here. After this change, existing mrope training scripts that never passed --no-rope-fusion will now silently start using the fused RoPE path (Triton fused mRoPE when available, else the TE/unfused fallbacks introduced in this PR).
Why it matters: this is a default-behavior change for the exact user population the PR targets. Numerically the fused path should match, but it flips the active code path for existing configs without an opt-in, and on hosts where Triton/TE support the config but the numerics differ subtly (bf16 rounding), it can perturb previously-validated runs.
Suggestion: confirm this default flip is intended and call it out in the PR description / release notes, and ideally verify golden values are unaffected for an existing mrope functional test. If a safer rollout is desired, gate the mrope fusion behind an explicit flag for one release before defaulting it on.
Signed-off-by: qizhang <qizhang@nvidia.com>
|
/ok to test eb73e39 |
Signed-off-by: qizhang <qizhang@nvidia.com>
3cdd48d to
5df1563
Compare
|
/ok to test 3454bb1 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30255237784 |
3 participants
Summary
This PR adds a Triton fused multimodal RoPE(mRoPE) implementation and integrates it with the Qwen3.5-VL training paths.
The fused kernel consumes raw T/H/W frequency tensors and performs axis selection, position mapping, sin/cos evaluation, rotary application, passthrough copying, and output casting in one kernel. This avoids materializing the full
RoPE embedding and several intermediate pointwise tensors.
Changes
[0, row_dim, column_dim]mRoPE and propagateapply_rope_fusionto the vision configuration.max_seqlenthrough the newly introduced THD fallback paths.ceil(L / 2)andfloor(L / 2).Unsupported configurations, including
mscale != 1, MLA-style interleaving, explicit inverse application, unsupported tensor layouts, or unavailable Triton support, fall back to the existing implementation.Testing
Added unit coverage for:
Validation results:
70 passed, 4 skipped.4 passedon each rank.