diff --git a/src/mobius/_configs/_vision_defaults.py b/src/mobius/_configs/_vision_defaults.py index a72e0ac2a..892fc11b7 100644 --- a/src/mobius/_configs/_vision_defaults.py +++ b/src/mobius/_configs/_vision_defaults.py @@ -101,6 +101,9 @@ def apply_vision_defaults(config, parent_config, model_type: str, fields: dict) # per-model hook may overwrite later via fields.update(...). fields["mm_tokens_per_image"] = getattr(vision_source, "mm_tokens_per_image", None) fields["image_token_id"] = getattr(vision_source, "image_token_id", None) + fields["video_token_id"] = getattr(vision_source, "video_token_id", None) + fields["vision_start_token_id"] = getattr(vision_source, "vision_start_token_id", None) + fields["vision_end_token_id"] = getattr(vision_source, "vision_end_token_id", None) # MRoPE section — only for composite VL models (parent_config != config). if parent_config is not None and parent_config is not config: diff --git a/src/mobius/components/_qwen3_vl_vision.py b/src/mobius/components/_qwen3_vl_vision.py index 1fc983989..7da0d5850 100644 --- a/src/mobius/components/_qwen3_vl_vision.py +++ b/src/mobius/components/_qwen3_vl_vision.py @@ -28,11 +28,6 @@ from mobius._build_context import ep_capabilities, get_build_dtype from mobius.components._common import LayerNorm, Linear, build_packed_token_offset from mobius.components._mlp import FCMLP -from mobius.components._scan_utils import ( - compact_scan_output, - create_body_graph, - rename_subgraph_values, -) class Qwen3VLPatchEmbed(nn.Module): @@ -403,69 +398,6 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return self.linear_fc2(op, x) -def _qwen3_rotary_pos_ids_one_image(op, T, H, W, ms): # noqa: N803 - """Compute 2D rotary position IDs for one image (Qwen3-VL style). - - Uses block_rows * ms + intra indexing for spatial-merge groups. - Works with any OpBuilder (main graph or Scan body graph). - - Args: - op: OpBuilder instance. - T, H, W: Scalar INT64 values. - ms: Python int — spatial merge size. - - Returns: - ``(T*H*W, 2)`` INT64 position IDs. - """ - H_m = op.Div(H, op.Constant(value_int=ms)) # noqa: N806 - W_m = op.Div(W, op.Constant(value_int=ms)) # noqa: N806 - - # Block row/col indices and intra-merge indices - block_rows = op.Range( - op.Constant(value_int=0), - H_m, - op.Constant(value_int=1), - ) - block_cols = op.Range( - op.Constant(value_int=0), - W_m, - op.Constant(value_int=1), - ) - intra = op.Range( - op.Constant(value_int=0), - op.Constant(value_int=ms), - op.Constant(value_int=1), - ) - - # row_idx = block_rows[:,None,None,None] * ms + intra[None,None,:,None] - br = op.Mul(op.Unsqueeze(block_rows, [1, 2, 3]), op.Constant(value_int=ms)) - ir_row = op.Unsqueeze(intra, [0, 1, 3]) - row_idx = op.Add(br, ir_row) - - bc = op.Mul(op.Unsqueeze(block_cols, [0, 2, 3]), op.Constant(value_int=ms)) - ir_col = op.Unsqueeze(intra, [0, 1, 2]) - col_idx = op.Add(bc, ir_col) - - # Expand to (H_m, W_m, ms, ms) and flatten - row_shape = op.Concat( - op.Reshape(H_m, [1]), - op.Reshape(W_m, [1]), - op.Constant(value_ints=[ms, ms]), - axis=0, - ) - row_flat = op.Reshape(op.Expand(row_idx, row_shape), [-1]) - col_flat = op.Reshape(op.Expand(col_idx, row_shape), [-1]) - - # Stack to (H*W, 2) and tile T times - pos_ids = op.Concat( - op.Unsqueeze(row_flat, [1]), - op.Unsqueeze(col_flat, [1]), - axis=1, - ) - tile_t = op.Concat(op.Reshape(T, [1]), op.Constant(value_ints=[1]), axis=0) - return op.Tile(pos_ids, tile_t) # (T*H*W, 2) - - class Qwen3VLVisionModel(nn.Module): """Full Qwen3-VL vision encoder with DeepStack outputs. @@ -558,291 +490,147 @@ def __init__( ] ) - def _interpolate_pos_embed(self, op, grid_thw): - """Bilinear interpolation of learned position embeddings for all images. - - Iterates over ``grid_thw`` via ONNX Scan, computing per-image - bilinear interpolation from the learned position grid and - concatenating results. + def _flat_grid_coordinates(self, op, grid_thw): + """Map each packed patch to its media row and merge-permuted H/W coordinates. - Matches HuggingFace ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. - - Args: - op: OpBuilder instance. - grid_thw: ``(num_images, 3)`` INT64 with ``[T, H, W]`` per image. - - Returns: - Position embeddings ``(total_patches, hidden_size)``. + Qwen3-VL stores each media item's patches in + ``(T, H // ms, W // ms, ms, ms)`` order. Computing coordinates over + the concatenated patch stream avoids control-flow subgraphs while + preserving arbitrary image/video sizes and order. """ - n = self.num_grid_per_side ms = self.spatial_merge_size - hidden_size = self.hidden_size - n_minus_1 = float(n - 1) - - # Per-image patch counts - T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 - W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 - patches_per_image = op.Mul(T_col, op.Mul(H_col, W_col)) - max_patches = op.ReduceMax(patches_per_image, keepdims=False) - - # --- Scan body: interpolate pos embeddings for one image --- - body_thw = ir.Value( - name="body_thw", - shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), - ) - body_graph, body_builder = create_body_graph([], [body_thw]) - body_op = body_builder.op - - bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) # noqa: N806 - bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) # noqa: N806 - bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) # noqa: N806 - - # linspace(0, n-1, H) and linspace(0, n-1, W) - H_f = body_op.Cast(bH, to=1) # noqa: N806 - W_f = body_op.Cast(bW, to=1) # noqa: N806 - h_range = body_op.Cast( - body_op.Range( - body_op.Constant(value_int=0), - bH, - body_op.Constant(value_int=1), - ), - to=1, - ) - w_range = body_op.Cast( - body_op.Range( - body_op.Constant(value_int=0), - bW, - body_op.Constant(value_int=1), - ), - to=1, - ) - - h_idxs = body_op.Div( - body_op.Mul(h_range, n_minus_1), - body_op.Sub(H_f, 1.0), - ) - w_idxs = body_op.Div( - body_op.Mul(w_range, n_minus_1), - body_op.Sub(W_f, 1.0), - ) - - # Floor/ceil indices - h_floor = body_op.Cast(body_op.Floor(h_idxs), to=7) - w_floor = body_op.Cast(body_op.Floor(w_idxs), to=7) - clip_max = body_op.Constant(value_int=n - 1) - h_ceil = body_op.Min( - body_op.Add(h_floor, body_op.Constant(value_int=1)), - clip_max, - ) - w_ceil = body_op.Min( - body_op.Add(w_floor, body_op.Constant(value_int=1)), - clip_max, - ) - - # Bilinear weights - dh = body_op.Sub(h_idxs, body_op.Cast(h_floor, to=1)) - dw = body_op.Sub(w_idxs, body_op.Cast(w_floor, to=1)) - - n_const = body_op.Constant(value_int=n) - base_h_floor = body_op.Mul(h_floor, n_const) - base_h_ceil = body_op.Mul(h_ceil, n_const) - - bh_f2 = body_op.Unsqueeze(base_h_floor, [1]) - bh_c2 = body_op.Unsqueeze(base_h_ceil, [1]) - wf2 = body_op.Unsqueeze(w_floor, [0]) - wc2 = body_op.Unsqueeze(w_ceil, [0]) - - idx_00 = body_op.Reshape(body_op.Add(bh_f2, wf2), [-1]) - idx_01 = body_op.Reshape(body_op.Add(bh_f2, wc2), [-1]) - idx_10 = body_op.Reshape(body_op.Add(bh_c2, wf2), [-1]) - idx_11 = body_op.Reshape(body_op.Add(bh_c2, wc2), [-1]) - - one_minus_dh = body_op.Sub(1.0, dh) - one_minus_dw = body_op.Sub(1.0, dw) - dh2 = body_op.Unsqueeze(dh, [1]) - omdh2 = body_op.Unsqueeze(one_minus_dh, [1]) - dw2 = body_op.Unsqueeze(dw, [0]) - omdw2 = body_op.Unsqueeze(one_minus_dw, [0]) - - w_00 = body_op.Reshape(body_op.Mul(omdh2, omdw2), [-1, 1]) - w_01 = body_op.Reshape(body_op.Mul(omdh2, dw2), [-1, 1]) - w_10 = body_op.Reshape(body_op.Mul(dh2, omdw2), [-1, 1]) - w_11 = body_op.Reshape(body_op.Mul(dh2, dw2), [-1, 1]) - - # Gather from learned pos_embed (implicit input from parent graph). - # Cast to float32 for bilinear interpolation (pos_embed may be bf16/f16). - e_00 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_00), to=1), w_00) - e_01 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_01), to=1), w_01) - e_10 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_10), to=1), w_10) - e_11 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_11), to=1), w_11) - pos_embeds = body_op.Add( - body_op.Add(e_00, e_01), - body_op.Add(e_10, e_11), + T_col = op.Gather(grid_thw, op.Constant(value_int=0), axis=1) # noqa: N806 + H_col = op.Gather(grid_thw, op.Constant(value_int=1), axis=1) # noqa: N806 + W_col = op.Gather(grid_thw, op.Constant(value_int=2), axis=1) # noqa: N806 + patches_per_media = op.Mul(T_col, op.Mul(H_col, W_col)) + patch_ends = op.CumSum(patches_per_media, op.Constant(value_int=0)) + patch_starts = op.Pad( + patch_ends, + op.Constant(value_ints=[1, 0]), + op.Constant(value_int=0), ) - # Tile T times: (H*W, D) → (T*H*W, D) - T_tile = body_op.Concat( # noqa: N806 - body_op.Reshape(bT, [1]), - body_op.Constant(value_ints=[1]), - axis=0, - ) - pos_embeds = body_op.Tile(pos_embeds, T_tile) - - # Spatial merge permutation: - # (T, H//ms, ms, W//ms, ms, D) → (T, H//ms, W//ms, ms, ms, D) - H_m = body_op.Div(bH, body_op.Constant(value_int=ms)) # noqa: N806 - W_m = body_op.Div(bW, body_op.Constant(value_int=ms)) # noqa: N806 - shape_6d = body_op.Concat( - body_op.Reshape(bT, [1]), - body_op.Reshape(H_m, [1]), - body_op.Constant(value_ints=[ms]), - body_op.Reshape(W_m, [1]), - body_op.Constant(value_ints=[ms]), - body_op.Constant(value_ints=[hidden_size]), - axis=0, + total_patches = op.ReduceSum(patches_per_media, keepdims=False) + patch_ids = op.Range( + op.Constant(value_int=0), + total_patches, + op.Constant(value_int=1), ) - pos_embeds = body_op.Reshape(pos_embeds, shape_6d) - pos_embeds = body_op.Transpose(pos_embeds, perm=[0, 1, 3, 2, 4, 5]) - pos_embeds = body_op.Reshape(pos_embeds, [-1, hidden_size]) - - # Pad to (max_patches, hidden_size) — implicit input from main graph - num_p = body_op.Mul(bT, body_op.Mul(bH, bW)) - pad_len = body_op.Reshape(body_op.Sub(max_patches, num_p), [1]) - pads = body_op.Concat( - body_op.Constant(value_ints=[0, 0]), - pad_len, - body_op.Constant(value_ints=[0]), + # Mark each nonzero media boundary, then prefix-sum the markers. This + # maps patches to media in O(total_patches + num_media) rather than + # materializing an O(total_patches * num_media) comparison matrix. + media_boundaries = op.Slice(patch_ends, [0], [-1]) + boundary_updates = op.ConstantOfShape( + op.Shape(media_boundaries), + value=ir.tensor(np.array([1], dtype=np.int64)), + ) + boundary_markers = op.ScatterElements( + op.Mul(patch_ids, op.Constant(value_int=0)), + media_boundaries, + boundary_updates, axis=0, ) - padded = body_op.Pad(pos_embeds, pads, 0.0) - padded.name = "padded_pos_embed" - body_graph.outputs.append(padded) + media_ids = op.CumSum(boundary_markers, op.Constant(value_int=0)) + local_ids = op.Sub(patch_ids, op.Gather(patch_starts, media_ids)) + + H = op.Gather(H_col, media_ids) # noqa: N806 + W = op.Gather(W_col, media_ids) # noqa: N806 + patches_per_frame = op.Mul(H, W) + frame_local_ids = op.Mod(local_ids, patches_per_frame) + + merge_area = op.Constant(value_int=ms * ms) + merge_block_ids = op.Div(frame_local_ids, merge_area) + intra_merge_ids = op.Mod(frame_local_ids, merge_area) + W_m = op.Div(W, op.Constant(value_int=ms)) # noqa: N806 + block_rows = op.Div(merge_block_ids, W_m) + block_cols = op.Mod(merge_block_ids, W_m) + intra_rows = op.Div(intra_merge_ids, op.Constant(value_int=ms)) + intra_cols = op.Mod(intra_merge_ids, op.Constant(value_int=ms)) + rows = op.Add(op.Mul(block_rows, op.Constant(value_int=ms)), intra_rows) + cols = op.Add(op.Mul(block_cols, op.Constant(value_int=ms)), intra_cols) + return rows, cols, H, W, frame_local_ids, patch_ids, total_patches + + def _interpolate_pos_embed(self, op, coordinates): + """Bilinearly interpolate learned positions for the packed media stream. - rename_subgraph_values(body_graph, "posemb_body_") - - scan_result = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, - ) # (num_images, max_patches, hidden_size) + Matches HuggingFace ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. - return compact_scan_output(op, scan_result, patches_per_image) + Args: + op: OpBuilder instance. + coordinates: Shared packed ``(rows, cols, H, W)`` coordinate values. - def _compute_rotary_pos_ids(self, op, grid_thw): - """Compute 2D rotary position IDs for all images via ONNX Scan. + Returns: + Position embeddings ``(total_patches, hidden_size)``. + """ + n = self.num_grid_per_side + rows, cols, H, W = coordinates # noqa: N806 + rows_f = op.Cast(rows, to=1) + cols_f = op.Cast(cols, to=1) + H_f = op.Cast(H, to=1) # noqa: N806 + W_f = op.Cast(W, to=1) # noqa: N806 + rows_scaled = op.Div(op.Mul(rows_f, float(n - 1)), op.Sub(H_f, 1.0)) + cols_scaled = op.Div(op.Mul(cols_f, float(n - 1)), op.Sub(W_f, 1.0)) + + row_floor_f = op.Floor(rows_scaled) + col_floor_f = op.Floor(cols_scaled) + row_floor = op.Cast(row_floor_f, to=7) + col_floor = op.Cast(col_floor_f, to=7) + clip_max = op.Constant(value_int=n - 1) + row_ceil = op.Min(op.Add(row_floor, op.Constant(value_int=1)), clip_max) + col_ceil = op.Min(op.Add(col_floor, op.Constant(value_int=1)), clip_max) + + row_delta = op.Unsqueeze(op.Sub(rows_scaled, row_floor_f), [1]) + col_delta = op.Unsqueeze(op.Sub(cols_scaled, col_floor_f), [1]) + + row_floor_base = op.Mul(row_floor, op.Constant(value_int=n)) + row_ceil_base = op.Mul(row_ceil, op.Constant(value_int=n)) + idx_00 = op.Add(row_floor_base, col_floor) + idx_01 = op.Add(row_floor_base, col_ceil) + idx_10 = op.Add(row_ceil_base, col_floor) + idx_11 = op.Add(row_ceil_base, col_ceil) + + # Interpolate in float32 even when the learned table is f16/bf16. + pos_embed_f = op.Cast(self.pos_embed, to=1) + e_00 = op.Gather(pos_embed_f, idx_00) + e_01 = op.Gather(pos_embed_f, idx_01) + e_10 = op.Gather(pos_embed_f, idx_10) + e_11 = op.Gather(pos_embed_f, idx_11) + + # Two horizontal lerps followed by one vertical lerp are equivalent to + # the four explicit bilinear weights with fewer graph operations. + top = op.Add(e_00, op.Mul(op.Sub(e_01, e_00), col_delta)) + bottom = op.Add(e_10, op.Mul(op.Sub(e_11, e_10), col_delta)) + return op.Add(top, op.Mul(op.Sub(bottom, top), row_delta)) + + def _compute_rotary_pos_ids(self, op, coordinates): + """Combine shared packed coordinates into 2D rotary position IDs. Matches HF ``Qwen3VLVisionModel.rot_pos_emb()`` position indexing. - Iterates over ``grid_thw`` rows, computing per-image spatial-merge- - permuted position IDs and concatenating. Returns ``(total_patches, 2)`` INT64 with ``[h_pos, w_pos]`` per patch. """ - ms = self.spatial_merge_size - - # Per-image patch counts for padding/compaction - T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 - W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 - patches_per_image = op.Mul(T_col, op.Mul(H_col, W_col)) - max_patches = op.ReduceMax(patches_per_image, keepdims=False) - - # --- Scan body: compute pos_ids for one image, pad to max_patches --- - body_thw = ir.Value( - name="body_thw", - shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), - ) - body_graph, body_builder = create_body_graph([], [body_thw]) - body_op = body_builder.op - - bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) # noqa: N806 - bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) # noqa: N806 - bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) # noqa: N806 - - pos_ids = _qwen3_rotary_pos_ids_one_image(body_op, bT, bH, bW, ms) - - # Pad to (max_patches, 2) — implicit input from main graph - num_p = body_op.Mul(bT, body_op.Mul(bH, bW)) - pad_len = body_op.Reshape(body_op.Sub(max_patches, num_p), [1]) - pads = body_op.Concat( - body_op.Constant(value_ints=[0, 0]), - pad_len, - body_op.Constant(value_ints=[0]), - axis=0, - ) - padded = body_op.Pad(pos_ids, pads, body_op.Constant(value_int=-1)) - padded.name = "padded_pos_ids" - body_graph.outputs.append(padded) + rows, cols = coordinates + return op.Concat(op.Unsqueeze(rows, [1]), op.Unsqueeze(cols, [1]), axis=1) - rename_subgraph_values(body_graph, "q3_rotary_body_") - - scan_result = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, - ) - return compact_scan_output(op, scan_result, patches_per_image) - - def _compute_cu_seqlens(self, op, grid_thw): + def _compute_cu_seqlens(self, op, frame_boundaries): """Compute full-attention cu_seqlens for all images. - Produces per-frame boundaries across all images using ONNX Scan - to handle per-image ``repeat_interleave(hw, T)`` + CumSum. + Each packed frame starts where its frame-local patch index is zero. + Compacting those patch IDs and appending the total patch count produces + the same boundaries as ``repeat_interleave(H * W, T)`` + CumSum. Returns ``(total_frames + 1,)`` INT64. """ - T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - max_T = op.ReduceMax(T_col, keepdims=False) # noqa: N806 - - # Scan body: for each image, output T copies of hw, padded to max_T - body_thw = ir.Value( - name="body_thw", - shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), - ) - body_graph, body_builder = create_body_graph([], [body_thw]) - body_op = body_builder.op - - bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) # noqa: N806 - bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) # noqa: N806 - bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) # noqa: N806 - - hw = body_op.Mul(bH, bW) - ones = body_op.Expand( - body_op.Constant(value_int=1), - body_op.Reshape(bT, [1]), - ) - hw_repeated = body_op.Mul(ones, hw) - - pad_len = body_op.Reshape(body_op.Sub(max_T, bT), [1]) - pads = body_op.Concat( - body_op.Constant(value_ints=[0]), - pad_len, + frame_local_ids, patch_ids, total_patches = frame_boundaries + frame_starts = op.Compress( + patch_ids, + op.Equal(frame_local_ids, op.Constant(value_int=0)), + ) + return op.Concat( + frame_starts, + op.Unsqueeze(total_patches, [0]), axis=0, ) - padded = body_op.Pad( - hw_repeated, - pads, - body_op.Constant(value_int=0), - ) - padded.name = "padded_hw" - body_graph.outputs.append(padded) - - rename_subgraph_values(body_graph, "q3_cu_body_") - - scan_hw = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, - ) - hw_flat = compact_scan_output(op, scan_hw, T_col) - cu = op.CumSum(hw_flat, op.Constant(value_int=0)) - return op.Pad(cu, op.Constant(value_ints=[1, 0]), op.Constant(value_int=0)) def forward( self, @@ -865,18 +653,22 @@ def forward( # Patch embedding hidden_states = self.patch_embed(op, hidden_states) + # Compute the packed patch coordinates once. Position interpolation, + # rotary IDs, and frame boundaries share these values. + coordinates = self._flat_grid_coordinates(op, grid_thw) + # Bilinear-interpolated position embeddings from learned grid. - # Cast to match hidden_states dtype (Scan body computes in float32). - pos_embeds = self._interpolate_pos_embed(op, grid_thw) + # Cast to match hidden_states dtype (interpolation computes in float32). + pos_embeds = self._interpolate_pos_embed(op, coordinates[:4]) pos_embeds = op.CastLike(pos_embeds, hidden_states) hidden_states = op.Add(hidden_states, pos_embeds) # Compute rotary position IDs and embeddings from grid_thw - rotary_pos_ids = self._compute_rotary_pos_ids(op, grid_thw) + rotary_pos_ids = self._compute_rotary_pos_ids(op, coordinates[:2]) position_embeddings = self.rotary_pos_emb(op, rotary_pos_ids) # Compute cu_seqlens from grid_thw - cu_seqlens = self._compute_cu_seqlens(op, grid_thw) + cu_seqlens = self._compute_cu_seqlens(op, coordinates[4:]) # Transformer blocks deepstack_features = [] diff --git a/src/mobius/components/_qwen3_vl_vision_test.py b/src/mobius/components/_qwen3_vl_vision_test.py new file mode 100644 index 000000000..cfd09ffba --- /dev/null +++ b/src/mobius/components/_qwen3_vl_vision_test.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import onnx_ir as ir + +from mobius._testing import count_op_type, create_test_builder, create_test_input +from mobius.components._qwen3_vl_vision import Qwen3VLVisionModel + +_PATCH_DIM = 3 * 2 * 16 * 16 + + +def _build_vision_graph() -> ir.Graph: + module = Qwen3VLVisionModel( + depth=1, + hidden_size=32, + intermediate_size=64, + num_heads=4, + patch_size=16, + temporal_patch_size=2, + in_channels=3, + out_hidden_size=64, + spatial_merge_size=2, + num_position_embeddings=16, + deepstack_visual_indexes=[], + ) + builder, op, graph = create_test_builder() + pixel_values = create_test_input( + builder, + "pixel_values", + ["total_patches", _PATCH_DIM], + dtype=ir.DataType.FLOAT, + ) + grid_thw = create_test_input( + builder, + "grid_thw", + ["num_media", 3], + dtype=ir.DataType.INT64, + ) + image_features = module(op, pixel_values, grid_thw)[0] + image_features.name = "image_features" + graph.outputs.append(image_features) + return graph + + +def test_packed_coordinates_are_linear_and_shared(): + graph = _build_vision_graph() + + # Media ownership uses boundary scatter + prefix sum, not a quadratic + # [total_patches, num_media] comparison matrix. + assert count_op_type(graph, "ScatterElements") == 1 + # The only remaining comparison belongs to the single attention block. + assert count_op_type(graph, "GreaterOrEqual") == 1 + + # The two row/column coordinate values each feed both interpolation (Cast) + # and rotary IDs (Unsqueeze), proving the coordinate graph is emitted once. + shared_coordinates = [] + for node in graph: + if node.op_type != "Add": + continue + for output in node.outputs: + consumer_types = {consumer.op_type for consumer, _ in output.uses()} + if {"Cast", "Unsqueeze"} <= consumer_types: + shared_coordinates.append(output) + assert len(shared_coordinates) == 2 diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index b281517a4..f24fb68a0 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -92,8 +92,9 @@ "qwen2_vl": "qwen2_5_vl", "qwen3_vl": "qwen3_vl", "qwen3_vl_text": "qwen3_vl", - "qwen3_5": "qwen2_5_vl", - "qwen3_5_vl": "qwen2_5_vl", + "qwen3_5": "qwen3_5", + "qwen3_5_vl": "qwen3_5", + "qwen3_5_text": "qwen3_5_text", # MiniCPM uses standard 1D decoder position IDs (unlike Qwen-VL MRoPE). # The phi3v multimodal runtime provides that contract; callers supply # HF-preprocessed packed pixels through Generator.set_inputs(). @@ -136,6 +137,7 @@ "qwen3_vl_text", "qwen3_5", "qwen3_5_vl", + "qwen3_5_text", "qwen3_5_moe", "videochat_flash_qwen", } @@ -184,6 +186,11 @@ def _select_ort_model_type( """ if is_decoder_only and config_model_type in _ORT_GENAI_MODEL_TYPE: return _ORT_GENAI_MODEL_TYPE[config_model_type] + if not is_decoder_only and config_model_type == "qwen3_5_text": + # Qwen3.5/Qwen3.8 multimodal builds unwrap the parent config to its + # text subtype, but their vision/embedding package uses the multimodal + # ORT pipeline and processor metadata. + return "qwen3_5" return _resolve_ort_genai_model_type(hf_model_type or "unknown") @@ -1377,7 +1384,11 @@ def write_ort_genai_config( # does not bind, so borrowing that type would mis-wire the graph. ort_model_type = "gemma3n" else: - ort_model_type = _resolve_ort_genai_model_type(raw_type) + ort_model_type = _select_ort_model_type( + raw_type, + raw_type, + is_decoder_only=is_decoder_only, + ) if ort_model_type == "unknown": logger.warning( "Could not determine ORT-GenAI model type: pkg.config.model_type " diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 1731572b3..1dc76cafa 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -82,6 +82,9 @@ class FakeConfig: class TestResolveOrtGenaiModelType: def test_known_model_type(self): assert _resolve_ort_genai_model_type("qwen3") == "qwen2" + assert _resolve_ort_genai_model_type("qwen3_5") == "qwen3_5" + assert _resolve_ort_genai_model_type("qwen3_5_vl") == "qwen3_5" + assert _resolve_ort_genai_model_type("qwen3_5_text") == "qwen3_5_text" assert _resolve_ort_genai_model_type("gemma2") == "gemma" assert _resolve_ort_genai_model_type("llama") == "llama" @@ -141,6 +144,24 @@ def test_multimodal_keeps_hf_type(self): def test_decoder_only_falls_back_to_hf_when_config_missing(self): assert _select_ort_model_type(None, "qwen3", is_decoder_only=True) == "qwen2" + def test_qwen35_text_type_depends_on_package_topology(self): + assert ( + _select_ort_model_type( + "qwen3_5_text", + "qwen3_5", + is_decoder_only=True, + ) + == "qwen3_5_text" + ) + assert ( + _select_ort_model_type( + "qwen3_5_text", + "qwen3_5_text", + is_decoder_only=False, + ) + == "qwen3_5" + ) + def test_decoder_only_unknown_config_falls_back_to_hf(self): # An unrecognised config.model_type (not in _ORT_GENAI_MODEL_TYPE) must # not pass straight through as an invalid ORT type; fall back to the @@ -354,6 +375,38 @@ def test_muse_glimmer_uses_packed_qwen_image_pipeline(self, tmp_path): "merge_size": 2, } + def test_qwen35_text_subtype_uses_packed_qwen_image_pipeline(self, tmp_path): + vision = types.SimpleNamespace( + image_size=448, + patch_size=16, + spatial_merge_size=2, + model_type="qwen3_5", + ) + config = types.SimpleNamespace( + vision=vision, + model_type="qwen3_5_text", + spatial_merge_size=2, + temporal_patch_size=2, + ) + + path = _write_vision_processor_config(config, str(tmp_path)) + + assert path is not None + with open(path, encoding="utf-8") as config_file: + processor = json.load(config_file)["processor"] + assert processor["name"] == "qwen2_5_image_processor" + transforms = processor["transforms"] + assert transforms[-1]["operation"] == { + "name": "patch_image", + "type": "PatchImage", + "attrs": { + "patch_size": 16, + "temporal_patch_size": 2, + "merge_size": 2, + }, + } + assert transforms[-2]["operation"]["attrs"]["qwen2_5_vl"] == 1 + def test_gemma3_vision_config(self, tmp_path): """Gemma3 gets a fixed-size resize + Permute3D (not the generic branch). @@ -1085,6 +1138,107 @@ class FakeConfig: "present_conv_names": "present.%d.conv_state", } + def test_qwen35_vl_hybrid_metadata_is_emitted_without_runtime_gate(self, tmp_path): + import dataclasses + + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "qwen3_5_text" + vocab_size: int = 256 + hidden_size: int = 64 + num_hidden_layers: int = 4 + num_attention_heads: int = 4 + num_key_value_heads: int = 2 + head_dim: int = 16 + max_position_embeddings: int = 128 + pad_token_id: int = 0 + + package = ModelPackage( + { + "decoder": _mock_model( + inputs=[ + "inputs_embeds", + "attention_mask", + "position_ids", + "past_key_values.0.conv_state", + "past_key_values.0.recurrent_state", + "past_key_values.3.key", + "past_key_values.3.value", + ], + outputs=[ + "logits", + "present.0.conv_state", + "present.0.recurrent_state", + "present.3.key", + "present.3.value", + ], + ), + "embedding": _mock_model(inputs=["input_ids", "image_features"]), + "vision_encoder": _mock_model( + inputs=["pixel_values", "image_grid_thw"], + outputs=["image_features"], + ), + }, + config=FakeConfig(), + ) + + result = write_ort_genai_config(package, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + decoder = generated["model"]["decoder"] + assert generated["model"]["type"] == "qwen3_5" + assert decoder["num_hidden_layers"] == 4 + assert decoder["inputs"]["past_key_names"] == "past_key_values.%d.key" + assert decoder["inputs"]["past_conv_names"] == "past_key_values.%d.conv_state" + + def test_qwen35_text_package_preserves_decoder_only_model_type(self, tmp_path): + import dataclasses + + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "qwen3_5_text" + vocab_size: int = 256 + hidden_size: int = 64 + num_hidden_layers: int = 4 + num_attention_heads: int = 4 + num_key_value_heads: int = 2 + head_dim: int = 16 + max_position_embeddings: int = 128 + pad_token_id: int = 0 + + package = ModelPackage( + { + "model": _mock_model( + inputs=[ + "input_ids", + "attention_mask", + "position_ids", + "past_key_values.0.conv_state", + "past_key_values.3.key", + "past_key_values.3.value", + ], + outputs=[ + "logits", + "present.0.conv_state", + "present.3.key", + "present.3.value", + ], + ) + }, + config=FakeConfig(), + ) + + result = write_ort_genai_config(package, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + assert generated["model"]["type"] == "qwen3_5_text" + def test_olive_renamed_logits_output_is_emitted(self, tmp_path): pkg = _make_fake_llm_pkg("qwen2") pkg["model"] = _mock_model( diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index ea46d8932..d85ea5f5f 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -223,9 +223,10 @@ def preprocess_weights( - Stripping ``language_model.`` prefix from HF checkpoint keys (HF stores weights as ``model.language_model.*`` in safetensors) - Dropping visual encoder keys (``model.visual.*``) - - Dropping multi-token prediction (MTP) keys (``mtp*``): - MTP heads are auxiliary decoding heads used only during - HuggingFace training; they are not needed for inference. + - Dropping multi-token prediction (MTP) keys (``mtp*``). The target + model's normal forward path does not consume this optional + self-speculative drafter; it is packaged separately as + :class:`Qwen35MtpModel` when speculative decoding is requested. - Weight tying (``tie_word_embeddings``) """ cleaned: dict[str, torch.Tensor] = {} @@ -376,9 +377,9 @@ def preprocess_weights( """Preprocess HuggingFace state dict for Qwen3.5-MoE. Handles: - - Dropping multi-token prediction (MTP) keys (``mtp_*``, ``mtp.*``): - MTP heads are auxiliary decoding heads used only during - HuggingFace training; they are not needed for inference. + - Dropping multi-token prediction (MTP) keys (``mtp_*``, ``mtp.*``). + The target model's normal forward path does not consume this optional + self-speculative drafter, which has a separate package contract. - Stripping ``language_model.`` prefix from HF checkpoint keys (HF stores weights as ``model.language_model.*`` in safetensors) - Dropping visual encoder keys (``model.visual.*``) @@ -537,9 +538,8 @@ def preprocess_weights( """ renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - # Drop multi-token prediction (MTP) keys: MTP heads are - # auxiliary decoding heads used only during HuggingFace - # training; they are not needed for inference. + # The standard target package excludes the optional + # self-speculative MTP drafter, which has a separate graph contract. if key.startswith(("mtp_", "mtp.")): continue @@ -611,7 +611,7 @@ def preprocess_weights( """Route language_model weights for standalone decoder build.""" renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - # Drop MTP heads (training-only auxiliary decoders) + # The optional self-speculative MTP drafter is packaged separately. if key.startswith(("mtp_", "mtp.")): continue stripped = key diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index 5ede02207..33cfe7914 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -13,14 +13,283 @@ from __future__ import annotations +import numpy as np +import onnx_ir as ir import torch +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config -from mobius._configs import QuantizationConfig +from mobius._configs import ArchitectureConfig, QuantizationConfig, Qwen35MtpConfig +from mobius._registry import registry from mobius._testing import make_config -from mobius.models.qwen35 import Qwen35MoECausalLMModel +from mobius._testing.ort_inference import OnnxModelSession +from mobius.models.qwen35 import Qwen35MoECausalLMModel, Qwen35VL3ModelCausalLMModel +from mobius.models.qwen_vl import Qwen3VLEmbeddingModel +from mobius.tasks import build_embedding_from_features _E, _H, _INT, _BLK, _BITS = 8, 32, 16, 16, 4 _FC1_OUT = 2 * _INT +_QWEN38_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + + +def _qwen38_config() -> Qwen3_5Config: + layer_types = [ + layer_type + for _ in range(16) + for layer_type in ( + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ) + ] + return Qwen3_5Config( + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + tie_word_embeddings=False, + text_config={ + "model_type": "qwen3_5_text", + "vocab_size": 248320, + "hidden_size": 5120, + "intermediate_size": 17408, + "num_hidden_layers": 64, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "head_dim": 256, + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "max_position_embeddings": 262144, + "layer_types": layer_types, + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "partial_rotary_factor": 0.25, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10_000_000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + }, + "mtp_num_hidden_layers": 1, + "tie_word_embeddings": False, + }, + vision_config={ + "model_type": "qwen3_5", + "depth": 27, + "hidden_size": 1152, + "intermediate_size": 4304, + "num_heads": 16, + "patch_size": 16, + "temporal_patch_size": 2, + "spatial_merge_size": 2, + "out_hidden_size": 5120, + "num_position_embeddings": 2304, + "deepstack_visual_indexes": [], + }, + ) + + +class TestQwen38Alias: + def test_exact_config_extracts_dense_hybrid_vl_architecture(self): + hf_config = _qwen38_config() + config = ArchitectureConfig.from_transformers( + hf_config.text_config, + parent_config=hf_config, + ) + + assert _QWEN38_REVISION == "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + assert registry.get("qwen3_5") is Qwen35VL3ModelCausalLMModel + assert registry.get("qwen3_5_vl") is Qwen35VL3ModelCausalLMModel + assert registry.get_registration("qwen3_5").test_model_id == "Qwen/Qwen3.5-2B" + assert registry.get_registration("qwen3_5_vl").test_model_id == "Qwen/Qwen3.5-2B" + assert config.hidden_size == 5120 + assert config.intermediate_size == 17408 + assert config.num_hidden_layers == 64 + assert config.layer_types == hf_config.text_config.layer_types + assert config.layer_types.count("linear_attention") == 48 + assert config.layer_types.count("full_attention") == 16 + assert config.num_attention_heads == 24 + assert config.num_key_value_heads == 4 + assert config.head_dim == 256 + assert np.isclose(config.partial_rotary_factor, 0.25) + assert config.mrope_interleaved is True + assert config.mrope_section == [11, 11, 10] + assert config.linear_num_key_heads == 16 + assert config.linear_num_value_heads == 48 + assert config.linear_key_head_dim == 128 + assert config.linear_value_head_dim == 128 + assert config.linear_conv_kernel_dim == 4 + assert config.vocab_size == 248320 + assert config.image_token_id == 248056 + assert config.video_token_id == 248057 + assert config.vision_start_token_id == 248053 + assert config.vision_end_token_id == 248054 + assert config.vision is not None + assert config.vision.num_hidden_layers == 27 + assert config.vision.hidden_size == 1152 + assert config.vision.intermediate_size == 4304 + assert config.vision.num_attention_heads == 16 + assert config.vision.patch_size == 16 + assert config.vision.temporal_patch_size == 2 + assert config.vision.spatial_merge_size == 2 + assert config.vision.out_hidden_size == 5120 + assert config.vision.deepstack_visual_indexes == [] + + def test_one_layer_mtp_is_classified_as_separate_optional_drafter(self): + hf_config = _qwen38_config() + assert hf_config.text_config.mtp_num_hidden_layers == 1 + + mtp_config = Qwen35MtpConfig.from_transformers(hf_config) + assert mtp_config.num_hidden_layers == 1 + assert mtp_config.layer_types == ["full_attention"] + assert registry.get_registration("Qwen35MtpModel").task == "qwen35-mtp" + + def test_weight_routing_excludes_separately_packaged_mtp(self): + config = ArchitectureConfig.from_transformers( + _qwen38_config().text_config, + parent_config=_qwen38_config(), + ) + model = Qwen35VL3ModelCausalLMModel(config) + state_dict = { + "model.language_model.embed_tokens.weight": torch.ones(2, 2), + "model.language_model.layers.0.linear_attn.A_log": torch.ones(2), + "model.language_model.layers.3.self_attn.q_proj.weight": torch.ones(2, 2), + "model.visual.blocks.0.mlp.linear_fc1.weight": torch.ones(2, 2), + "lm_head.weight": torch.ones(2, 2), + "mtp.layers.0.self_attn.q_proj.weight": torch.ones(2, 2), + "mtp.fc.weight": torch.ones(2, 2), + } + + result = model.preprocess_weights(state_dict) + + assert "decoder.model.embed_tokens.weight" in result + assert "embedding.embed_tokens.weight" in result + assert "decoder.model.layers.0.linear_attn.A_log" in result + assert "decoder.model.layers.3.self_attn.q_proj.weight" in result + assert "vision_encoder.visual.blocks.0.mlp.up_proj.weight" in result + assert "decoder.lm_head.weight" in result + assert not any(key.startswith("mtp") or ".mtp." in key for key in result) + + def test_qwen_vl_processor_boundary_stays_float32_for_bf16_export(self): + hf_config = _qwen38_config() + config = ArchitectureConfig.from_transformers( + hf_config.text_config, + parent_config=hf_config, + ) + config.dtype = ir.DataType.BFLOAT16 + package = Qwen35VL3ModelCausalLMModel(config) + task = registry.get_registration("qwen3_5").task + + from mobius.tasks import get_task + + vision_model = get_task(task).build(package, config)["vision_encoder"] + + assert vision_model.graph.inputs[0].name == "pixel_values" + assert vision_model.graph.inputs[0].dtype == ir.DataType.FLOAT + assert any(node.op_type == "Cast" for node in vision_model.graph) + + def test_embedding_scatter_matches_separate_image_then_video_streams(self): + config = ArchitectureConfig( + vocab_size=16, + hidden_size=4, + pad_token_id=0, + image_token_id=10, + video_token_id=11, + dtype=ir.DataType.FLOAT, + ) + graph = build_embedding_from_features( + Qwen3VLEmbeddingModel(config), + config, + feature_name="image_features", + feature_dim=config.hidden_size, + ) + embedding_weight = np.arange( + config.vocab_size * config.hidden_size, + dtype=np.float32, + ).reshape(config.vocab_size, config.hidden_size) + for name, initializer in graph.graph.initializers.items(): + if name.endswith("embed_tokens.weight"): + initializer.const_value = ir.tensor(embedding_weight) + + input_ids = np.array( + [ + [config.video_token_id, 1, config.image_token_id], + [2, config.image_token_id, config.video_token_id], + ], + dtype=np.int64, + ) + # HF scatters the two image rows first, then the two video rows. + media_features = np.arange(100, 116, dtype=np.float32).reshape(4, 4) + session = OnnxModelSession(graph) + result = session.run( + { + "input_ids": input_ids, + "image_features": media_features, + } + )["inputs_embeds"] + + expected = embedding_weight[input_ids].copy() + expected[0, 2] = media_features[0] + expected[1, 1] = media_features[1] + expected[0, 0] = media_features[2] + expected[1, 2] = media_features[3] + np.testing.assert_array_equal(result, expected) + + decode_ids = np.array([[3], [4]], dtype=np.int64) + decode = session.run( + { + "input_ids": decode_ids, + "image_features": np.empty((0, config.hidden_size), dtype=np.float32), + } + )["inputs_embeds"] + session.close() + np.testing.assert_array_equal(decode, embedding_weight[decode_ids]) + + def test_embedding_scatter_without_video_token_id(self): + config = ArchitectureConfig( + vocab_size=16, + hidden_size=4, + pad_token_id=0, + image_token_id=10, + video_token_id=None, + dtype=ir.DataType.FLOAT, + ) + graph = build_embedding_from_features( + Qwen3VLEmbeddingModel(config), + config, + feature_name="image_features", + feature_dim=config.hidden_size, + ) + embedding_weight = np.arange( + config.vocab_size * config.hidden_size, + dtype=np.float32, + ).reshape(config.vocab_size, config.hidden_size) + for name, initializer in graph.graph.initializers.items(): + if name.endswith("embed_tokens.weight"): + initializer.const_value = ir.tensor(embedding_weight) + + input_ids = np.array( + [[config.image_token_id, 1], [2, config.image_token_id]], + dtype=np.int64, + ) + image_features = np.arange(100, 108, dtype=np.float32).reshape(2, 4) + session = OnnxModelSession(graph) + result = session.run( + { + "input_ids": input_ids, + "image_features": image_features, + } + )["inputs_embeds"] + session.close() + + expected = embedding_weight[input_ids].copy() + expected[0, 0] = image_features[0] + expected[1, 1] = image_features[1] + np.testing.assert_array_equal(result, expected) def _moe_config(quantization: QuantizationConfig | None) -> object: diff --git a/src/mobius/models/qwen_vl.py b/src/mobius/models/qwen_vl.py index b97a0d4e5..85950f1c7 100644 --- a/src/mobius/models/qwen_vl.py +++ b/src/mobius/models/qwen_vl.py @@ -995,12 +995,12 @@ def preprocess_weights( class Qwen3VLEmbeddingModel(Qwen25VLEmbeddingModel): """Qwen3-VL embedding model for the 3-model split. - Scatters merged image features at image-token positions (like - Qwen2.5-VL) and, when the vision encoder produces DeepStack features, - also scatters each intermediate DeepStack map into a full-length - ``[batch, seq, hidden]`` tensor (zero at non-image positions). The - stacked ``deepstack_embeds`` output is consumed by the decoder, which - adds them to the hidden states of its first ``D`` layers. + Scatters packed image-then-video features at their respective placeholder + positions. When the vision encoder produces DeepStack features, each + intermediate map is scattered with the same media ordering into a + full-length ``[batch, seq, hidden]`` tensor. The stacked + ``deepstack_embeds`` output is consumed by the decoder, which adds them to + the hidden states of its first ``D`` layers. Inputs: - input_ids: (batch, seq_len) INT64 @@ -1013,6 +1013,10 @@ class Qwen3VLEmbeddingModel(Qwen25VLEmbeddingModel): (only when DeepStack is active) """ + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + self.video_token_id = config.video_token_id + def forward( self, op: OpBuilder, @@ -1022,21 +1026,45 @@ def forward( ): text_embeds = self.embed_tokens(op, input_ids) - # Image-token positions and their running index into the packed - # feature tensors (shared by the main image scatter and every - # DeepStack scatter). + # Hugging Face scatters image and video streams independently. The + # package therefore packs every image feature first, then every video + # feature, regardless of placeholder order or batch row. image_mask = op.Equal(input_ids, op.Constant(value_int=self.image_token_id)) - image_mask_3d = op.Unsqueeze(image_mask, [-1]) - mask_int = op.Cast(image_mask, to=7) # INT64 - cumsum = op.CumSum(mask_int, op.Constant(value_int=1)) - indices = op.Clip( - op.Sub(cumsum, op.Constant(value_int=1)), - op.Constant(value_int=0), + if self.video_token_id is None: + video_mask = op.Not(op.Equal(input_ids, input_ids)) + else: + video_mask = op.Equal( + input_ids, + op.Constant(value_int=self.video_token_id), + ) + media_mask = op.Or(image_mask, video_mask) + media_mask_3d = op.Unsqueeze(media_mask, [-1]) + + flat_image_mask_bool = op.Reshape(image_mask, [-1]) + flat_image_mask = op.Cast(flat_image_mask_bool, to=7) + flat_video_mask = op.Cast(op.Reshape(video_mask, [-1]), to=7) + image_indices = op.Sub( + op.CumSum(flat_image_mask, op.Constant(value_int=0)), + op.Constant(value_int=1), + ) + video_indices = op.Add( + op.Sub( + op.CumSum(flat_video_mask, op.Constant(value_int=0)), + op.Constant(value_int=1), + ), + op.ReduceSum(flat_image_mask, keepdims=0), + ) + flat_indices = op.Where( + flat_image_mask_bool, + image_indices, + video_indices, ) + flat_indices = op.Clip(flat_indices, op.Constant(value_int=0)) + indices = op.Reshape(flat_indices, op.Shape(input_ids)) def _scatter(features: ir.Value, fallback: ir.Value) -> ir.Value: - # Pad with one zero row so Gather stays in-bounds for text-only - # input (num_image_tokens == 0); the Where mask discards it. + # Keep Gather valid for text-only/decode calls with zero media rows; + # the Where mask discards the synthetic row. pad_row = op.Expand( op.CastLike(0.0, features), op.Concat( @@ -1047,7 +1075,7 @@ def _scatter(features: ir.Value, fallback: ir.Value) -> ir.Value: ) padded = op.Concat(features, pad_row, axis=0) gathered = op.Gather(padded, indices, axis=0) - return op.Where(image_mask_3d, gathered, fallback) + return op.Where(media_mask_3d, gathered, fallback) inputs_embeds = _scatter(image_features, text_embeds) diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index fa961ede4..f2d5ce650 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -195,7 +195,7 @@ def _build_vision( op = builder.op pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[total_patches, pixel_dim], ) image_grid_thw = builder.input( @@ -203,10 +203,11 @@ def _build_vision( dtype=ir.DataType.INT64, shape=[num_images, 3], ) + model_pixel_values = op.Cast(pixel_values, to=config.dtype) outputs = vision( op, - pixel_values=pixel_values, + pixel_values=model_pixel_values, image_grid_thw=image_grid_thw, ) diff --git a/testdata/cases/vision-language/qwen3_8-27b.yaml b/testdata/cases/vision-language/qwen3_8-27b.yaml new file mode 100644 index 000000000..1b370cc82 --- /dev/null +++ b/testdata/cases/vision-language/qwen3_8-27b.yaml @@ -0,0 +1,26 @@ +model_id: "Qwen/Qwen3.8-27B" +model_type: "qwen3_5" +revision: "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" +task_type: "image-text-to-text" +dtype: "bfloat16" + +inputs: + prompts: + - "Describe this image in detail." + images: + - "pipeline-cat-chonk.jpeg" + +level: "L4+L5" + +generation: + max_new_tokens: 30 + do_sample: false + +ci_skip_reason: >- + The pinned official checkpoint is 55.6 GB and exceeds hosted CI storage and + GPU memory; no reduced or quantized private fixture is committed. +notes: >- + Qwen3.8-27B native image/video model. Dense Qwen3.5 alias with 64 hybrid + layers (48 Gated DeltaNet + 16 gated GQA), a 27-block vision encoder, and an + optional one-layer self-speculative MTP drafter. The standard target package + omits that drafter; Mobius exposes it through the separate qwen35-mtp task. diff --git a/tests/integration_test.py b/tests/integration_test.py index 25816460d..059c85678 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -2028,6 +2028,9 @@ def test_encoder_matches_diffusers(self): # Qwen3.5 hybrid (DeltaNet + full attention) — random-weight tests # --------------------------------------------------------------------------- +_QWEN38_MODEL_ID = "Qwen/Qwen3.8-27B" +_QWEN38_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + def _build_and_compare_qwen35(hf_model, text_config, onnx_module_cls): """Shared helper: build ONNX model, load HF weights, compare logits.""" @@ -2104,7 +2107,10 @@ def test_qwen35_prefill_logits_match(): Qwen3_5ForCausalLM, ) - c = transformers.AutoConfig.from_pretrained("Qwen/Qwen3.5-27B") + c = transformers.AutoConfig.from_pretrained( + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, + ) tc = c.text_config tc.num_hidden_layers = 4 tc.layer_types = [ @@ -3518,14 +3524,17 @@ def make_feeds(token_id, conv_states, rec_states, kv_cache, step): # --------------------------------------------------------------------------- -def _make_tiny_qwen35_vl_config(): +def _make_tiny_qwen35_vl_config(*, keep_production_vocab: bool = False): """Create a tiny Qwen3.5-VL config for fast HF parity testing. Downloads the real Qwen3.5-27B config structure, then overrides all dimensions to be tiny. Also overrides rope_theta to float to avoid a pre-existing float64 rotary cache bug (int ** np.float32 → float64). """ - c = transformers.AutoConfig.from_pretrained("Qwen/Qwen3.5-27B") + c = transformers.AutoConfig.from_pretrained( + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, + ) tc = c.text_config # Truncate layers: 3 DeltaNet + 1 full attention @@ -3543,7 +3552,8 @@ def _make_tiny_qwen35_vl_config(): tc.num_attention_heads = 4 tc.num_key_value_heads = 2 tc.head_dim = 16 - tc.vocab_size = 256 + if not keep_production_vocab: + tc.vocab_size = 256 tc.linear_num_value_heads = 4 tc.linear_num_key_heads = 4 tc.linear_key_head_dim = 8 @@ -3745,7 +3755,8 @@ def test_qwen35_vl_vision_features_match(): # Process real image (resized small for speed — 256 patches) processor = transformers.AutoProcessor.from_pretrained( - "Qwen/Qwen3.5-27B", + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, ) image = Image.open("testdata/pipeline-cat-chonk.jpeg").resize( (64, 64), @@ -3815,6 +3826,181 @@ def test_qwen35_vl_vision_features_match(): assert max_diff < 0.01, f"Vision features max_diff={max_diff:.6f} (expected < 0.01)" +@pytest.mark.integration +def test_qwen38_vl_image_video_mixed_pipeline_matches_huggingface(): + """Pinned Qwen3.8 image/video processor contract matches the ONNX pipeline. + + Runs image-only, video-only, and a two-row mixed batch whose rows use + opposite media placeholder order. Hugging Face scatters image and video + feature streams independently, so the ONNX embedding input packs all image + features first and all video features second. + """ + import onnx_ir as ir + from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5ForConditionalGeneration, + ) + + from mobius import build_from_module + from mobius._weight_loading import apply_weights + + hf_config = _make_tiny_qwen35_vl_config(keep_production_vocab=True) + arch_config = ArchitectureConfig.from_transformers( + hf_config.text_config, + parent_config=hf_config, + ) + arch_config.dtype = ir.DataType.FLOAT + onnx_module = models.Qwen35VL3ModelCausalLMModel(arch_config) + package = build_from_module( + onnx_module, + arch_config, + task="hybrid-qwen-vl", + ) + + torch.manual_seed(1) + hf_model = ( + Qwen3_5ForConditionalGeneration._from_config( + hf_config, + dtype=torch.float32, + ) + .float() + .eval() + ) + weights = onnx_module.preprocess_weights(dict(hf_model.state_dict())) + for model_name, model in package.items(): + apply_weights(model, weights) + unset = [ + name + for name, initializer in model.graph.initializers.items() + if initializer.const_value is None + ] + assert not unset, f"{model_name} has unset target parameters: {unset[:5]}" + + processor = transformers.AutoProcessor.from_pretrained( + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, + ) + image_a = Image.open("testdata/pipeline-cat-chonk.jpeg").convert("RGB").resize((64, 64)) + image_b = image_a.transpose(Image.Transpose.FLIP_LEFT_RIGHT) + video_a = np.stack( + [np.full((64, 64, 3), value, dtype=np.uint8) for value in (16, 64, 128, 224)] + ) + video_b = np.flip(video_a, axis=0).copy() + + cases = { + "image-only": { + "text": ["<|vision_start|><|image_pad|><|vision_end|> Describe."], + "images": [image_a], + }, + "video-only": { + "text": ["<|vision_start|><|video_pad|><|vision_end|> Describe."], + "videos": [video_a], + }, + "mixed-two-row": { + "text": [ + ( + "<|vision_start|><|video_pad|><|vision_end|> Then " + "<|vision_start|><|image_pad|><|vision_end|>." + ), + ( + "<|vision_start|><|image_pad|><|vision_end|> Then " + "<|vision_start|><|video_pad|><|vision_end|>." + ), + ], + "images": [image_a, image_b], + "videos": [video_a, video_b], + }, + } + + vision_session = _make_session(package["vision_encoder"]) + embedding_session = _make_session(package["embedding"]) + decoder_session = _make_session(package["decoder"]) + try: + for case_name, processor_inputs in cases.items(): + hf_inputs = processor( + **processor_inputs, + padding=True, + return_tensors="pt", + ) + with torch.no_grad(): + hf_logits = hf_model(**hf_inputs).logits.numpy() + text_embeds = hf_model.model.language_model.embed_tokens( + hf_inputs["input_ids"] + ) + position_ids = hf_model.model.compute_3d_position_ids( + input_ids=hf_inputs["input_ids"], + inputs_embeds=text_embeds, + image_grid_thw=hf_inputs.get("image_grid_thw"), + video_grid_thw=hf_inputs.get("video_grid_thw"), + attention_mask=hf_inputs["attention_mask"], + past_key_values=None, + mm_token_type_ids=hf_inputs["mm_token_type_ids"], + ) + + media_features = [] + if "pixel_values" in hf_inputs: + media_features.append( + vision_session.run( + { + "pixel_values": hf_inputs["pixel_values"].numpy(), + "image_grid_thw": hf_inputs["image_grid_thw"].numpy(), + } + )["image_features"] + ) + if "pixel_values_videos" in hf_inputs: + media_features.append( + vision_session.run( + { + "pixel_values": hf_inputs["pixel_values_videos"].numpy(), + "image_grid_thw": hf_inputs["video_grid_thw"].numpy(), + } + )["image_features"] + ) + packed_features = np.concatenate(media_features, axis=0) + onnx_embeds = embedding_session.run( + { + "input_ids": hf_inputs["input_ids"].numpy(), + "image_features": packed_features, + } + )["inputs_embeds"] + + feeds: dict[str, np.ndarray] = { + "inputs_embeds": onnx_embeds, + "attention_mask": hf_inputs["attention_mask"].numpy(), + "position_ids": position_ids.numpy(), + } + batch_size = hf_inputs["input_ids"].shape[0] + for graph_input in package["decoder"].graph.inputs: + if graph_input.name in feeds: + continue + shape = tuple( + dim if isinstance(dim, int) else batch_size if axis == 0 else 0 + for axis, dim in enumerate(graph_input.shape) + ) + feeds[graph_input.name] = np.zeros(shape, dtype=np.float32) + + onnx_logits = decoder_session.run(feeds)["logits"] + max_abs = float(np.max(np.abs(onnx_logits - hf_logits))) + cosine = float( + np.dot(onnx_logits.ravel(), hf_logits.ravel()) + / (np.linalg.norm(onnx_logits) * np.linalg.norm(hf_logits)) + ) + print(f"Qwen3.8 {case_name}: max_abs={max_abs:.8f}, cosine={cosine:.9f}") + assert max_abs < 1e-2, case_name + assert cosine > 0.99999, case_name + assert_logits_close(onnx_logits, hf_logits, rtol=2e-2, atol=2e-2) + + attention_mask = hf_inputs["attention_mask"].numpy() + for row in range(attention_mask.shape[0]): + last_index = np.flatnonzero(attention_mask[row])[-1] + assert np.argmax(onnx_logits[row, last_index]) == np.argmax( + hf_logits[row, last_index] + ), case_name + finally: + decoder_session.close() + embedding_session.close() + vision_session.close() + + @pytest.mark.integration @pytest.mark.integration_fast def test_qwen35_deltanet_single_layer_parity():