Fused MRoPE - #8
Merged
wplf merged 5 commits intoMay 30, 2026
Merged
Conversation
Reviewer's GuideIntegrates a new Triton-based fused multimodal RoPE (mRoPE) implementation for both standard and THD-packed layouts, wires it into GPT and Qwen3.5-VL vision paths with configuration/dispatch logic and detailed fallbacks, and extends scripts/tests to cover fusion behavior, NVTX profiling, and multimodal configs. Sequence diagram for fused mRoPE dispatch in apply_rotary_pos_embsequenceDiagram
participant GPTModel as GPTModel_preprocess
participant MRoPE as MultimodalRotaryEmbedding_forward
participant Rope as apply_rotary_pos_emb
participant Triton as fused_apply_mrope
participant TE as fused_apply_rotary_pos_emb
participant Unfused as _apply_rotary_pos_emb_bshd
GPTModel->>MRoPE: rotary_pos_emb(position_ids, mrope_section,\n cp_group, return_raw_freqs, packed_seq)
MRoPE-->>GPTModel: freqs_or_emb
GPTModel->>Rope: apply_rotary_pos_emb(t, freqs, config,\n cu_seqlens=None,...)
alt apply_rope_fusion and raw mRoPE freqs
Rope->>Triton: fused_apply_mrope(t, freqs, mrope_section,\n interleaved_mrope, rotary_interleaved=False)
Triton-->>Rope: rotated_t
Rope-->>GPTModel: rotated_t
else TE fused RoPE available
Rope->>TE: fused_apply_rotary_pos_emb(t, freqs,\n interleaved=config.rotary_interleaved)
TE-->>Rope: rotated_t
Rope-->>GPTModel: rotated_t
else
Rope->>Unfused: _apply_rotary_pos_emb_bshd(t, freqs,...)
Unfused-->>Rope: rotated_t
Rope-->>GPTModel: rotated_t
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 3 issues, and left some high level feedback:
- The THD context-parallel index arithmetic for zigzag splits is now implemented in multiple places (e.g.,
_get_thd_freqs_on_this_cp_rank,_get_thd_raw_mrope_freqs_on_this_cp_rank,_get_thd_cp_splits, and_get_thd_token_idxinfused_mla_yarn_rope_apply); consider factoring this into a single shared helper to avoid drift and make future changes less error-prone. - The logic that toggles between raw mRoPE freqs and materialized rotary embeddings is spread across
MultimodalRotaryEmbedding.forward,apply_rotary_pos_emb, andGPTModel._preprocess; it may be worth centralizing this policy (e.g., a helper that decides when to return raw vs materialized based on config/inference mode) to reduce duplication and keep behavior consistent across call sites.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The THD context-parallel index arithmetic for zigzag splits is now implemented in multiple places (e.g., `_get_thd_freqs_on_this_cp_rank`, `_get_thd_raw_mrope_freqs_on_this_cp_rank`, `_get_thd_cp_splits`, and `_get_thd_token_idx` in `fused_mla_yarn_rope_apply`); consider factoring this into a single shared helper to avoid drift and make future changes less error-prone.
- The logic that toggles between raw mRoPE freqs and materialized rotary embeddings is spread across `MultimodalRotaryEmbedding.forward`, `apply_rotary_pos_emb`, and `GPTModel._preprocess`; it may be worth centralizing this policy (e.g., a helper that decides when to return raw vs materialized based on config/inference mode) to reduce duplication and keep behavior consistent across call sites.
## Individual Comments
### Comment 1
<location path="megatron/core/models/common/embeddings/rope_utils.py" line_range="505-507" />
<code_context>
if cp_group is None:
cp_group = parallel_state.get_context_parallel_group()
+ is_raw_mrope_freqs = (
+ _is_raw_mrope_freqs(t, freqs, config)
+ if cu_seqlens is None
+ else _is_raw_mrope_freqs_thd(t, freqs, cu_seqlens, config, cp_group.size())
+ )
</code_context>
<issue_to_address>
**issue (bug_risk):** Guard `cp_group` before calling `.size()` in THD mRoPE detection to avoid `NoneType` errors.
In the THD branch, `_is_raw_mrope_freqs_thd` is called with `cp_group.size()`, but `cp_group` comes from `parallel_state.get_context_parallel_group()`, which can be `None` in non-CP setups. That will raise an AttributeError before the later `_apply_rotary_pos_emb_thd` check that handles `cp_group` explicitly.
Consider either failing fast with a clear assertion/error when `cp_group` is `None` in the THD path, or treating `cp_group is None` as `cp_size == 1` and passing `cp_size=1` into `_is_raw_mrope_freqs_thd`, so the behavior and error messaging remain intentional rather than a `NoneType` attribute error.
</issue_to_address>
### Comment 2
<location path="tests/unit_tests/fusions/test_fused_mrope.py" line_range="195-196" />
<code_context>
+ assert len(set(actual)) == len(actual)
+
+
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(use_packed_seq):
+ captured_kwargs = {}
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add a regression test that covers the fallback path when fused mRoPE is not available in GPT `_preprocess`.
To increase coverage, please also exercise the branch where `_fused_mrope_available` (or `apply_rope_fusion`) is `False` and confirm `_preprocess` continues to use materialized RoPE embeddings instead of raw freqs. For example, parametrize this test (or add a sibling test) to run with `_fused_mrope_available=False` and assert that `return_raw_freqs` stays `False` and `rotary_pos_emb` comes from the materialized embedding path.
Suggested implementation:
```python
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+@pytest.mark.parametrize("fused_available", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(
+ use_packed_seq, fused_available, monkeypatch
+):
+ """
+ When fused mRoPE is available, GPT._preprocess should request raw freqs.
+ When it is not available, GPT._preprocess should fall back to materialized RoPE
+ embeddings and keep return_raw_freqs=False.
+ """
+ captured_kwargs = {}
+
+ # Patch the fused mRoPE availability / fusion entry point that _preprocess consults.
+ # We capture the arguments to inspect return_raw_freqs and rotary_pos_emb.
+ def _capturing_fused_apply_mrope(*args, **kwargs):
+ captured_kwargs.update(kwargs)
+ # We do not care about the numerical result here; let the original function run
+ # when fusion is "available", and raise when it's disabled so we know that path
+ # is not mistakenly taken.
+ if not fused_available:
+ raise AssertionError("fused_apply_mrope should not be called when fusion is disabled")
+ return _orig_fused_apply_mrope(*args, **kwargs)
+
+ _orig_fused_apply_mrope = fused_apply_mrope
+ monkeypatch.setattr(
+ rope_utils,
+ "apply_rope_fusion",
+ lambda: fused_available,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "megatron.core.fusions.fused_mrope.fused_apply_mrope",
+ _capturing_fused_apply_mrope,
+ )
+
+ # Build a tiny GPT model configured to use mRoPE in eval mode.
+ # The helper should return a model where forward() goes through _preprocess.
+ gpt = build_gpt_for_test(use_packed_seq=use_packed_seq, use_mrope=True)
+ gpt.eval()
+
+ # Run one forward pass to trigger _preprocess and (potential) fusion.
+ batch_size, seq_len = 2, 8
+ input_ids = torch.randint(0, gpt.config.vocab_size, (batch_size, seq_len), device="cuda")
+ with torch.no_grad():
+ _ = gpt(input_ids)
+
+ if fused_available:
+ # Fusion path: raw freqs should be requested.
+ assert captured_kwargs, "Expected fused_apply_mrope to be called when fusion is available"
+ assert captured_kwargs.get("return_raw_freqs", False) is True
+ else:
+ # Fallback path: fused_apply_mrope must not be called, and _preprocess should
+ # have used materialized rotary_pos_emb instead of raw freqs.
+ assert captured_kwargs == {}, "fused_apply_mrope should not be used when fusion is disabled"
+ # Inspect the rope state attached to the model to ensure materialized embeddings were used.
+ rope_state = getattr(gpt, "_rope_state", None)
+ assert rope_state is not None, "Expected GPT to store RoPE state for materialized embeddings"
+ assert rope_state.return_raw_freqs is False
+ assert rope_state.rotary_pos_emb is not None
+ assert isinstance(rope_state.rotary_pos_emb, torch.Tensor)
+
```
`.
```xml
<file_operations>
<file_operation operation="edit" file_path="tests/unit_tests/fusions/test_fused_mrope.py">
<<<<<<< SEARCH
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(use_packed_seq):
+ captured_kwargs = {}
+
=======
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+@pytest.mark.parametrize("fused_available", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(
+ use_packed_seq, fused_available, monkeypatch
+):
+ """
+ When fused mRoPE is available, GPT._preprocess should request raw freqs.
+ When it is not available, GPT._preprocess should fall back to materialized RoPE
+ embeddings and keep return_raw_freqs=False.
+ """
+ captured_kwargs = {}
+
+ # Patch the fused mRoPE availability / fusion entry point that _preprocess consults.
+ # We capture the arguments to inspect return_raw_freqs and rotary_pos_emb.
+ def _capturing_fused_apply_mrope(*args, **kwargs):
+ captured_kwargs.update(kwargs)
+ # We do not care about the numerical result here; let the original function run
+ # when fusion is "available", and raise when it's disabled so we know that path
+ # is not mistakenly taken.
+ if not fused_available:
+ raise AssertionError("fused_apply_mrope should not be called when fusion is disabled")
+ return _orig_fused_apply_mrope(*args, **kwargs)
+
+ _orig_fused_apply_mrope = fused_apply_mrope
+ monkeypatch.setattr(
+ rope_utils,
+ "apply_rope_fusion",
+ lambda: fused_available,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "megatron.core.fusions.fused_mrope.fused_apply_mrope",
+ _capturing_fused_apply_mrope,
+ )
+
+ # Build a tiny GPT model configured to use mRoPE in eval mode.
+ # The helper should return a model where forward() goes through _preprocess.
+ gpt = build_gpt_for_test(use_packed_seq=use_packed_seq, use_mrope=True)
+ gpt.eval()
+
+ # Run one forward pass to trigger _preprocess and (potential) fusion.
+ batch_size, seq_len = 2, 8
+ input_ids = torch.randint(0, gpt.config.vocab_size, (batch_size, seq_len), device="cuda")
+ with torch.no_grad():
+ _ = gpt(input_ids)
+
+ if fused_available:
+ # Fusion path: raw freqs should be requested.
+ assert captured_kwargs, "Expected fused_apply_mrope to be called when fusion is available"
+ assert captured_kwargs.get("return_raw_freqs", False) is True
+ else:
+ # Fallback path: fused_apply_mrope must not be called, and _preprocess should
+ # have used materialized rotary_pos_emb instead of raw freqs.
+ assert captured_kwargs == {}, "fused_apply_mrope should not be used when fusion is disabled"
+ # Inspect the rope state attached to the model to ensure materialized embeddings were used.
+ rope_state = getattr(gpt, "_rope_state", None)
+ assert rope_state is not None, "Expected GPT to store RoPE state for materialized embeddings"
+ assert rope_state.return_raw_freqs is False
+ assert rope_state.rotary_pos_emb is not None
+ assert isinstance(rope_state.rotary_pos_emb, torch.Tensor)
+
>>>>>>> REPLACE
</file_operation>
</file_operations>
<additional_changes>
1. Ensure there is an import for `build_gpt_for_test` or adjust the construction of the GPT model to match your existing test helpers. For example, if you already use a `build_model` or `get_gpt_model` helper elsewhere, use that instead of `build_gpt_for_test`.
2. The test assumes that:
- `rope_utils.apply_rope_fusion()` (or a similarly named function) is what `_preprocess` checks to determine fusion availability. If the real name differs, update the `monkeypatch.setattr(rope_utils, "apply_rope_fusion", ...)` call accordingly.
- GPT attaches RoPE-related information to `gpt._rope_state` with attributes `return_raw_freqs` and `rotary_pos_emb`. If your implementation exposes these under a different name or location, change the `rope_state` lookup and assertions to match.
3. If `_preprocess` does not call `fused_apply_mrope` directly but goes through another wrapper, adjust the `monkeypatch.setattr("megatron.core.fusions.fused_mrope.fused_apply_mrope", ...)` target string to hook the actually-called symbol.
4. If your tests do not run on CUDA, change `device="cuda"` to `"cpu"` or use `gpt.device` / `next(gpt.parameters()).device` so the tensor is on the correct device.
</issue_to_address>
### Comment 3
<location path="examples/multimodal_dev/tests/test_vision_rope_fusion.py" line_range="71-80" />
<code_context>
+def test_vision_fp32_wrapper_dispatches_raw_freqs_to_fused_mrope_thd(monkeypatch):
</code_context>
<issue_to_address>
**suggestion (testing):** Complement this positive-path dispatch test with a fallback-path test where fused mRoPE THD is unavailable.
To complete coverage of this wrapper, please also add a test where `get_fused_mrope_thd_unavailable_reason` returns a non-`None` string and verify that `_apply_rope_fp32_no_cp` falls back to the unfused `apply_rotary_pos_emb` path (e.g., by asserting `fake_fused_apply_mrope_thd` is not called and the rope-utils path runs instead). This will help catch regressions in the fallback branch for vision THD RoPE.
Suggested implementation:
```python
torch.testing.assert_close(converted, expected)
def test_vision_fp32_wrapper_dispatches_fallback_when_fused_mrope_thd_unavailable(monkeypatch):
# We want to verify that when fused mRoPE THD is unavailable, the wrapper
# falls back to the unfused `apply_rotary_pos_emb` path.
# Local imports to avoid affecting other tests; adjust module paths as needed.
import torch
# The module under test and its dependency names will need to match the
# actual implementation; these strings may need adjustment.
module_path = "examples.multimodal_dev.vision_rope_fusion"
calls = {
"fused_mrope": 0,
"unfused_rope": 0,
}
def fake_get_fused_mrope_thd_unavailable_reason() -> str:
# Non-None return value should trigger the fallback path.
return "test: fused mRoPE THD unavailable"
def fake_fused_apply_mrope_thd(
t,
cu_seqlens,
freqs,
mrope_section,
interleaved_mrope=False,
rotary_interleaved=False,
cp_size=1,
):
# If this gets called in this test, the wrapper did not take the fallback.
calls["fused_mrope"] += 1
return t
def fake_apply_rotary_pos_emb(x, freqs, *args, **kwargs):
# Track that the unfused RoPE path ran.
calls["unfused_rope"] += 1
return x
# Monkeypatch the unavailability probe to force the fallback path.
monkeypatch.setattr(
f"{module_path}.get_fused_mrope_thd_unavailable_reason",
lambda: fake_get_fused_mrope_thd_unavailable_reason(),
raising=True,
)
# Monkeypatch both the fused mRoPE implementation and the unfused RoPE util.
monkeypatch.setattr(
f"{module_path}.apply_fused_mrope_thd",
fake_fused_apply_mrope_thd,
raising=True,
)
monkeypatch.setattr(
f"{module_path}.apply_rotary_pos_emb",
fake_apply_rotary_pos_emb,
raising=True,
)
# Import the wrapper under test after monkeypatching so it picks up patched symbols.
from examples.multimodal_dev.vision_rope_fusion import _apply_rope_fp32_no_cp
# Minimal tensor inputs; precise shapes are not important for this behavioral test.
t = torch.randn(2, 4, 8, dtype=torch.float32)
cu_seqlens = torch.tensor([0, 2], dtype=torch.int32)
freqs = torch.randn(1, 1, 4, 4, dtype=torch.float32)
mrope_section = [0, 4, 4]
# Call the wrapper; with fused mRoPE marked unavailable, it should take the fallback.
_apply_rope_fp32_no_cp(
t=t,
cu_seqlens=cu_seqlens,
freqs=freqs,
mrope_section=mrope_section,
interleaved_mrope=False,
rotary_interleaved=False,
)
# Assert that the fused path was not used and the unfused RoPE path was taken.
assert calls["fused_mrope"] == 0
assert calls["unfused_rope"] == 1
```
1. Update `module_path = "examples.multimodal_dev.vision_rope_fusion"` and the two monkeypatch targets (`apply_fused_mrope_thd`, `apply_rotary_pos_emb`) to match the actual module and function names used by `_apply_rope_fp32_no_cp` in your codebase. The key is that:
- `get_fused_mrope_thd_unavailable_reason` is patched to return a non-`None` string.
- The fused mRoPE THD implementation used by `_apply_rope_fp32_no_cp` is patched as `fake_fused_apply_mrope_thd`.
- The unfused RoPE helper called by `_apply_rope_fp32_no_cp` is patched as `fake_apply_rotary_pos_emb`.
2. If `_apply_rope_fp32_no_cp` requires additional parameters in your implementation (e.g., `dtype`, `device`, or vision-specific flags), pass appropriate dummy values in the test call so it exercises the exact fallback branch you want to cover.
3. If the helper functions are imported into the module under different local names (e.g., `from ... import apply_rotary_pos_emb as apply_vision_rope`), adjust the monkeypatch targets to those names so the patched functions are actually used.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+505
to
+507
| is_raw_mrope_freqs = ( | ||
| _is_raw_mrope_freqs(t, freqs, config) | ||
| if cu_seqlens is None |
There was a problem hiding this comment.
issue (bug_risk): Guard cp_group before calling .size() in THD mRoPE detection to avoid NoneType errors.
In the THD branch, _is_raw_mrope_freqs_thd is called with cp_group.size(), but cp_group comes from parallel_state.get_context_parallel_group(), which can be None in non-CP setups. That will raise an AttributeError before the later _apply_rotary_pos_emb_thd check that handles cp_group explicitly.
Consider either failing fast with a clear assertion/error when cp_group is None in the THD path, or treating cp_group is None as cp_size == 1 and passing cp_size=1 into _is_raw_mrope_freqs_thd, so the behavior and error messaging remain intentional rather than a NoneType attribute error.
Comment on lines
+195
to
+196
| @pytest.mark.parametrize("use_packed_seq", [False, True]) | ||
| def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(use_packed_seq): |
There was a problem hiding this comment.
suggestion (testing): Add a regression test that covers the fallback path when fused mRoPE is not available in GPT _preprocess.
To increase coverage, please also exercise the branch where _fused_mrope_available (or apply_rope_fusion) is False and confirm _preprocess continues to use materialized RoPE embeddings instead of raw freqs. For example, parametrize this test (or add a sibling test) to run with _fused_mrope_available=False and assert that return_raw_freqs stays False and rotary_pos_emb comes from the materialized embedding path.
Suggested implementation:
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+@pytest.mark.parametrize("fused_available", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(
+ use_packed_seq, fused_available, monkeypatch
+):
+ """
+ When fused mRoPE is available, GPT._preprocess should request raw freqs.
+ When it is not available, GPT._preprocess should fall back to materialized RoPE
+ embeddings and keep return_raw_freqs=False.
+ """
+ captured_kwargs = {}
+
+ # Patch the fused mRoPE availability / fusion entry point that _preprocess consults.
+ # We capture the arguments to inspect return_raw_freqs and rotary_pos_emb.
+ def _capturing_fused_apply_mrope(*args, **kwargs):
+ captured_kwargs.update(kwargs)
+ # We do not care about the numerical result here; let the original function run
+ # when fusion is "available", and raise when it's disabled so we know that path
+ # is not mistakenly taken.
+ if not fused_available:
+ raise AssertionError("fused_apply_mrope should not be called when fusion is disabled")
+ return _orig_fused_apply_mrope(*args, **kwargs)
+
+ _orig_fused_apply_mrope = fused_apply_mrope
+ monkeypatch.setattr(
+ rope_utils,
+ "apply_rope_fusion",
+ lambda: fused_available,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "megatron.core.fusions.fused_mrope.fused_apply_mrope",
+ _capturing_fused_apply_mrope,
+ )
+
+ # Build a tiny GPT model configured to use mRoPE in eval mode.
+ # The helper should return a model where forward() goes through _preprocess.
+ gpt = build_gpt_for_test(use_packed_seq=use_packed_seq, use_mrope=True)
+ gpt.eval()
+
+ # Run one forward pass to trigger _preprocess and (potential) fusion.
+ batch_size, seq_len = 2, 8
+ input_ids = torch.randint(0, gpt.config.vocab_size, (batch_size, seq_len), device="cuda")
+ with torch.no_grad():
+ _ = gpt(input_ids)
+
+ if fused_available:
+ # Fusion path: raw freqs should be requested.
+ assert captured_kwargs, "Expected fused_apply_mrope to be called when fusion is available"
+ assert captured_kwargs.get("return_raw_freqs", False) is True
+ else:
+ # Fallback path: fused_apply_mrope must not be called, and _preprocess should
+ # have used materialized rotary_pos_emb instead of raw freqs.
+ assert captured_kwargs == {}, "fused_apply_mrope should not be used when fusion is disabled"
+ # Inspect the rope state attached to the model to ensure materialized embeddings were used.
+ rope_state = getattr(gpt, "_rope_state", None)
+ assert rope_state is not None, "Expected GPT to store RoPE state for materialized embeddings"
+ assert rope_state.return_raw_freqs is False
+ assert rope_state.rotary_pos_emb is not None
+ assert isinstance(rope_state.rotary_pos_emb, torch.Tensor)
+`.
<file_operations>
<file_operation operation="edit" file_path="tests/unit_tests/fusions/test_fused_mrope.py">
<<<<<<< SEARCH
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(use_packed_seq):
+ captured_kwargs = {}
+
=======
+@pytest.mark.parametrize("use_packed_seq", [False, True])
+@pytest.mark.parametrize("fused_available", [False, True])
+def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(
+ use_packed_seq, fused_available, monkeypatch
+):
+ """
+ When fused mRoPE is available, GPT._preprocess should request raw freqs.
+ When it is not available, GPT._preprocess should fall back to materialized RoPE
+ embeddings and keep return_raw_freqs=False.
+ """
+ captured_kwargs = {}
+
+ # Patch the fused mRoPE availability / fusion entry point that _preprocess consults.
+ # We capture the arguments to inspect return_raw_freqs and rotary_pos_emb.
+ def _capturing_fused_apply_mrope(*args, **kwargs):
+ captured_kwargs.update(kwargs)
+ # We do not care about the numerical result here; let the original function run
+ # when fusion is "available", and raise when it's disabled so we know that path
+ # is not mistakenly taken.
+ if not fused_available:
+ raise AssertionError("fused_apply_mrope should not be called when fusion is disabled")
+ return _orig_fused_apply_mrope(*args, **kwargs)
+
+ _orig_fused_apply_mrope = fused_apply_mrope
+ monkeypatch.setattr(
+ rope_utils,
+ "apply_rope_fusion",
+ lambda: fused_available,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "megatron.core.fusions.fused_mrope.fused_apply_mrope",
+ _capturing_fused_apply_mrope,
+ )
+
+ # Build a tiny GPT model configured to use mRoPE in eval mode.
+ # The helper should return a model where forward() goes through _preprocess.
+ gpt = build_gpt_for_test(use_packed_seq=use_packed_seq, use_mrope=True)
+ gpt.eval()
+
+ # Run one forward pass to trigger _preprocess and (potential) fusion.
+ batch_size, seq_len = 2, 8
+ input_ids = torch.randint(0, gpt.config.vocab_size, (batch_size, seq_len), device="cuda")
+ with torch.no_grad():
+ _ = gpt(input_ids)
+
+ if fused_available:
+ # Fusion path: raw freqs should be requested.
+ assert captured_kwargs, "Expected fused_apply_mrope to be called when fusion is available"
+ assert captured_kwargs.get("return_raw_freqs", False) is True
+ else:
+ # Fallback path: fused_apply_mrope must not be called, and _preprocess should
+ # have used materialized rotary_pos_emb instead of raw freqs.
+ assert captured_kwargs == {}, "fused_apply_mrope should not be used when fusion is disabled"
+ # Inspect the rope state attached to the model to ensure materialized embeddings were used.
+ rope_state = getattr(gpt, "_rope_state", None)
+ assert rope_state is not None, "Expected GPT to store RoPE state for materialized embeddings"
+ assert rope_state.return_raw_freqs is False
+ assert rope_state.rotary_pos_emb is not None
+ assert isinstance(rope_state.rotary_pos_emb, torch.Tensor)
+
>>>>>>> REPLACE
</file_operation>
</file_operations>
<additional_changes>
1. Ensure there is an import for `build_gpt_for_test` or adjust the construction of the GPT model to match your existing test helpers. For example, if you already use a `build_model` or `get_gpt_model` helper elsewhere, use that instead of `build_gpt_for_test`.
2. The test assumes that:
- `rope_utils.apply_rope_fusion()` (or a similarly named function) is what `_preprocess` checks to determine fusion availability. If the real name differs, update the `monkeypatch.setattr(rope_utils, "apply_rope_fusion", ...)` call accordingly.
- GPT attaches RoPE-related information to `gpt._rope_state` with attributes `return_raw_freqs` and `rotary_pos_emb`. If your implementation exposes these under a different name or location, change the `rope_state` lookup and assertions to match.
3. If `_preprocess` does not call `fused_apply_mrope` directly but goes through another wrapper, adjust the `monkeypatch.setattr("megatron.core.fusions.fused_mrope.fused_apply_mrope", ...)` target string to hook the actually-called symbol.
4. If your tests do not run on CUDA, change `device="cuda"` to `"cpu"` or use `gpt.device` / `next(gpt.parameters()).device` so the tensor is on the correct device.
Comment on lines
+71
to
+80
| def test_vision_fp32_wrapper_dispatches_raw_freqs_to_fused_mrope_thd(monkeypatch): | ||
| calls = {} | ||
|
|
||
| def fake_fused_apply_mrope_thd( | ||
| t, | ||
| cu_seqlens, | ||
| freqs, | ||
| mrope_section, | ||
| interleaved_mrope=False, | ||
| rotary_interleaved=False, |
There was a problem hiding this comment.
suggestion (testing): Complement this positive-path dispatch test with a fallback-path test where fused mRoPE THD is unavailable.
To complete coverage of this wrapper, please also add a test where get_fused_mrope_thd_unavailable_reason returns a non-None string and verify that _apply_rope_fp32_no_cp falls back to the unfused apply_rotary_pos_emb path (e.g., by asserting fake_fused_apply_mrope_thd is not called and the rope-utils path runs instead). This will help catch regressions in the fallback branch for vision THD RoPE.
Suggested implementation:
torch.testing.assert_close(converted, expected)
def test_vision_fp32_wrapper_dispatches_fallback_when_fused_mrope_thd_unavailable(monkeypatch):
# We want to verify that when fused mRoPE THD is unavailable, the wrapper
# falls back to the unfused `apply_rotary_pos_emb` path.
# Local imports to avoid affecting other tests; adjust module paths as needed.
import torch
# The module under test and its dependency names will need to match the
# actual implementation; these strings may need adjustment.
module_path = "examples.multimodal_dev.vision_rope_fusion"
calls = {
"fused_mrope": 0,
"unfused_rope": 0,
}
def fake_get_fused_mrope_thd_unavailable_reason() -> str:
# Non-None return value should trigger the fallback path.
return "test: fused mRoPE THD unavailable"
def fake_fused_apply_mrope_thd(
t,
cu_seqlens,
freqs,
mrope_section,
interleaved_mrope=False,
rotary_interleaved=False,
cp_size=1,
):
# If this gets called in this test, the wrapper did not take the fallback.
calls["fused_mrope"] += 1
return t
def fake_apply_rotary_pos_emb(x, freqs, *args, **kwargs):
# Track that the unfused RoPE path ran.
calls["unfused_rope"] += 1
return x
# Monkeypatch the unavailability probe to force the fallback path.
monkeypatch.setattr(
f"{module_path}.get_fused_mrope_thd_unavailable_reason",
lambda: fake_get_fused_mrope_thd_unavailable_reason(),
raising=True,
)
# Monkeypatch both the fused mRoPE implementation and the unfused RoPE util.
monkeypatch.setattr(
f"{module_path}.apply_fused_mrope_thd",
fake_fused_apply_mrope_thd,
raising=True,
)
monkeypatch.setattr(
f"{module_path}.apply_rotary_pos_emb",
fake_apply_rotary_pos_emb,
raising=True,
)
# Import the wrapper under test after monkeypatching so it picks up patched symbols.
from examples.multimodal_dev.vision_rope_fusion import _apply_rope_fp32_no_cp
# Minimal tensor inputs; precise shapes are not important for this behavioral test.
t = torch.randn(2, 4, 8, dtype=torch.float32)
cu_seqlens = torch.tensor([0, 2], dtype=torch.int32)
freqs = torch.randn(1, 1, 4, 4, dtype=torch.float32)
mrope_section = [0, 4, 4]
# Call the wrapper; with fused mRoPE marked unavailable, it should take the fallback.
_apply_rope_fp32_no_cp(
t=t,
cu_seqlens=cu_seqlens,
freqs=freqs,
mrope_section=mrope_section,
interleaved_mrope=False,
rotary_interleaved=False,
)
# Assert that the fused path was not used and the unfused RoPE path was taken.
assert calls["fused_mrope"] == 0
assert calls["unfused_rope"] == 1- Update
module_path = "examples.multimodal_dev.vision_rope_fusion"and the two monkeypatch targets (apply_fused_mrope_thd,apply_rotary_pos_emb) to match the actual module and function names used by_apply_rope_fp32_no_cpin your codebase. The key is that:get_fused_mrope_thd_unavailable_reasonis patched to return a non-Nonestring.- The fused mRoPE THD implementation used by
_apply_rope_fp32_no_cpis patched asfake_fused_apply_mrope_thd. - The unfused RoPE helper called by
_apply_rope_fp32_no_cpis patched asfake_apply_rotary_pos_emb.
- If
_apply_rope_fp32_no_cprequires additional parameters in your implementation (e.g.,dtype,device, or vision-specific flags), pass appropriate dummy values in the test call so it exercises the exact fallback branch you want to cover. - If the helper functions are imported into the module under different local names (e.g.,
from ... import apply_rotary_pos_emb as apply_vision_rope), adjust the monkeypatch targets to those names so the patched functions are actually used.
…aunch path The fused THD dispatch (rope_utils.apply_rotary_pos_emb -> fused_apply_mrope_thd) calls the kernel directly and only validated total seqlen % cp_size, not each packed sub-sequence. The unfused per-sequence check in _get_thd_cp_splits() is bypassed on the fused path, so for CP>1 variable-length packing where the total is divisible but an individual sub-sequence is not, the kernel would silently compute wrong local->global CP token indices (global_start // cp_size). Add the per-sequence guard in get_fused_mrope_thd_unavailable_reason (only on the cp_size>1 path); when it triggers, the dispatch falls back to the unfused path.
…ivisibility - Add fwd/bwd parity at the real deployment shape head_dim=256, rotary_dim=64 (rotary_percent=0.25, 75% pass-through) with mrope_section=[11,11,10], plus the non-interleaved and full-rotary variants, for both BSHD and THD. The existing parametrized tests only covered head_dim=16/20 with rotary_dim=16 (~80% rotated). - Add a regression test that get_fused_mrope_thd_unavailable_reason rejects a packed batch whose total length is CP-divisible but an individual sub-sequence is not (and accepts the all-divisible control).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Integrate Triton-based fused multimodal RoPE (mRoPE) kernels (standard BSHD and THD-packed layouts) and wire them into the GPT decoder and the Qwen3.5-VL vision encoder, with full autograd support and a safe unfused fallback whenever fusion is unavailable.
Implementation
megatron/core/fusions/fused_mrope.py: Triton forward/backward kernels for BSHD and THD layouts, a raw→materialized conversion helper (mrope_freqs_to_rotary_emb), and capability/availability gates (get_fused_mrope_unavailable_reason,get_fused_mrope_thd_unavailable_reason).gpt_model._preprocessrequests raw per-axis mRoPE frequencies when fusion is available;rope_utils.apply_rotary_pos_embdispatches fused vs. unfused and centralizes one-time fallback warnings.apply_rope_fusionvalidation relaxed to accept Triton fused mRoPE as a backend alongside Transformer Engine; CP-aware THD handling, incl. the odd-local-seqlen fix.A detailed auto-generated breakdown is in Summary by Sourcery below.
Validation & numerical verification
Operator-level backward — verified correct three independent ways (all errors at fp32 epsilon ~1e-7; covers BSHD + THD, interleaved + non-interleaved, including the real
head_dim=256 / rotary_dim=64 / mrope_section=[11,11,10]shape):torch.autograd.gradcheck: PASSED<R·v, g> == <v, grad>(RoPE is linear): ~1e-7Coverage: added forward/backward parity tests at the real Qwen3.5-VL deployment shape (
head_dim=256, rotary_dim=64, 75% pass-through,[11,11,10]), plus non-interleaved and full-rotary variants, for both BSHD and THD — the previous parametrized tests only coveredhead_dim=16/20withrotary_dim=16(~80% rotated). Unit suite: 65 passed on GB200.Robustness: the fused THD launch path now enforces per-sequence CP divisibility in
get_fused_mrope_thd_unavailable_reason(not just total length). Variable-length packing where the total is CP-divisible but an individual sub-sequence is not now cleanly falls back to the unfused path instead of silently miscomputingglobal_start // cp_sizetoken indices.End-to-end (Qwen3.5-VL 397B proxy, GB200): fused vs. unfused are numerically equivalent in the forward and backward of every isolated component and of the full single-GPU multimodal model (grad-norm ratio 1.000, loss matches, ~+3–4% throughput).
Issue tracking
Linked issue:
Contribution process
Pre-checks
Merge
All PRs start as draft; mark Ready for Review once merge-conflicts are resolved and CI is passing. For PRs that change
megatron/core, expert + final reviewers are assigned automatically via.github/CODEOWNERS. Any member of mcore-engineers can merge.Summary by Sourcery
Integrate Triton-based fused multimodal RoPE (mRoPE) kernels and wire them into GPT and Qwen3.5-VL vision models while preserving unfused fallbacks and improving RoPE fusion configurability.
New Features:
Bug Fixes:
Enhancements:
Tests: