Lit/qwen35 benchmark

[wplf]

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @megatron-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @megatron-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

Summary by Sourcery

Integrate Triton-based fused multimodal RoPE (mRoPE) kernels into Megatron, wire them through GPT and Qwen3.5-VL vision paths, and extend scripts and configs to benchmark fused vs unfused RoPE while relaxing an overly-strict HybridEP safety check.

New Features:

• Add Triton fused multimodal RoPE kernels for both standard and THD layouts, with public helpers for availability checks and raw-frequency to embedding conversion.
• Enable GPT mRoPE to request and consume raw per-axis frequency tensors for fused application, including context-parallel and packed-sequence support.
• Update Qwen3.5-VL vision encoder to emit sectioned raw 2D mRoPE frequencies and optionally route RoPE through fused THD mRoPE with fp32 compute.
• Extend Qwen3.5-VL training script with additional knobs for RoPE fusion toggling, MTP, attention layout, tokenizer/dataset selection, profiling, and checkpoint control.

Bug Fixes:

• Fix context-parallel THD token index calculations for MLA+YaRN fused RoPE and correct CP slicing for THD frequency tensors with odd local lengths.
• Relax HybridEP RDMA Queue Pair depth validation to avoid false positives on valid workloads.
• Ensure apply_rope_fusion is disabled for non-RoPE position embedding types except mRoPE and keep packed-sequence mRoPE embeddings unsliced by context parallelism.

Enhancements:

• Improve RoPE fusion warning handling with categorized, one-time fallback messages across TE and Triton mRoPE paths.
• Tighten validation and capability checks in TransformerConfig and fused mRoPE utilities, including CUDA, dtype, stride, and TE/Triton availability.
• Refine Qwen3.5-VL vision configuration and rotary embedding plumbing to model its 2D RoPE as sectioned mRoPE and to propagate language-model RoPE fusion settings to the vision tower.

Tests:

• Add an extensive fused mRoPE unit test suite covering Triton kernels, THD layouts, context parallelism, fallbacks, and GPT integration.
• Add Qwen3.5-VL vision-specific tests to validate raw 2D mRoPE frequency construction, fused vs unfused RoPE equivalence, and proper dispatch to fused THD mRoPE.

[sourcery-ai]

Reviewer's Guide

Integrates a Triton-based fused multimodal RoPE (mRoPE) path into Megatron for both standard and THD-packed layouts, wires it through GPT and Qwen3.5-VL vision models with new configuration/CLI knobs and NVTX profiling support, and adds extensive unit tests while relaxing a HybridEP RDMA guard.

Sequence diagram for runtime selection of fused vs unfused mRoPE

sequenceDiagram
    participant Caller
    participant apply_rotary_pos_emb
    participant fused_mrope
    participant TE_fused as fused_apply_rotary_pos_emb
    participant Unfused as _apply_rotary_pos_emb_bshd

    Caller->>apply_rotary_pos_emb: apply_rotary_pos_emb(t, freqs, config,...)
    apply_rotary_pos_emb->>apply_rotary_pos_emb: is_raw_mrope_freqs
    alt raw mRoPE and cu_seqlens is None
        apply_rotary_pos_emb->>fused_mrope: get_fused_mrope_unavailable_reason
        alt fused available
            apply_rotary_pos_emb->>fused_mrope: fused_apply_mrope(t, freqs,...)
            fused_mrope-->>apply_rotary_pos_emb: rotated_t
        else fused unavailable or options unsupported
            apply_rotary_pos_emb->>fused_mrope: mrope_freqs_to_rotary_emb
            fused_mrope-->>apply_rotary_pos_emb: rotary_emb
            apply_rotary_pos_emb->>Unfused: _apply_rotary_pos_emb_bshd(t, rotary_emb,...)
            Unfused-->>apply_rotary_pos_emb: rotated_t
        end
    else TE RoPE fusion path
        apply_rotary_pos_emb->>TE_fused: fused_apply_rotary_pos_emb or fused_apply_rotary_pos_emb_thd
        TE_fused-->>apply_rotary_pos_emb: rotated_t
    end
    apply_rotary_pos_emb-->>Caller: rotated_t

Loading

Flow diagram for fused mRoPE integration across components

flowchart LR
    GPT[GPTModel._preprocess]
    MRE[MultimodalRotaryEmbedding.forward]
    RU[apply_rotary_pos_emb]
    FM[fused_mrope module]

    GPT -->|position_embedding_type=mrope\nuse_raw_mrope_freqs| MRE
    MRE -->|return_raw_freqs=True\npacked_seq flag| RU
    RU -->|is_raw_mrope_freqs| FM
    FM -->|fused_apply_mrope /\nfused_apply_mrope_thd| RU
    RU -->|rotated tensor| GPT

Loading

File-Level Changes

Change	Details	Files
Add Triton fused multimodal RoPE implementation and integrate it with existing RoPE utilities and Transformer configuration.	Introduce megatron.core.fusions.fused_mrope module implementing raw-frequency based fused mRoPE kernels for BSHD and THD layouts with autograd support and availability checks. Extend rope_utils with raw mRoPE detection, conversion, THD packing helpers, and dispatch logic that prefers fused mRoPE, falls back to TE fused RoPE, then to unfused kernels with de-duplicated warnings. Update TransformerConfig to recognize fused mRoPE availability (even without TE) and relax apply_rope_fusion validation accordingly. Adjust attention module to pass mrope_interleaved into MultimodalRotaryEmbedding construction so fused/conversion paths know the layout.	megatron/core/fusions/fused_mrope.py megatron/core/models/common/embeddings/rope_utils.py megatron/core/transformer/transformer_config.py megatron/core/transformer/attention.py
Plumb fused mRoPE support through GPT mRoPE usage, ensuring correct raw/materialized RoPE handling across training, eval, and inference modes.	Add a GPTModel flag caching fused mRoPE availability and use it to decide when to request raw mRoPE frequencies from MultimodalRotaryEmbedding. Modify GPTModel._preprocess mrope path to distinguish packed vs unpacked sequences, avoid raw freqs in inference and fused_single_qkv_rope cases, and set packed_seq and return_raw_freqs flags accordingly. Allow arguments validation to keep rope fusion enabled for both rope and mrope position embedding types.	megatron/core/models/gpt/gpt_model.py megatron/training/arguments.py
Update multimodal Qwen3.5-VL vision encoder and specs to use raw 2D mRoPE frequencies compatible with fused mRoPE THD, plus NVTX profiling hooks.	Change Qwen35VLVisionEncoder to optionally emit raw sectioned 3-axis vision RoPE frequencies when mrope_section is configured and consume them directly in the decoder instead of materialized embeddings. Set up the vision TransformerConfig to describe 2D RoPE as sectioned raw mRoPE (zero temporal section, non-interleaved) and validate per-head dimension splits. Refactor the _apply_rope_fp32 helper to delegate to the generic apply_rotary_pos_emb, adding a fast path that calls fused mRoPE THD when possible, and keep output casting behavior. Wrap the no-CP RoPE application for the vision encoder with NVTX push/pop ranges for profiling. Minor import cleanup in the vision encoder.	examples/multimodal_dev/models/qwen35_vl/vision_encoder.py examples/multimodal_dev/models/qwen35_vl/specs.py examples/multimodal_dev/models/qwen35_vl/configuration.py examples/multimodal_dev/models/base.py
Enhance Qwen3.5-VL training script with additional benchmarking, profiling, and configurability options, including RoPE fusion toggles.	Introduce environment-driven knobs for MTP layer count, linear attention frequency, dataset provider, tokenizer type (including NullTokenizer), RoPE fusion disable flag, checkpoint saving toggle, NVTX range emission, and torchrun Python executable. Wire new knobs into TRAINING_ARGS, PROFILE_ARGS, EVAL_AND_LOGGING_ARGS, TOKENIZER_ARGS, and MULTIMODAL_ARGS, conditionally appending flags like --mtp-num-layers, --nvtx-ranges, --no-rope-fusion, custom image seq length, and optional HF processor path. Refine experiment naming and console summary to reflect new options (MTP layers, linear attention freq, RoPE fusion state, dataset, tokenizer, checkpoint saving, NVTX ranges). Change torchrun launcher invocation to use the selected TORCHRUN_PYTHON -m torch.distributed.run API.	examples/multimodal_dev/scripts/run_qwen35_vl.sh examples/multimodal_dev/pretrain_multimodal.py
Adjust THD context-parallel RoPE indexing and MLA-YARN helper to support odd local sequence lengths and keep CP mappings consistent with fused mRoPE.	Update _get_thd_freqs_on_this_cp_rank in rope_utils to compute first/second CP segments with ceil/floor splits, mirroring the fused THD mRoPE mapping for odd-length local sequences. Add THD-specific helpers for raw mRoPE: per-rank slicing, CP splits from cu_seqlens, and packing raw freqs into THD token order for fallback implementations. Fix fused_mla_yarn_rope_apply THD token index calculation to use the same first/second segment logic as the updated RoPE CP mapping.	megatron/core/models/common/embeddings/rope_utils.py megatron/core/fusions/fused_mla_yarn_rope_apply.py
Relax a HybridEP RDMA queue depth guard that was producing false positives at high token counts.	Comment out the tx_depth >= 65536 ValueError in fused_a2a, documenting that it was a false-positive guard for HybridEP and should be re-enabled only if real RDMA QP failures are reproduced.	megatron/core/transformer/moe/fused_a2a.py
Extend MultimodalRotaryEmbedding to better support raw mRoPE outputs and packed-sequence behavior.	Add return_raw_freqs and packed_seq flags to forward; when return_raw_freqs is set, return contiguous raw per-axis frequencies with CP slicing only for non-packed sequences. Ensure packed_seq=True skips CP slicing (to preserve global positions for THD packing) in both raw and materialized modes, and fix reshape logic for interleaved layouts to use reshape instead of view with correct dims.	megatron/core/models/common/embeddings/rotary_pos_embedding.py
Add extensive unit tests validating fused mRoPE correctness, fallbacks, and Qwen3.5-VL vision RoPE integration.	Create tests for fused mRoPE kernels covering forward/backward parity with unfused RoPE, BSHD and THD layouts, context-parallel behavior, dtype/device/stride capability checks, and option-driven fallbacks with single-shot warnings. Add tests ensuring GPT mRoPE paths request raw vs materialized frequencies correctly under various fusion and inference configurations, and that TransformerConfig validation honors fused mRoPE availability. Add Qwen3.5-VL vision-specific tests for 2D RoPE configuration, raw-vs-legacy frequency equivalence, fused THD dispatch from the fp32 helper, and full fused-vs-unfused forward/backward parity on CUDA.	tests/unit_tests/fusions/test_fused_mrope.py examples/multimodal_dev/tests/test_vision_rope_fusion.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've found 1 issue, and left some high level feedback:

• In _apply_rope_fp32_no_cp you call nvtx_range_pop(range_name), but nvtx_range_pop is typically a no-arg function that pops the last range; passing the name may be ignored or error-prone, so consider dropping the argument and just calling nvtx_range_pop().
• The new validation helpers in fused_mrope.py (e.g. _validate_mrope_section, _validate_mrope_inputs, _validate_mrope_thd_inputs) rely heavily on assert, which will be stripped when Python is run with optimizations; consider converting these to explicit ValueError/RuntimeError checks so invalid shapes and sections still generate clear errors in production.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In `_apply_rope_fp32_no_cp` you call `nvtx_range_pop(range_name)`, but `nvtx_range_pop` is typically a no-arg function that pops the last range; passing the name may be ignored or error-prone, so consider dropping the argument and just calling `nvtx_range_pop()`.
- The new validation helpers in `fused_mrope.py` (e.g. `_validate_mrope_section`, `_validate_mrope_inputs`, `_validate_mrope_thd_inputs`) rely heavily on `assert`, which will be stripped when Python is run with optimizations; consider converting these to explicit `ValueError`/`RuntimeError` checks so invalid shapes and sections still generate clear errors in production.

## Individual Comments

### Comment 1
<location path="examples/multimodal_dev/tests/test_vision_rope_fusion.py" line_range="71" />
<code_context>
+    torch.testing.assert_close(converted, expected)
+
+
+def test_vision_fp32_wrapper_dispatches_raw_freqs_to_fused_mrope_thd(monkeypatch):
+    calls = {}
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add a complementary test for the fallback path when fused THD mRoPE is unavailable in the vision FP32 wrapper.

The current test only exercises the case where `fused_apply_mrope_thd` is available. Please also add a test where `get_fused_mrope_thd_unavailable_reason` returns a non-None value so the code takes the unfused `apply_rotary_pos_emb` path. That test should verify that `fused_apply_mrope_thd` is not called, no `fp32_fused` flag is propagated incorrectly, and the numerical output matches the unfused implementation, so both branches of the wrapper are covered.

Suggested implementation:

```python
def test_vision_fp32_wrapper_falls_back_to_unfused_when_fused_thd_unavailable(
    monkeypatch,
):
    # Import the wrapper and rotary helpers from the module under test.
    # The exact import path / names may need to be aligned with the rest of this file.
    from examples.multimodal_dev.vision_rope_fusion import (  # type: ignore[import]
        vision_fp32_wrapper,
        apply_rotary_pos_emb,
        fused_apply_mrope_thd,
        get_fused_mrope_thd_unavailable_reason,
    )

    calls = {"fused": 0, "unfused": 0}

    # Force the wrapper down the unfused path by reporting that the fused kernel
    # is unavailable.
    def fake_unavailable_reason(*args, **kwargs):
        return "no fused kernel"

    # If the fused path is taken in this test, fail loudly.
    def fake_fused_apply(*args, **kwargs):
        calls["fused"] += 1
        raise AssertionError("fused_apply_mrope_thd must not be called in fallback path")

    # Keep a handle to the real unfused implementation so we can both wrap it and
    # use it to compute the reference output.
    real_apply_rotary_pos_emb = apply_rotary_pos_emb

    def wrapped_apply_rotary_pos_emb(q, k, freqs, *args, fp32_fused=None, **kwargs):
        # The fallback path must not propagate an fp32_fused flag.
        assert fp32_fused is None
        calls["unfused"] += 1
        return real_apply_rotary_pos_emb(q, k, freqs, *args, **kwargs)

    # Patch the module-under-test symbols so the wrapper sees the fake behavior.
    monkeypatch.setattr(
        "examples.multimodal_dev.vision_rope_fusion.get_fused_mrope_thd_unavailable_reason",
        fake_unavailable_reason,
        raising=True,
    )
    monkeypatch.setattr(
        "examples.multimodal_dev.vision_rope_fusion.fused_apply_mrope_thd",
        fake_fused_apply,
        raising=True,
    )
    monkeypatch.setattr(
        "examples.multimodal_dev.vision_rope_fusion.apply_rotary_pos_emb",
        wrapped_apply_rotary_pos_emb,
        raising=True,
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Build small but non‑trivial test inputs; shapes can be adjusted to match
    # the rest of the vision tests if needed.
    q = torch.randn(2, 1, 4, 8, dtype=torch.float16, device=device)
    k = torch.randn_like(q)
    freqs = torch.randn(1, 1, 4, 4, dtype=torch.float16, device=device)

    # Reference output using the real unfused implementation.
    ref_q, ref_k = real_apply_rotary_pos_emb(q, k, freqs)

    # Call the FP32 wrapper, which should dispatch to the unfused path.
    out_q, out_k = vision_fp32_wrapper(q, k, freqs)

    # Verify dispatch behavior.
    assert calls["fused"] == 0
    assert calls["unfused"] == 1

    # And verify numerics match the unfused implementation.
    torch.testing.assert_close(out_q, ref_q)
    torch.testing.assert_close(out_k, ref_k)


"""Tests for Qwen3.5-VL vision RoPE fusion dispatch."""

```

1. Adjust the import in the new test to use the actual module and symbol names for:
   - The vision FP32 wrapper under test (`vision_fp32_wrapper` is a placeholder).
   - `apply_rotary_pos_emb`, `fused_apply_mrope_thd`, and `get_fused_mrope_thd_unavailable_reason`.
   If these are already imported at module scope in this test file, you can drop the local `from examples... import ...` and just use the existing names.
2. Update the `monkeypatch.setattr` target strings (`"examples.multimodal_dev.vision_rope_fusion.*"`) to match the real module path where the wrapper and helpers are defined.
3. If the wrapper’s calling convention differs (e.g., additional arguments, different tensor shapes, or named parameters for `raw_freqs`), adjust the construction of `q`, `k`, `freqs` and the `vision_fp32_wrapper` call accordingly, mirroring the existing `test_vision_fp32_wrapper_dispatches_raw_freqs_to_fused_mrope_thd` test.
4. If the unfused helper already accepts or requires an `fp32_fused` keyword, adapt the assertion in `wrapped_apply_rotary_pos_emb` so that it checks for the expected value in the fallback path (the key requirement is that the “fp32_fused” flag is not incorrectly forced to True when taking the unfused path).
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

suggestion (testing): Add a complementary test for the fallback path when fused THD mRoPE is unavailable in the vision FP32 wrapper.

The current test only exercises the case where fused_apply_mrope_thd is available. Please also add a test where get_fused_mrope_thd_unavailable_reason returns a non-None value so the code takes the unfused apply_rotary_pos_emb path. That test should verify that fused_apply_mrope_thd is not called, no fp32_fused flag is propagated incorrectly, and the numerical output matches the unfused implementation, so both branches of the wrapper are covered.

Suggested implementation:

def test_vision_fp32_wrapper_falls_back_to_unfused_when_fused_thd_unavailable(
    monkeypatch,
):
    # Import the wrapper and rotary helpers from the module under test.
    # The exact import path / names may need to be aligned with the rest of this file.
    from examples.multimodal_dev.vision_rope_fusion import (  # type: ignore[import]
        vision_fp32_wrapper,
        apply_rotary_pos_emb,
        fused_apply_mrope_thd,
        get_fused_mrope_thd_unavailable_reason,
    )

    calls = {"fused": 0, "unfused": 0}

    # Force the wrapper down the unfused path by reporting that the fused kernel
    # is unavailable.
    def fake_unavailable_reason(*args, **kwargs):
        return "no fused kernel"

    # If the fused path is taken in this test, fail loudly.
    def fake_fused_apply(*args, **kwargs):
        calls["fused"] += 1
        raise AssertionError("fused_apply_mrope_thd must not be called in fallback path")

    # Keep a handle to the real unfused implementation so we can both wrap it and
    # use it to compute the reference output.
    real_apply_rotary_pos_emb = apply_rotary_pos_emb

    def wrapped_apply_rotary_pos_emb(q, k, freqs, *args, fp32_fused=None, **kwargs):
        # The fallback path must not propagate an fp32_fused flag.
        assert fp32_fused is None
        calls["unfused"] += 1
        return real_apply_rotary_pos_emb(q, k, freqs, *args, **kwargs)

    # Patch the module-under-test symbols so the wrapper sees the fake behavior.
    monkeypatch.setattr(
        "examples.multimodal_dev.vision_rope_fusion.get_fused_mrope_thd_unavailable_reason",
        fake_unavailable_reason,
        raising=True,
    )
    monkeypatch.setattr(
        "examples.multimodal_dev.vision_rope_fusion.fused_apply_mrope_thd",
        fake_fused_apply,
        raising=True,
    )
    monkeypatch.setattr(
        "examples.multimodal_dev.vision_rope_fusion.apply_rotary_pos_emb",
        wrapped_apply_rotary_pos_emb,
        raising=True,
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Build small but non‑trivial test inputs; shapes can be adjusted to match
    # the rest of the vision tests if needed.
    q = torch.randn(2, 1, 4, 8, dtype=torch.float16, device=device)
    k = torch.randn_like(q)
    freqs = torch.randn(1, 1, 4, 4, dtype=torch.float16, device=device)

    # Reference output using the real unfused implementation.
    ref_q, ref_k = real_apply_rotary_pos_emb(q, k, freqs)

    # Call the FP32 wrapper, which should dispatch to the unfused path.
    out_q, out_k = vision_fp32_wrapper(q, k, freqs)

    # Verify dispatch behavior.
    assert calls["fused"] == 0
    assert calls["unfused"] == 1

    # And verify numerics match the unfused implementation.
    torch.testing.assert_close(out_q, ref_q)
    torch.testing.assert_close(out_k, ref_k)


"""Tests for Qwen3.5-VL vision RoPE fusion dispatch."""

1. Adjust the import in the new test to use the actual module and symbol names for:
   • The vision FP32 wrapper under test (vision_fp32_wrapper is a placeholder).
   • apply_rotary_pos_emb, fused_apply_mrope_thd, and get_fused_mrope_thd_unavailable_reason.
     If these are already imported at module scope in this test file, you can drop the local from examples... import ... and just use the existing names.
2. Update the monkeypatch.setattr target strings ("examples.multimodal_dev.vision_rope_fusion.*") to match the real module path where the wrapper and helpers are defined.
3. If the wrapper’s calling convention differs (e.g., additional arguments, different tensor shapes, or named parameters for raw_freqs), adjust the construction of q, k, freqs and the vision_fp32_wrapper call accordingly, mirroring the existing test_vision_fp32_wrapper_dispatches_raw_freqs_to_fused_mrope_thd test.
4. If the unfused helper already accepts or requires an fp32_fused keyword, adapt the assertion in wrapped_apply_rotary_pos_emb so that it checks for the expected value in the fallback path (the key requirement is that the “fp32_fused” flag is not incorrectly forced to True when taking the unfused path).