mstar.model.qwen3_5.components.rope

Qwen3.5’s interleaved 3D MRoPE, over a partial rotary dim.

Model-side because the position resource serves only 1D RoPE. The three grids are interleaved [T,H,W,T,H,W,...] over the frequency dims, and only head_dim * partial_rotary_factor of each head is rotated (64 of 256 for every released size).

TODO: this duplicates qwen3_omni/components/rope.py apart from the partial rotation. Both belong in model/components/rope.py.

Functions

apply_partial_mrope(q, k, cos, sin[, ...])	Rotate the leading cos.shape[-1] dims of each head, pass the rest.
compute_3d_cos_sin(position_ids_3d, ...[, ...])	cos/sin of shape [seq_len, rotary_dim] from [3, seq_len] positions.
compute_inv_freq(rotary_dim, rope_theta[, ...])	1 / theta^(2i/d) over the rotated dims only; [rotary_dim // 2].
text_position_ids(seq_len[, start_pos, device])	[3, seq_len] for a pure-text span: all three grids advance together.
vision_position_advance(grid_thw, ...)	How far one image moves the position cursor.
vision_position_ids(grid_thw, spatial_merge_size)	[3, t * h' * w'] for one image, where h' = h // merge.

mstar.model.qwen3_5.components.rope.apply_partial_mrope(q, k, cos, sin, unsqueeze_dim=1)

Rotate the leading cos.shape[-1] dims of each head, pass the rest.

q/k are [tokens, heads, head_dim]; cos/sin are [tokens, rotary_dim] and broadcast over the head axis.

Parameters:

• q (Tensor)

• k (Tensor)

• cos (Tensor)

• sin (Tensor)

• unsqueeze_dim (int)

Return type:

tuple[Tensor, Tensor]

mstar.model.qwen3_5.components.rope.compute_3d_cos_sin(position_ids_3d, inv_freq, mrope_section, target_dtype=None)

cos/sin of shape [seq_len, rotary_dim] from [3, seq_len] positions.

Parameters:

• position_ids_3d (Tensor)

• inv_freq (Tensor)

• mrope_section (list[int])

• target_dtype (dtype | None)

Return type:

tuple[Tensor, Tensor]

mstar.model.qwen3_5.components.rope.compute_inv_freq(rotary_dim, rope_theta, device=None)

1 / theta^(2i/d) over the rotated dims only; [rotary_dim // 2].

Parameters:

• rotary_dim (int)

• rope_theta (float)

• device (device | None)

Return type:

Tensor

mstar.model.qwen3_5.components.rope.text_position_ids(seq_len, start_pos=0, device=None)

[3, seq_len] for a pure-text span: all three grids advance together.

Parameters:

• seq_len (int)

• start_pos (float)

• device (device | None)

Return type:

Tensor

mstar.model.qwen3_5.components.rope.vision_position_advance(grid_thw, spatial_merge_size)

How far one image moves the position cursor.

max(h', w') positions, not its t * h' * w' tokens: the grids run in parallel.

Parameters:

• grid_thw (Tensor)

• spatial_merge_size (int)

Return type:

int

mstar.model.qwen3_5.components.rope.vision_position_ids(grid_thw, spatial_merge_size, start_pos=0, device=None)

[3, t * h' * w'] for one image, where h' = h // merge.

T is constant across a frame while H and W sweep the merged patch grid, all based at start_pos; row-major over (t, h', w'), the merger’s token order.

Parameters:

• grid_thw (Tensor)

• spatial_merge_size (int)

• start_pos (float)

• device (device | None)

Return type:

Tensor