Multi-RoPE (MRoPE) Implementation

Introduction

Multi-RoPE (MRoPE) is a specialized positional embedding mechanism utilized within the Qwen3-ASR architecture to handle multi-dimensional positional information. In this project, MRoPE is implemented using the MLX framework, optimized for Apple Silicon. Its primary purpose is to apply rotary positional embeddings across interleaved sections of head dimensions, specifically catering to temporal, height, and width dimensions often found in multimodal or complex sequence tasks.

The implementation ensures correctness by following a strict interleaved frequency assignment rather than a chunked approach. This is critical for maintaining parity with the original Qwen3-ASR model weights and ensuring accurate transcription output.

Sources: CLAUDE.md:88-89, mlx_qwen3_asr/mrope.py

Architecture and Core Logic

The MRoPE implementation revolves around the InterleavedMRoPE class, which manages the generation of cosine and sine embeddings. Unlike standard RoPE, MRoPE divides the head dimension into sections that correspond to different positional axes.

Interleaved Frequency Assignment

The implementation uses an interleaved assignment with a stride-3 frequency distribution. This means frequencies are not assigned in contiguous blocks but are distributed across dimensions to maintain structural information.

Key Components

Sources: mlx_qwen3_asr/mrope.py:10-50, tests/test_mrope.py:38-40

Implementation Details

Positional Embedding Computation

The InterleavedMRoPE class expects position_ids in the shape (B, 3, L), representing the Batch, three dimensions (Temporal, Height, Width), and Sequence Length. The embedding generation follows these steps:

1. Inverse Frequency Calculation: Computes the base frequencies using a power-law distribution.
2. Sectional Overwriting: The temporal frequencies serve as the base, which are then partially overwritten by Height and Width frequencies at indices defined by the mrope_section and a stride of 3.
3. Rotation: The resulting frequency tensor is concatenated with itself and transformed via cos and sin operations.

Sources: mlx_qwen3_asr/mrope.py:53-83, tests/test_mrope.py:90-105

Application to Attention

The embeddings are integrated into the text decoder's attention mechanism. As the decoder processes sequences, MRoPE is interleaved into the 28 Qwen3-style layers.

Technical Constraints and Invariants

To maintain model parity and avoid silently incorrect outputs, the following invariants must be respected:

• Interleaving: Sections must be interleaved using stride-3 frequency assignment across dimensions.
• Section Summation: The values in mrope_section must sum exactly to head_dim // 2.
• Computed Embeddings: Sinusoidal position embeddings must be computed at runtime as they are not stored in the weight files.
• Dtype Consistency: The positional embeddings should typically match the model's computation dtype (e.g., float16 or float32).

Sources: CLAUDE.md:92-96, tests/test_mrope.py:42-47, scripts/eval_nonnear_parity_matrix.py:145-155

Conclusion

The Multi-RoPE (MRoPE) implementation is a vital component of the mlx-qwen3-asr package, enabling the model to process complex positional data through interleaved frequency assignments. By utilizing the MLX framework, it provides a high-performance, Metal-optimized positional embedding system that ensures full parity with the original Qwen3-ASR architecture.