From 5eebb59aee56cbc23421201fb766defc29d407a4 Mon Sep 17 00:00:00 2001
From: Slawomir Kierat <skierat@nvidia.com>
Date: Tue, 14 Jul 2026 02:54:16 -0700
Subject: [PATCH 1/5] feat(speculative): add Qwen3-VL support for DFlash
 training

Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
---
 examples/speculative_decoding/eagle_utils.py  |   3 +
 .../torch/export/plugins/hf_spec_export.py    |  29 ++-
 .../torch/speculative/plugins/hf_dflash.py    | 202 ++++++++++++++++--
 .../torch/export/test_hf_spec_rope_export.py  |  13 ++
 .../speculative/plugins/test_hf_dflash.py     |  81 +++++++
 .../plugins/test_hf_speculative_offline.py    |  34 +++
 6 files changed, 339 insertions(+), 23 deletions(-)

diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py
index b12b9da1a52..68c6db45235 100644
--- a/examples/speculative_decoding/eagle_utils.py
+++ b/examples/speculative_decoding/eagle_utils.py
@@ -141,6 +141,9 @@ def make_speculative_data_module(
                 train_len=train_len,
                 local_image_path=data_args.vlm_img_dir,
                 return_labels=True,
+                answer_only_loss=answer_only_loss,
+                shift_labels=shift_labels,
+                chat_template=chat_template,
             )
 
     else:
diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py
index caa93db3634..eb80744d199 100644
--- a/modelopt/torch/export/plugins/hf_spec_export.py
+++ b/modelopt/torch/export/plugins/hf_spec_export.py
@@ -29,6 +29,23 @@
 
 ALL_SPEC_MODES = ["eagle", "dflash"]
 
+
+def _get_rope_theta(config, default=None):
+    """Get RoPE theta from either legacy or Transformers 5 config fields."""
+    rope_theta = getattr(config, "rope_theta", None)
+    if rope_theta is not None:
+        return rope_theta
+
+    # Transformers 5 stores this under rope_parameters (and exposes the same
+    # data through rope_scaling for backwards compatibility).
+    for attr in ("rope_parameters", "rope_scaling"):
+        rope_config = getattr(config, attr, None)
+        if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None:
+            return rope_config["rope_theta"]
+
+    return default
+
+
 LLAMA_EAGLE_SINGLE_LAYER = {
     "required": {
         "layers.0.self_attn.q_proj",
@@ -376,13 +393,11 @@ def _export_config(self):
             "initializer_range": getattr(base_config, "initializer_range", 0.02),
             "attention_bias": getattr(draft_config, "attention_bias", False),
             "attention_dropout": getattr(draft_config, "attention_dropout", 0.0),
-            # Inherit the target's rope_theta: DFlash injects the target's KV into every
-            # draft layer, so the draft's RoPE base must match the target's. (The draft
-            # arch config carries no rope_theta of its own.)
-            "rope_theta": (
-                getattr(base_config, "rope_theta", None)
-                if getattr(base_config, "rope_theta", None) is not None
-                else getattr(draft_config, "rope_theta", 1000000.0)
+            # Inherit the target's RoPE base: DFlash injects target KV into every draft
+            # layer, so their RoPE bases must match. Transformers 5 stores rope_theta
+            # in rope_parameters rather than a top-level config attribute.
+            "rope_theta": _get_rope_theta(
+                base_config, _get_rope_theta(draft_config, 1000000.0)
             ),
             # YaRN long-context scaling is injected below (see the rope_scaling block).
             "rope_scaling": None,
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 7679ff0020b..eb55642d9d6 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -74,6 +74,7 @@
 import logging
 
 import torch
+import transformers
 import torch.nn.functional as F
 from transformers import PreTrainedModel
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config
@@ -100,6 +101,30 @@
 __all__ = ["HFDFlashModel"]
 
 
+def _expand_qwen3_video_grid_thw(video_grid_thw: torch.Tensor) -> torch.Tensor:
+    """Return the per-frame video grid representation used by Qwen3-VL RoPE.
+
+    Qwen3-VL's video processor emits one ``[T, H, W]`` row per source video, but
+    its rendered prompt contains a separate visual-token group for every temporal
+    frame.  Transformers 5.3's ``get_rope_index`` consumes one grid row per
+    rendered group, while the vision encoder still requires the original one-row-
+    per-video representation.  This helper is therefore used *only* for mRoPE
+    position construction; callers must keep the original tensor for the model
+    forward.
+    """
+    if video_grid_thw.ndim != 2 or video_grid_thw.shape[-1] != 3:
+        raise ValueError(
+            "Qwen3-VL video_grid_thw must have shape [num_videos, 3], got "
+            f"{tuple(video_grid_thw.shape)}."
+        )
+    if torch.any(video_grid_thw[:, 0] <= 0):
+        raise ValueError("Qwen3-VL video_grid_thw temporal lengths must be positive.")
+
+    expanded_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+    expanded_grid_thw[:, 0] = 1
+    return expanded_grid_thw
+
+
 def _dpace_position_weights(
     confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None
 ) -> torch.Tensor:
@@ -183,6 +208,98 @@ def _base_llm_config(self):
             or self.config
         )
 
+    def _qwen3_vl_position_ids(
+        self,
+        input_ids,
+        attention_mask,
+        position_ids,
+        past_key_values,
+        inputs_embeds,
+        model_kwargs,
+    ):
+        """Precompute Qwen3-VL mRoPE positions for Transformers 5.3 batches.
+
+        The video encoder consumes one grid row per source video, whereas mRoPE
+        consumes one row per rendered temporal-frame group.  Calling the
+        top-level model with the original video grid makes the two contracts
+        conflict.  Construct the mRoPE positions with a frame-expanded copy,
+        then pass the original grid to the vision encoder in ``forward``.
+
+        Prefer ``get_rope_index`` over ``compute_3d_position_ids``.  The latter
+        writes ``rope_deltas`` into the base model even though DFlash training
+        never supplies a cache.  Keeping this calculation side-effect free is
+        important when the frozen target is reused for consecutive training
+        batches or validation.
+        """
+        if (
+            position_ids is not None
+            or getattr(self.config, "model_type", None) != "qwen3_vl"
+            or not transformers.__version__.startswith("5.3.")
+            # Cached decoding uses the base model's rope_deltas path.  DFlash
+            # training has no cache and is the only path that needs the
+            # frame-expanded construction below.
+            or past_key_values is not None
+        ):
+            return position_ids
+
+        image_grid_thw = model_kwargs.get("image_grid_thw")
+        video_grid_thw = model_kwargs.get("video_grid_thw")
+        mm_token_type_ids = model_kwargs.get("mm_token_type_ids")
+        backbone = getattr(self, "model", None)
+        get_rope_index = getattr(backbone, "get_rope_index", None)
+        compute_position_ids = getattr(backbone, "compute_3d_position_ids", None)
+        if not isinstance(image_grid_thw, torch.Tensor) and not isinstance(
+            video_grid_thw, torch.Tensor
+        ):
+            return position_ids
+        if (
+            not isinstance(mm_token_type_ids, torch.Tensor)
+            or input_ids is None
+            or (not callable(get_rope_index) and not callable(compute_position_ids))
+        ):
+            raise ValueError(
+                "Qwen3-VL DFlash training requires input_ids, mm_token_type_ids, and "
+                "a Qwen3-VL model with get_rope_index or compute_3d_position_ids. "
+                "Use the Qwen3-VL AutoProcessor without dropping mm_token_type_ids."
+            )
+
+        if mm_token_type_ids.shape != input_ids.shape:
+            raise ValueError(
+                "Qwen3-VL mm_token_type_ids must have the same shape as input_ids, got "
+                f"{tuple(mm_token_type_ids.shape)} and {tuple(input_ids.shape)}."
+            )
+
+        rope_video_grid_thw = video_grid_thw
+        if isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.numel() > 0:
+            rope_video_grid_thw = _expand_qwen3_video_grid_thw(video_grid_thw)
+
+        rope_kwargs = {
+            "input_ids": input_ids,
+            "image_grid_thw": image_grid_thw,
+            "video_grid_thw": rope_video_grid_thw,
+            "attention_mask": attention_mask,
+            "mm_token_type_ids": mm_token_type_ids,
+        }
+        if callable(get_rope_index):
+            position_ids, _ = get_rope_index(**rope_kwargs)
+        else:
+            position_ids = compute_position_ids(
+                **rope_kwargs,
+                inputs_embeds=inputs_embeds,
+                past_key_values=past_key_values,
+            )
+
+        expected_shape = (3, *input_ids.shape)
+        valid_position_ids = isinstance(position_ids, torch.Tensor) and (
+            tuple(position_ids.shape) == expected_shape
+        )
+        if not valid_position_ids:
+            raise RuntimeError(
+                "Qwen3-VL produced invalid mRoPE position ids: expected shape "
+                f"{expected_shape}, got {getattr(position_ids, 'shape', None)}."
+            )
+        return position_ids
+
     def _find_base_model_parts(self):
         """Locate base model submodules (backbone, embeddings, lm_head) by probing known paths.
 
@@ -592,6 +709,15 @@ def forward(
         - Label alignment: position k predicts token at anchor+k
         - Optional loss decay weighting
         """
+        position_ids = self._qwen3_vl_position_ids(
+            input_ids,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            inputs_embeds,
+            kwargs,
+        )
+
         if not self.training:
             if self.dflash_offline:
                 raise RuntimeError(
@@ -638,12 +764,49 @@ def forward(
             )
             target_hidden = base_outputs.target_hidden
         else:
-            # TODO: For co-training the base model, remove no_grad and eval() switch.
+            # Multimodal models need the top-level conditional-generation forward so their
+            # image/video features are inserted before the language model runs.  Keep the
+            # long-standing narrow call for text-only models.
+            multimodal_keys = {
+                "pixel_values",
+                "pixel_values_videos",
+                "image_grid_thw",
+                "video_grid_thw",
+                "image_sizes",
+                "images",
+                "videos",
+            }
+            use_top_level_forward = any(kwargs.get(key) is not None for key in multimodal_keys)
             with torch.no_grad():
-                raw_outputs = super().forward(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
+                if use_top_level_forward:
+                    base_forward_kwargs = dict(kwargs)
+                    # Training-only data keys are not accepted by Hugging Face model forwards.
+                    base_forward_kwargs.pop("assistant_masks", None)
+                    base_forward_kwargs.pop("loss_mask", None)
+                    raw_outputs = super().forward(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        position_ids=position_ids,
+                        past_key_values=past_key_values,
+                        inputs_embeds=inputs_embeds,
+                        use_cache=False,
+                        output_attentions=output_attentions,
+                        output_hidden_states=True,
+                        cache_position=cache_position,
+                        return_dict=True,
+                        **base_forward_kwargs,
+                    )
+                else:
+                    raw_outputs = super().forward(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        output_hidden_states=True,
+                    )
+
+            if not getattr(raw_outputs, "hidden_states", None):
+                raise RuntimeError(
+                    "The base model did not return hidden states required for DFlash training. "
+                    "Ensure its top-level multimodal forward supports output_hidden_states=True."
                 )
             offset = 1
             selected = [raw_outputs.hidden_states[lid + offset] for lid in self.target_layer_ids]
@@ -652,16 +815,15 @@ def forward(
                 target_hidden=target_hidden, logits=raw_outputs.logits
             )
 
-        # 2. Build loss mask.
-        # When labels are provided (answer_only_loss), they already encode both
-        # assistant masking and padding (-100 for both). When labels are not
-        # provided, fall back to attention_mask for padding only.
+        # 2. Build loss mask. Labels carry optional answer-only masking, but do
+        # not in general mark padded tokens with -100 (the VLM collator creates
+        # them from padded input_ids). Always intersect with attention_mask so
+        # anchor sampling and loss never include the padded tail.
+        loss_mask = torch.ones(bsz, seq_len, device=device)
         if labels is not None:
-            loss_mask = (labels != LabelSmoother.ignore_index).float()
-        elif attention_mask is not None:
-            loss_mask = attention_mask.float()
-        else:
-            loss_mask = torch.ones(bsz, seq_len, device=device)
+            loss_mask = loss_mask * (labels != LabelSmoother.ignore_index).float()
+        if attention_mask is not None:
+            loss_mask = loss_mask * attention_mask.float()
 
         # In offline training, assistant mask is dumped and passed as kwarg.
         if kwargs.get("loss_mask") is not None:
@@ -674,8 +836,16 @@ def forward(
         n_blocks = anchor_positions.shape[1]
 
         if n_blocks == 0 or not block_keep_mask.any():
-            # Zero loss that still flows through dflash_module for DDP gradient sync
-            dummy = self.dflash_module.fc.weight.sum() * 0.0
+            # Keep all trainable draft parameters in the graph so DDP can reduce a rank
+            # that receives an all-masked answer-only batch.
+            dummy = sum(
+                (
+                    parameter.reshape(-1)[0] * 0.0
+                    for parameter in self.dflash_module.parameters()
+                    if parameter.requires_grad
+                ),
+                torch.zeros((), device=device),
+            )
             return ModelOutput(loss=dummy, logits=base_outputs.logits, train_acc=[[0.0]])
 
         # 4. Build draft inputs
diff --git a/tests/unit/torch/export/test_hf_spec_rope_export.py b/tests/unit/torch/export/test_hf_spec_rope_export.py
index 720082bc617..fbeb218793e 100644
--- a/tests/unit/torch/export/test_hf_spec_rope_export.py
+++ b/tests/unit/torch/export/test_hf_spec_rope_export.py
@@ -139,3 +139,16 @@ def test_dflash_rope_theta_inherits_base():
     """rope_theta is inherited from the target/base config (draft drafts for the base)."""
     config = _make_dflash_exporter(base_rope_theta=5000000.0)._export_config()
     assert config["rope_theta"] == 5000000.0
+
+
+def test_dflash_rope_theta_inherits_base_rope_parameters():
+    """Transformers 5 stores the target RoPE base in rope_parameters."""
+    exporter = _make_dflash_exporter(base_rope_theta=None)
+    exporter.model.config.rope_parameters = {
+        "rope_type": "default",
+        "rope_theta": 5000000.0,
+    }
+
+    config = exporter._export_config()
+
+    assert config["rope_theta"] == 5000000.0
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index ef2eec2ad07..6b69bce20a8 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -35,6 +35,7 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.speculative as mtsp
+import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash
 from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG
 from modelopt.torch.speculative.plugins.hf_dflash import (
     DFlashAttention,
@@ -119,6 +120,86 @@ def test_convert_sets_mask_token_id(self):
         assert model.mask_token_id == 0
 
 
+def test_qwen3_vl_transformers_53_position_ids_expand_video_grid(monkeypatch):
+    """Only mRoPE receives a per-frame video grid on Transformers 5.3."""
+    original_grid = torch.tensor([[3, 4, 5], [2, 6, 7]])
+    expected_position_ids = torch.ones(3, 1, 12, dtype=torch.long)
+    get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1)))
+    compute_position_ids = MagicMock()
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl"),
+        model=SimpleNamespace(
+            get_rope_index=get_rope_index,
+            compute_3d_position_ids=compute_position_ids,
+        ),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0")
+
+    position_ids = HFDFlashModel._qwen3_vl_position_ids(
+        fake_model,
+        input_ids=torch.ones(1, 12, dtype=torch.long),
+        attention_mask=torch.ones(1, 12, dtype=torch.long),
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        model_kwargs={
+            "video_grid_thw": original_grid,
+            "mm_token_type_ids": torch.zeros(1, 12, dtype=torch.long),
+        },
+    )
+
+    assert position_ids is expected_position_ids
+    assert not compute_position_ids.called
+    assert torch.equal(original_grid, torch.tensor([[3, 4, 5], [2, 6, 7]]))
+    assert torch.equal(
+        get_rope_index.call_args.kwargs["video_grid_thw"],
+        torch.tensor([[1, 4, 5], [1, 4, 5], [1, 4, 5], [1, 6, 7], [1, 6, 7]]),
+    )
+
+
+def test_qwen3_vl_transformers_53_position_ids_require_mm_token_types(monkeypatch):
+    """Never silently fall back to one-dimensional positions for a visual batch."""
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl"),
+        model=SimpleNamespace(get_rope_index=MagicMock()),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0")
+
+    with pytest.raises(ValueError, match="mm_token_type_ids"):
+        HFDFlashModel._qwen3_vl_position_ids(
+            fake_model,
+            input_ids=torch.ones(1, 12, dtype=torch.long),
+            attention_mask=torch.ones(1, 12, dtype=torch.long),
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=None,
+            model_kwargs={"image_grid_thw": torch.tensor([[1, 4, 4]])},
+        )
+
+
+def test_qwen3_vl_transformers_53_position_ids_reject_bad_mm_token_shape(monkeypatch):
+    """Keep processor-produced modality ids aligned with the padded text sequence."""
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl"),
+        model=SimpleNamespace(get_rope_index=MagicMock()),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0")
+
+    with pytest.raises(ValueError, match="same shape as input_ids"):
+        HFDFlashModel._qwen3_vl_position_ids(
+            fake_model,
+            input_ids=torch.ones(1, 12, dtype=torch.long),
+            attention_mask=torch.ones(1, 12, dtype=torch.long),
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=None,
+            model_kwargs={
+                "image_grid_thw": torch.tensor([[1, 4, 4]]),
+                "mm_token_type_ids": torch.zeros(1, 11, dtype=torch.long),
+            },
+        )
+
+
 class TestDPaceWeights:
     """Test the D-PACE position-weighting objective (arXiv:2605.18810)."""
 
diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
index 2deefadd9e3..852a226d9f3 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
@@ -57,6 +57,40 @@
 make_speculative_data_module = _eagle_utils.make_speculative_data_module
 
 
+# ---------------------------------------------------------------------------
+# online VLM data-module wiring
+# ---------------------------------------------------------------------------
+
+
+def test_vlm_data_module_passes_dflash_label_mode(monkeypatch):
+    """VLM batches must use unshifted labels for DFlash and preserve OSL settings."""
+    data_args = argparse.Namespace(
+        mode="online",
+        data_path="unused.jsonl",
+        vlm_processor="dummy-vlm-processor",
+        vlm_img_dir="/images",
+        chat_template=None,
+    )
+    collator = MagicMock()
+    monkeypatch.setattr(_eagle_utils, "ShardedDataset", MagicMock())
+    monkeypatch.setattr(_eagle_utils, "VisionLanguageDataCollator", collator)
+
+    module = make_speculative_data_module(
+        MagicMock(), data_args, train_len=16, answer_only_loss=True, shift_labels=False
+    )
+
+    collator.assert_called_once_with(
+        processor="dummy-vlm-processor",
+        train_len=16,
+        local_image_path="/images",
+        return_labels=True,
+        answer_only_loss=True,
+        shift_labels=False,
+        chat_template=None,
+    )
+    assert module["data_collator"] is collator.return_value
+
+
 # ---------------------------------------------------------------------------
 # sample_size truncation tests
 # ---------------------------------------------------------------------------

From 301c7c19ce27570e0835734589775684e2ce45e4 Mon Sep 17 00:00:00 2001
From: Slawomir Kierat <skierat@nvidia.com>
Date: Tue, 28 Jul 2026 02:03:08 -0700
Subject: [PATCH 2/5] fix(speculative): harden Qwen3-VL DFlash online training

Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
---
 .../torch/speculative/plugins/hf_dflash.py    | 101 ++++++++----
 modelopt/torch/speculative/utils.py           |   8 +-
 .../utils/plugins/transformers_dataset.py     | 132 +++++++++++++++-
 .../speculative/plugins/test_fakebase.py      |  24 +++
 .../speculative/plugins/test_hf_dflash.py     | 144 +++++++++++++++++-
 .../plugins/test_hf_speculative_offline.py    |  21 ++-
 6 files changed, 393 insertions(+), 37 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index eb55642d9d6..6c95a8ea1d0 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -74,8 +74,8 @@
 import logging
 
 import torch
-import transformers
 import torch.nn.functional as F
+import transformers
 from transformers import PreTrainedModel
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config
 from transformers.trainer_pt_utils import LabelSmoother
@@ -101,6 +101,30 @@
 __all__ = ["HFDFlashModel"]
 
 
+_QWEN3_VL_MROPE_WORKAROUND_VERSION = "5.3.0"
+_MULTIMODAL_FORWARD_KWARGS = frozenset(
+    {
+        "pixel_values",
+        "pixel_values_videos",
+        "image_grid_thw",
+        "video_grid_thw",
+        "mm_token_type_ids",
+        "image_sizes",
+        "images",
+        "videos",
+    }
+)
+
+
+def _multimodal_forward_kwargs(model_kwargs: dict) -> dict:
+    """Return collator fields accepted by Hugging Face multimodal forwards."""
+    return {
+        name: value
+        for name, value in model_kwargs.items()
+        if name in _MULTIMODAL_FORWARD_KWARGS and value is not None
+    }
+
+
 def _expand_qwen3_video_grid_thw(video_grid_thw: torch.Tensor) -> torch.Tensor:
     """Return the per-frame video grid representation used by Qwen3-VL RoPE.
 
@@ -217,7 +241,7 @@ def _qwen3_vl_position_ids(
         inputs_embeds,
         model_kwargs,
     ):
-        """Precompute Qwen3-VL mRoPE positions for Transformers 5.3 batches.
+        """Precompute Qwen3-VL mRoPE positions for Transformers 5.3.0 batches.
 
         The video encoder consumes one grid row per source video, whereas mRoPE
         consumes one row per rendered temporal-frame group.  Calling the
@@ -225,16 +249,20 @@ def _qwen3_vl_position_ids(
         conflict.  Construct the mRoPE positions with a frame-expanded copy,
         then pass the original grid to the vision encoder in ``forward``.
 
-        Prefer ``get_rope_index`` over ``compute_3d_position_ids``.  The latter
+        Transformers 5.4.0 performs this frame expansion in ``get_rope_index``
+        itself; only 5.3.0 needs the external workaround. See
+        https://github.com/huggingface/transformers/blob/v5.4.0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py
+
+        Prefer ``get_rope_index`` over ``compute_3d_position_ids``. The latter
         writes ``rope_deltas`` into the base model even though DFlash training
         never supplies a cache.  Keeping this calculation side-effect free is
         important when the frozen target is reused for consecutive training
         batches or validation.
         """
+        model_type = str(getattr(self.config, "model_type", ""))
         if (
             position_ids is not None
-            or getattr(self.config, "model_type", None) != "qwen3_vl"
-            or not transformers.__version__.startswith("5.3.")
+            or not model_type.startswith("qwen3_vl")
             # Cached decoding uses the base model's rope_deltas path.  DFlash
             # training has no cache and is the only path that needs the
             # frame-expanded construction below.
@@ -244,14 +272,24 @@ def _qwen3_vl_position_ids(
 
         image_grid_thw = model_kwargs.get("image_grid_thw")
         video_grid_thw = model_kwargs.get("video_grid_thw")
-        mm_token_type_ids = model_kwargs.get("mm_token_type_ids")
-        backbone = getattr(self, "model", None)
-        get_rope_index = getattr(backbone, "get_rope_index", None)
-        compute_position_ids = getattr(backbone, "compute_3d_position_ids", None)
         if not isinstance(image_grid_thw, torch.Tensor) and not isinstance(
             video_grid_thw, torch.Tensor
         ):
             return position_ids
+
+        if transformers.__version__ != _QWEN3_VL_MROPE_WORKAROUND_VERSION:
+            if transformers.__version__.startswith("5.3."):
+                raise RuntimeError(
+                    "Qwen3-VL DFlash mRoPE supports Transformers 5.3.0 or >=5.4.0; "
+                    f"got {transformers.__version__}. A 5.3.x patch release may already "
+                    "expand video_grid_thw internally."
+                )
+            return position_ids
+
+        mm_token_type_ids = model_kwargs.get("mm_token_type_ids")
+        backbone = getattr(self, "model", None)
+        get_rope_index = getattr(backbone, "get_rope_index", None)
+        compute_position_ids = getattr(backbone, "compute_3d_position_ids", None)
         if (
             not isinstance(mm_token_type_ids, torch.Tensor)
             or input_ids is None
@@ -271,6 +309,18 @@ def _qwen3_vl_position_ids(
 
         rope_video_grid_thw = video_grid_thw
         if isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.numel() > 0:
+            video_token_mask = mm_token_type_ids == 2
+            if isinstance(attention_mask, torch.Tensor):
+                video_token_mask = video_token_mask & attention_mask.bool()
+            video_group_starts = video_token_mask.clone()
+            video_group_starts[:, 1:] &= ~video_token_mask[:, :-1]
+            expected_video_groups = int(video_grid_thw[:, 0].sum())
+            actual_video_groups = int(video_group_starts.sum())
+            if actual_video_groups != expected_video_groups:
+                raise ValueError(
+                    "Qwen3-VL video frame groups do not match video_grid_thw: "
+                    f"expected {expected_video_groups}, found {actual_video_groups}."
+                )
             rope_video_grid_thw = _expand_qwen3_video_grid_thw(video_grid_thw)
 
         rope_kwargs = {
@@ -709,14 +759,15 @@ def forward(
         - Label alignment: position k predicts token at anchor+k
         - Optional loss decay weighting
         """
-        position_ids = self._qwen3_vl_position_ids(
-            input_ids,
-            attention_mask,
-            position_ids,
-            past_key_values,
-            inputs_embeds,
-            kwargs,
-        )
+        if self.training:
+            position_ids = self._qwen3_vl_position_ids(
+                input_ids,
+                attention_mask,
+                position_ids,
+                past_key_values,
+                inputs_embeds,
+                kwargs,
+            )
 
         if not self.training:
             if self.dflash_offline:
@@ -767,22 +818,10 @@ def forward(
             # Multimodal models need the top-level conditional-generation forward so their
             # image/video features are inserted before the language model runs.  Keep the
             # long-standing narrow call for text-only models.
-            multimodal_keys = {
-                "pixel_values",
-                "pixel_values_videos",
-                "image_grid_thw",
-                "video_grid_thw",
-                "image_sizes",
-                "images",
-                "videos",
-            }
-            use_top_level_forward = any(kwargs.get(key) is not None for key in multimodal_keys)
+            base_forward_kwargs = _multimodal_forward_kwargs(kwargs)
+            use_top_level_forward = bool(base_forward_kwargs)
             with torch.no_grad():
                 if use_top_level_forward:
-                    base_forward_kwargs = dict(kwargs)
-                    # Training-only data keys are not accepted by Hugging Face model forwards.
-                    base_forward_kwargs.pop("assistant_masks", None)
-                    base_forward_kwargs.pop("loss_mask", None)
                     raw_outputs = super().forward(
                         input_ids=input_ids,
                         attention_mask=attention_mask,
diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index 6a3c19b993d..8c5418bb6b8 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -610,7 +610,13 @@ def load_vlm_or_llm(
         return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code)
 
     if _is_vlm:
-        model_cls = transformers.AutoModelForVision2Seq
+        # Transformers 5 renamed AutoModelForVision2Seq to
+        # AutoModelForImageTextToText.  Prefer the pre-5 name so this loader
+        # continues to support the Transformers 4 environments used by older
+        # speculative-decoding jobs.
+        model_cls = getattr(transformers, "AutoModelForVision2Seq", None)
+        if model_cls is None:
+            model_cls = transformers.AutoModelForImageTextToText
     else:
         model_cls = transformers.AutoModelForCausalLM
 
diff --git a/modelopt/torch/utils/plugins/transformers_dataset.py b/modelopt/torch/utils/plugins/transformers_dataset.py
index c27a3d09aea..97ae4ea2d14 100644
--- a/modelopt/torch/utils/plugins/transformers_dataset.py
+++ b/modelopt/torch/utils/plugins/transformers_dataset.py
@@ -325,6 +325,7 @@ def __init__(
         chat_template: str | None = None,
         add_generation_prompt: bool = False,
         answer_only_loss: bool = False,
+        shift_labels: bool = True,
         local_image_path: str = "",
         return_labels: bool = False,
     ):
@@ -340,10 +341,97 @@ def __init__(
             chat_template=chat_template,
             add_generation_prompt=add_generation_prompt,
             answer_only_loss=answer_only_loss,
+            shift_labels=shift_labels,
             return_labels=return_labels,
         )
 
+    def _verify_generation_tags(self):
+        """Accept VLM templates whose assistant spans have stable chat markers.
+
+        Cosmos/Qwen ChatML templates do not necessarily use Hugging Face's
+        ``{% generation %}`` tags.  For those templates we derive the same
+        assistant-only loss mask from the tokenized assistant boundaries.
+        """
+        if self._assistant_marker_specs():
+            return
+        super()._verify_generation_tags()
+
+    def _assistant_marker_specs(self):
+        """Return tokenized assistant start/end boundaries for supported templates."""
+        if hasattr(self, "_cached_assistant_marker_specs"):
+            return self._cached_assistant_marker_specs
+
+        template = self.tokenizer.chat_template or ""
+        specs = []
+        if "<|im_start|>" in template and "<|im_end|>" in template:
+            specs.append(
+                (
+                    self.tokenizer("<|im_start|>assistant\n", add_special_tokens=False)[
+                        "input_ids"
+                    ],
+                    [
+                        self.tokenizer("<|im_end|>\n", add_special_tokens=False)["input_ids"],
+                        self.tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"],
+                    ],
+                )
+            )
+        self._cached_assistant_marker_specs = [
+            (start, [end for end in ends if end]) for start, ends in specs if start and any(ends)
+        ]
+        return self._cached_assistant_marker_specs
+
+    @staticmethod
+    def _find_subsequence(values, pattern, start=0, stop=None):
+        stop = len(values) if stop is None else stop
+        if not pattern or start >= stop:
+            return -1
+        for index in range(start, stop - len(pattern) + 1):
+            if values[index : index + len(pattern)] == pattern:
+                return index
+        return -1
+
+    def _build_assistant_masks(self, tokenized_messages):
+        """Build assistant-content masks from ChatML boundaries."""
+        input_ids = tokenized_messages["input_ids"]
+        attention_mask = tokenized_messages.get("attention_mask")
+        assistant_masks = torch.zeros_like(input_ids)
+
+        for row_index, row in enumerate(input_ids):
+            tokens = row.tolist()
+            if isinstance(attention_mask, torch.Tensor):
+                active = attention_mask[row_index].nonzero(as_tuple=False).flatten()
+                if active.numel() == 0:
+                    continue
+                sequence_start, sequence_end = int(active[0]), int(active[-1]) + 1
+            else:
+                sequence_start, sequence_end = 0, len(tokens)
+
+            for start_marker, end_markers in self._assistant_marker_specs():
+                search_from = sequence_start
+                while search_from < sequence_end:
+                    start = self._find_subsequence(tokens, start_marker, search_from, sequence_end)
+                    if start == -1:
+                        break
+                    content_start = start + len(start_marker)
+                    end_positions = [
+                        position
+                        for marker in end_markers
+                        if (
+                            position := self._find_subsequence(
+                                tokens, marker, content_start, sequence_end
+                            )
+                        )
+                        != -1
+                    ]
+                    content_end = min(end_positions) if end_positions else sequence_end
+                    if content_start < content_end:
+                        assistant_masks[row_index, content_start:content_end] = 1
+                    search_from = max(content_start + 1, content_end + 1)
+
+        return assistant_masks
+
     def _process_multimodal_sample(self, examples):
+        derive_masks_from_markers = self.answer_only_loss and bool(self._assistant_marker_specs())
         tokenized_messages = self.processor.apply_chat_template(
             examples,
             tokenize=True,
@@ -353,9 +441,36 @@ def _process_multimodal_sample(self, examples):
             truncation=True,
             max_length=self.train_len,
             add_generation_prompt=self.add_generation_prompt,
-            return_assistant_tokens_mask=self.answer_only_loss,
+            return_assistant_tokens_mask=self.answer_only_loss and not derive_masks_from_markers,
         )
 
+        if derive_masks_from_markers:
+            tokenized_messages["assistant_masks"] = self._build_assistant_masks(tokenized_messages)
+
+        if self.return_labels:
+            input_ids = tokenized_messages["input_ids"]
+            labels = input_ids.new_full(input_ids.shape, IGNORE_TOKEN_ID)
+            if self.shift_labels:
+                labels[..., :-1] = input_ids[..., 1:]
+            else:
+                # DFlash predicts the token at the current position rather
+                # than the next autoregressive token.
+                labels[:] = input_ids
+
+            if self.answer_only_loss:
+                if "assistant_masks" not in tokenized_messages:
+                    raise ValueError(
+                        "answer_only_loss requires assistant_masks from the VLM chat template."
+                    )
+                assistant_mask = tokenized_messages["assistant_masks"]
+                if not isinstance(assistant_mask, torch.Tensor) or not assistant_mask.any():
+                    labels[:] = IGNORE_TOKEN_ID
+                elif self.shift_labels:
+                    labels[..., :-1][assistant_mask[..., 1:] == 0] = IGNORE_TOKEN_ID
+                else:
+                    labels[assistant_mask == 0] = IGNORE_TOKEN_ID
+            tokenized_messages["labels"] = labels
+
         return tokenized_messages
 
     def __call__(self, examples):
@@ -385,6 +500,21 @@ def __call__(self, examples):
                     msg["content"] = [{"type": "text", "text": msg["content"]}]
 
                 for ctn in msg["content"]:
+                    # Some JSONL producers use a fixed multimodal-part schema
+                    # (text/image/video/fps on every part) so Arrow can load
+                    # heterogeneous image and video datasets together.  Drop
+                    # the inactive placeholders before handing a part to the
+                    # processor, which expects only fields relevant to its type.
+                    content_type = ctn.get("type")
+                    if content_type != "text" and ctn.get("text") == "":
+                        del ctn["text"]
+                    if content_type != "image" and ctn.get("image") == "":
+                        del ctn["image"]
+                    if content_type != "video":
+                        if ctn.get("video") == "":
+                            del ctn["video"]
+                        if ctn.get("fps") == 0:
+                            del ctn["fps"]
                     if ctn["type"] == "image" and "image" in ctn:
                         ctn["image"] = os.path.abspath(
                             os.path.join(self.local_image_path, ctn["image"])
diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py
index 2880bf4ef1c..52f0dba775d 100644
--- a/tests/unit/torch/speculative/plugins/test_fakebase.py
+++ b/tests/unit/torch/speculative/plugins/test_fakebase.py
@@ -134,3 +134,27 @@ def _fake_from_pretrained(*args, **kwargs):
     model = load_vlm_or_llm("fake-model", use_offline_training=True, use_fake_base=False)
     assert captured_kwargs.get("num_hidden_layers") == 0
     assert model.config.num_orig_hidden_layers == 4
+
+
+def test_load_vlm_or_llm_uses_transformers5_vlm_auto_class(monkeypatch):
+    """Transformers 5 loads VLMs through AutoModelForImageTextToText."""
+    cfg = transformers.PretrainedConfig()
+    cfg.model_type = "qwen3_vl"
+    cfg.text_config = object()
+    monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **kw: cfg)
+
+    captured = {}
+
+    class _FakeVLM:
+        @staticmethod
+        def from_pretrained(*args, **kwargs):
+            captured["args"] = args
+            captured["kwargs"] = kwargs
+            return object()
+
+    monkeypatch.delattr(transformers, "AutoModelForVision2Seq", raising=False)
+    monkeypatch.setattr(transformers, "AutoModelForImageTextToText", _FakeVLM)
+
+    assert load_vlm_or_llm("qwen3-vl", dtype="auto") is not None
+    assert captured["args"] == ("qwen3-vl",)
+    assert captured["kwargs"]["torch_dtype"] == "auto"
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 6b69bce20a8..959914ba807 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -120,8 +120,8 @@ def test_convert_sets_mask_token_id(self):
         assert model.mask_token_id == 0
 
 
-def test_qwen3_vl_transformers_53_position_ids_expand_video_grid(monkeypatch):
-    """Only mRoPE receives a per-frame video grid on Transformers 5.3."""
+def test_qwen3_vl_transformers_530_position_ids_expand_video_grid(monkeypatch):
+    """Only mRoPE receives a per-frame video grid on Transformers 5.3.0."""
     original_grid = torch.tensor([[3, 4, 5], [2, 6, 7]])
     expected_position_ids = torch.ones(3, 1, 12, dtype=torch.long)
     get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1)))
@@ -144,7 +144,9 @@ def test_qwen3_vl_transformers_53_position_ids_expand_video_grid(monkeypatch):
         inputs_embeds=None,
         model_kwargs={
             "video_grid_thw": original_grid,
-            "mm_token_type_ids": torch.zeros(1, 12, dtype=torch.long),
+            "mm_token_type_ids": torch.tensor(
+                [[2, 0, 2, 0, 2, 0, 2, 0, 2, 0, 0, 0]]
+            ),
         },
     )
 
@@ -157,6 +159,142 @@ def test_qwen3_vl_transformers_53_position_ids_expand_video_grid(monkeypatch):
     )
 
 
+def test_qwen3_vl_moe_transformers_530_position_ids_expand_video_grid(monkeypatch):
+    """Qwen3-VL family variants use the same 5.3.0 mRoPE workaround."""
+    expected_position_ids = torch.ones(3, 1, 4, dtype=torch.long)
+    get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1)))
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl_moe"),
+        model=SimpleNamespace(get_rope_index=get_rope_index),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0")
+
+    position_ids = HFDFlashModel._qwen3_vl_position_ids(
+        fake_model,
+        input_ids=torch.ones(1, 4, dtype=torch.long),
+        attention_mask=torch.ones(1, 4, dtype=torch.long),
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        model_kwargs={
+            "video_grid_thw": torch.tensor([[2, 4, 4]]),
+            "mm_token_type_ids": torch.tensor([[2, 0, 2, 0]]),
+        },
+    )
+
+    assert position_ids is expected_position_ids
+    assert torch.equal(
+        get_rope_index.call_args.kwargs["video_grid_thw"],
+        torch.tensor([[1, 4, 4], [1, 4, 4]]),
+    )
+
+
+def test_qwen3_vl_transformers_53_patch_release_raises(monkeypatch):
+    """Avoid double expansion when a 5.3 patch backports the upstream fix."""
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl"),
+        model=SimpleNamespace(get_rope_index=MagicMock()),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.1")
+
+    with pytest.raises(RuntimeError, match="5.3.0 or >=5.4.0"):
+        HFDFlashModel._qwen3_vl_position_ids(
+            fake_model,
+            input_ids=torch.ones(1, 4, dtype=torch.long),
+            attention_mask=torch.ones(1, 4, dtype=torch.long),
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=None,
+            model_kwargs={
+                "video_grid_thw": torch.tensor([[1, 4, 4]]),
+                "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]),
+            },
+        )
+
+
+def test_qwen3_vl_transformers_54_uses_native_position_ids(monkeypatch):
+    """Transformers 5.4+ performs the grid expansion inside get_rope_index."""
+    get_rope_index = MagicMock()
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl"),
+        model=SimpleNamespace(get_rope_index=get_rope_index),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.4.0")
+
+    position_ids = HFDFlashModel._qwen3_vl_position_ids(
+        fake_model,
+        input_ids=torch.ones(1, 4, dtype=torch.long),
+        attention_mask=torch.ones(1, 4, dtype=torch.long),
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        model_kwargs={
+            "video_grid_thw": torch.tensor([[1, 4, 4]]),
+            "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]),
+        },
+    )
+
+    assert position_ids is None
+    assert not get_rope_index.called
+
+
+def test_qwen3_vl_transformers_530_rejects_bad_video_frame_groups(monkeypatch):
+    """Fail before mRoPE construction when processor and video-grid contracts differ."""
+    fake_model = SimpleNamespace(
+        config=SimpleNamespace(model_type="qwen3_vl"),
+        model=SimpleNamespace(get_rope_index=MagicMock()),
+    )
+    monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0")
+
+    with pytest.raises(ValueError, match="video frame groups"):
+        HFDFlashModel._qwen3_vl_position_ids(
+            fake_model,
+            input_ids=torch.ones(1, 4, dtype=torch.long),
+            attention_mask=torch.ones(1, 4, dtype=torch.long),
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=None,
+            model_kwargs={
+                "video_grid_thw": torch.tensor([[2, 4, 4]]),
+                "mm_token_type_ids": torch.tensor([[2, 2, 0, 0]]),
+            },
+        )
+
+
+def test_multimodal_forward_kwargs_exclude_non_model_inputs():
+    """Do not forward Trainer or collator-only fields to Hugging Face models."""
+    pixel_values = torch.ones(1)
+    mm_token_type_ids = torch.zeros(1, 4, dtype=torch.long)
+
+    forwarded = hf_dflash._multimodal_forward_kwargs(
+        {
+            "pixel_values": pixel_values,
+            "mm_token_type_ids": mm_token_type_ids,
+            "assistant_masks": torch.ones(1, 4),
+            "loss_mask": torch.ones(1, 4),
+            "num_items_in_batch": 4,
+            "unexpected_dataset_column": "drop me",
+        }
+    )
+
+    assert set(forwarded) == {"pixel_values", "mm_token_type_ids"}
+    assert forwarded["pixel_values"] is pixel_values
+    assert forwarded["mm_token_type_ids"] is mm_token_type_ids
+
+
+def test_eval_does_not_precompute_qwen3_vl_position_ids(monkeypatch):
+    """Evaluation delegates mRoPE construction to the base model and its cache."""
+    model = get_tiny_llama(num_hidden_layers=4)
+    mtsp.convert(model, [("dflash", _get_dflash_config())])
+    precompute_position_ids = MagicMock()
+    monkeypatch.setattr(model, "_qwen3_vl_position_ids", precompute_position_ids)
+
+    model.eval()
+    model(input_ids=torch.tensor([[1, 2, 3, 4]]))
+
+    precompute_position_ids.assert_not_called()
+
+
 def test_qwen3_vl_transformers_53_position_ids_require_mm_token_types(monkeypatch):
     """Never silently fall back to one-dimensional positions for a visual batch."""
     fake_model = SimpleNamespace(
diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
index 852a226d9f3..5abaa124d68 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
@@ -30,7 +30,7 @@
 
 import pytest
 import torch
-from _test_utils.torch.transformers_models import get_tiny_llama
+from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer
 
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.speculative.eagle.default_config import default_eagle_config
@@ -38,6 +38,7 @@
     EagleOfflineDataCollator,
     OfflineSupervisedDataset,
 )
+from modelopt.torch.utils.plugins import transformers_dataset
 
 _mock_scripts = types.ModuleType("scripts")
 _mock_ar = types.ModuleType("scripts.ar_validate")
@@ -91,6 +92,24 @@ def test_vlm_data_module_passes_dflash_label_mode(monkeypatch):
     assert module["data_collator"] is collator.return_value
 
 
+def test_vlm_data_collator_accepts_unshifted_labels(monkeypatch):
+    """The real VLM collator must support DFlash's unshifted labels."""
+    processor = types.SimpleNamespace(tokenizer=get_tiny_tokenizer())
+    monkeypatch.setattr(
+        transformers_dataset.transformers.AutoProcessor,
+        "from_pretrained",
+        lambda *_args, **_kwargs: processor,
+    )
+
+    collator = transformers_dataset.VisionLanguageDataCollator(
+        processor="dummy-vlm-processor",
+        chat_template="{{ messages }}",
+        shift_labels=False,
+    )
+
+    assert collator.shift_labels is False
+
+
 # ---------------------------------------------------------------------------
 # sample_size truncation tests
 # ---------------------------------------------------------------------------

From b5aa42fd2da303f6205b878fffc6008dc4b6ceba Mon Sep 17 00:00:00 2001
From: Slawomir Kierat <skierat@nvidia.com>
Date: Tue, 28 Jul 2026 03:05:21 -0700
Subject: [PATCH 3/5] Quick fix

Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
---
 tests/unit/torch/speculative/plugins/test_fakebase.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py
index 52f0dba775d..5a3eb88f817 100644
--- a/tests/unit/torch/speculative/plugins/test_fakebase.py
+++ b/tests/unit/torch/speculative/plugins/test_fakebase.py
@@ -152,7 +152,11 @@ def from_pretrained(*args, **kwargs):
             captured["kwargs"] = kwargs
             return object()
 
-    monkeypatch.delattr(transformers, "AutoModelForVision2Seq", raising=False)
+    # ``transformers`` exposes auto classes lazily, so deleting this attribute
+    # lets its module-level ``__getattr__`` recreate the legacy class.  An
+    # explicit ``None`` models its absence and reliably exercises the v5
+    # fallback.
+    monkeypatch.setattr(transformers, "AutoModelForVision2Seq", None)
     monkeypatch.setattr(transformers, "AutoModelForImageTextToText", _FakeVLM)
 
     assert load_vlm_or_llm("qwen3-vl", dtype="auto") is not None

From c35d49019e5d6bf37c5c7d55fc88046cfa5eefac Mon Sep 17 00:00:00 2001
From: Slawomir Kierat <skierat@nvidia.com>
Date: Tue, 28 Jul 2026 05:03:24 -0700
Subject: [PATCH 4/5] Quick fix

Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
---
 tests/unit/torch/speculative/plugins/test_fakebase.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py
index 5a3eb88f817..cf6dfe1a6bc 100644
--- a/tests/unit/torch/speculative/plugins/test_fakebase.py
+++ b/tests/unit/torch/speculative/plugins/test_fakebase.py
@@ -156,7 +156,7 @@ def from_pretrained(*args, **kwargs):
     # lets its module-level ``__getattr__`` recreate the legacy class.  An
     # explicit ``None`` models its absence and reliably exercises the v5
     # fallback.
-    monkeypatch.setattr(transformers, "AutoModelForVision2Seq", None)
+    monkeypatch.setattr(transformers, "AutoModelForVision2Seq", None, raising=False)
     monkeypatch.setattr(transformers, "AutoModelForImageTextToText", _FakeVLM)
 
     assert load_vlm_or_llm("qwen3-vl", dtype="auto") is not None

From 19c23eb1a01e801c935f29bd886f3163cce05ffb Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Wed, 29 Jul 2026 07:51:22 +0000
Subject: [PATCH 5/5] fix: resolve ruff and mypy failures in Qwen3-VL DFlash

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 modelopt/torch/export/plugins/hf_spec_export.py        | 4 +---
 modelopt/torch/speculative/plugins/hf_dflash.py        | 6 ++++--
 tests/unit/torch/speculative/plugins/test_hf_dflash.py | 6 ++----
 3 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py
index eb80744d199..255b1d9ab04 100644
--- a/modelopt/torch/export/plugins/hf_spec_export.py
+++ b/modelopt/torch/export/plugins/hf_spec_export.py
@@ -396,9 +396,7 @@ def _export_config(self):
             # Inherit the target's RoPE base: DFlash injects target KV into every draft
             # layer, so their RoPE bases must match. Transformers 5 stores rope_theta
             # in rope_parameters rather than a top-level config attribute.
-            "rope_theta": _get_rope_theta(
-                base_config, _get_rope_theta(draft_config, 1000000.0)
-            ),
+            "rope_theta": _get_rope_theta(base_config, _get_rope_theta(draft_config, 1000000.0)),
             # YaRN long-context scaling is injected below (see the rope_scaling block).
             "rope_scaling": None,
             "tie_word_embeddings": False,
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 6c95a8ea1d0..e0d63bde136 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -72,6 +72,7 @@
 """
 
 import logging
+from typing import Any
 
 import torch
 import torch.nn.functional as F
@@ -288,8 +289,9 @@ def _qwen3_vl_position_ids(
 
         mm_token_type_ids = model_kwargs.get("mm_token_type_ids")
         backbone = getattr(self, "model", None)
-        get_rope_index = getattr(backbone, "get_rope_index", None)
-        compute_position_ids = getattr(backbone, "compute_3d_position_ids", None)
+        # Probed dynamically: which one exists depends on the Transformers version.
+        get_rope_index: Any = getattr(backbone, "get_rope_index", None)
+        compute_position_ids: Any = getattr(backbone, "compute_3d_position_ids", None)
         if (
             not isinstance(mm_token_type_ids, torch.Tensor)
             or input_ids is None
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 959914ba807..bd243421d2c 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -144,9 +144,7 @@ def test_qwen3_vl_transformers_530_position_ids_expand_video_grid(monkeypatch):
         inputs_embeds=None,
         model_kwargs={
             "video_grid_thw": original_grid,
-            "mm_token_type_ids": torch.tensor(
-                [[2, 0, 2, 0, 2, 0, 2, 0, 2, 0, 0, 0]]
-            ),
+            "mm_token_type_ids": torch.tensor([[2, 0, 2, 0, 2, 0, 2, 0, 2, 0, 0, 0]]),
         },
     )
 
@@ -197,7 +195,7 @@ def test_qwen3_vl_transformers_53_patch_release_raises(monkeypatch):
     )
     monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.1")
 
-    with pytest.raises(RuntimeError, match="5.3.0 or >=5.4.0"):
+    with pytest.raises(RuntimeError, match=r"5\.3\.0 or >=5\.4\.0"):
         HFDFlashModel._qwen3_vl_position_ids(
             fake_model,
             input_ids=torch.ones(1, 4, dtype=torch.long),