Fix Qwen3-VL THD packed mRoPE positions

[maocheng23]

Summary

• move the existing Qwen3-VL Bridge rotary compatibility monkey patch into a small helper
• synthesize packed [3, 1, total_T] mRoPE position_ids for Miles THD batches by splitting on packed_seq_params.cu_seqlens_q and calling Qwen3-VL get_rope_index per original sequence
• add focused tests for per-segment packed mRoPE construction, local-length mismatch skipping, and install-time patching

Scope

This is independent from #1244. It is based directly on main and only contains the Qwen3-VL THD mRoPE fix.

Tests

• uvx ruff check miles/backends/megatron_utils/qwen_vl_packed_mrope.py tests/fast/backends/megatron_utils/test_qwen_vl_packed_mrope.py
• python -m py_compile miles/backends/megatron_utils/qwen_vl_packed_mrope.py tests/fast/backends/megatron_utils/test_qwen_vl_packed_mrope.py
• PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q -c /dev/null -p no:cacheprovider --confcutdir=tests/fast/backends/megatron_utils tests/fast/backends/megatron_utils/test_qwen_vl_packed_mrope.py
• uvx pre-commit run --all-files --show-diff-on-failure --color=never

[gemini-code-assist]

Code Review

This pull request refactors the Qwen3-VL rotary embedding patch into a dedicated module (qwen_vl_packed_mrope.py) to support Miles THD packed batches by dynamically building packed mRoPE position IDs and rope deltas. It also adds comprehensive unit tests. The review feedback highlights two critical improvement opportunities: first, handling positional arguments in the patched model forward to prevent silent failures or TypeError exceptions, and second, optimizing performance by copying input tensors to the CPU once to avoid costly GPU-to-CPU synchronizations inside the segment loop.

I am having trouble creating individual review comments. Click here to see my feedback.

miles/backends/megatron_utils/qwen_vl_packed_mrope.py (77-100)

The current implementation of _patch_model_forward extracts arguments like input_ids and packed_seq_params solely from kwargs. In PyTorch/HuggingFace models, it is extremely common to pass these arguments positionally (e.g., model(input_ids)). If passed positionally, kwargs.get("input_ids") will be None, causing the patch to silently fail to generate position_ids.

Additionally, if position_ids is passed as a positional argument (even if it is None), directly adding it to kwargs will raise a TypeError: forward() got multiple values for keyword argument 'position_ids' when calling original_forward.

We can resolve both issues robustly by inspecting the signature of original_forward once during patching and dynamically mapping positional arguments to their corresponding names.

def _patch_model_forward(cls: type) -> None: 
    if cls.__dict__.get(_PATCHED_ATTR, False):
        return

    original_forward = cls.forward

    import inspect
    try:
        sig = inspect.signature(original_forward)
        param_to_idx = {}
        for idx, (name, param) in enumerate(sig.parameters.items()):
            if name == "self":
                continue
            param_to_idx[name] = idx - 1
    except Exception:
        param_to_idx = {}

    def patched_forward(self, *args, **kwargs):
        def _get_val(name, default=None):
            if name in kwargs:
                return kwargs[name]
            idx = param_to_idx.get(name)
            if idx is not None and idx < len(args):
                return args[idx]
            return default

        if _get_val("position_ids") is None:
            position_ids, rope_deltas = _try_build_packed_mrope_position_ids(
                self,
                input_ids=_get_val("input_ids"),
                image_grid_thw=_get_val("image_grid_thw"),
                video_grid_thw=_get_val("video_grid_thw"),
                packed_seq_params=_get_val("packed_seq_params"),
            )
            if rope_deltas is not None and hasattr(self, "rope_deltas"):
                self.rope_deltas = rope_deltas
            if position_ids is not None:
                pos_idx = param_to_idx.get("position_ids")
                if pos_idx is not None and pos_idx < len(args):
                    args_list = list(args)
                    args_list[pos_idx] = position_ids
                    args = tuple(args_list)
                else:
                    kwargs["position_ids"] = position_ids

        return original_forward(self, *args, **kwargs)

    cls.forward = patched_forward
    setattr(cls, _PATCHED_ATTR, True)

miles/backends/megatron_utils/qwen_vl_packed_mrope.py (103-199)

The current implementation performs multiple GPU-to-CPU synchronizations inside the loop over segments (e.g., starts.numel() == 0 and .item() inside _count_segment_media for every segment). In a high-performance training or inference loop, these synchronization points can severely degrade GPU utilization and throughput.

Since input_ids is a very small tensor containing only token IDs, we can copy the entire flat_input_ids to CPU once before the loop. This allows us to perform all segment slicing, media counting, and linear position ID generation entirely on the CPU with zero GPU-CPU synchronizations. We only slice the GPU tensor when we actually need to call model.get_rope_index for segments containing media.

def _try_build_packed_mrope_position_ids(
    model: Any,
    *,
    input_ids: torch.Tensor | None,
    image_grid_thw: torch.Tensor | None,
    video_grid_thw: torch.Tensor | None,
    packed_seq_params: Any,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
    if input_ids is None or packed_seq_params is None:
        return None, None
    if getattr(packed_seq_params, "qkv_format", None) != "thd":
        return None, None
    if input_ids.ndim != 2 or input_ids.size(0) != 1:
        return None, None
    if not hasattr(model, "get_rope_index"):
        return None, None

    cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q", None)
    if cu_seqlens is None:
        return None, None

    cu = [int(x) for x in cu_seqlens.detach().cpu().tolist()]
    flat_input_ids = input_ids.squeeze(0)
    if not cu or cu[0] != 0 or cu[-1] != flat_input_ids.numel():
        logger.debug(
            "Skipping Qwen3-VL packed mRoPE patch because cu_seqlens=%s does not match input length=%s",
            cu,
            flat_input_ids.numel(),
        )
        return None, None

    image_offset = 0
    video_offset = 0
    packed_position_ids: list[torch.Tensor] = []
    rope_deltas: list[torch.Tensor] = []

    flat_input_ids_cpu = flat_input_ids.cpu()
    config = getattr(model, "config", model)

    for start, end in zip(cu[:-1], cu[1:], strict=False):
        segment_cpu = flat_input_ids_cpu[start:end]
        numel = segment_cpu.numel()
        if numel == 0:
            continue

        image_count, video_count = _count_segment_media(config, segment_cpu)
        segment_image_grid = _slice_optional_grid(image_grid_thw, image_offset, image_count)
        segment_video_grid = _slice_optional_grid(video_grid_thw, video_offset, video_count)
        image_offset += image_count
        video_offset += video_count

        if image_count == 0 and video_count == 0:
            pos, delta = _linear_position_ids(numel, flat_input_ids.dtype, flat_input_ids.device)
        else:
            segment_gpu = flat_input_ids[start:end]
            pos, delta = model.get_rope_index(
                input_ids=segment_gpu.unsqueeze(0),
                image_grid_thw=segment_image_grid,
                video_grid_thw=segment_video_grid,
                attention_mask=torch.ones((1, numel), dtype=torch.long, device=flat_input_ids.device),
            )

        packed_position_ids.append(pos[:, 0, : numel])
        rope_deltas.append(delta.reshape(1, -1))

    if not packed_position_ids:
        return None, None

    return torch.cat(packed_position_ids, dim=1).unsqueeze(1), torch.cat(rope_deltas, dim=0)


def _count_segment_media(config: Any, segment_cpu: torch.Tensor) -> tuple[int, int]:
    vision_start_token_id = getattr(config, "vision_start_token_id", None)
    image_token_id = getattr(config, "image_token_id", None)
    video_token_id = getattr(config, "video_token_id", None)
    if vision_start_token_id is None or image_token_id is None or video_token_id is None:
        return 0, 0

    starts = torch.nonzero(segment_cpu == vision_start_token_id, as_tuple=False).flatten()
    starts = starts[starts + 1 < segment_cpu.numel()]
    if starts.numel() == 0:
        return 0, 0

    vision_tokens = segment_cpu[starts + 1]
    return int((vision_tokens == image_token_id).sum().item()), int((vision_tokens == video_token_id).sum().item())


def _slice_optional_grid(grid: torch.Tensor | None, offset: int, count: int) -> torch.Tensor | None:
    if grid is None or count == 0:
        return None
    return grid[offset : offset + count]


def _linear_position_ids(numel: int, dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
    position_ids = (
        torch.arange(numel, dtype=dtype, device=device)
        .view(1, 1, -1)
        .expand(3, 1, -1)
    )
    rope_delta = torch.zeros((1, 1), dtype=dtype, device=device)
    return position_ids, rope_delta

[TSunny007]

Rename this and the corresponding files as qwen3_vl...?

[TSunny007]

Looking at #1272, I think their analysis is right — forward resets position_ids = None internally, so injecting it via kwargs here won't survive the reset. Posting my CP-padded note over there instead.

[TSunny007]

Can confirm that this stabilizes our sequence packed runs and the rewards we collect after this change

[Zhichenzzz]

Nice — cleaner than what we have locally.

One gap: _try_build_packed_mrope_position_ids uses cu_seqlens_q, but with CP + THD packing the input_ids are laid out against cu_seqlens_q_padded (see bridge's own thd_to_bshd at modelling_qwen3_vl/model.py:271-275). Right now this no-ops in that case and CP+packing regresses to the broken cross-sample positions.

Plan to handle that here, or as a follow-up? We have it working in https://github.com/GymPod/miles/pull/23 and can fold it in once this lands so we can drop our local patch.

could you try this PR patch? #1272. btw, I have no access to your GymPod forked miles lol

[Zhichenzzz]

Opened another PR #1272 to solve the problems, will close this PR.