Text Decoder & Attention Mechanisms

Relevant source files

The following files were used as context for generating this wiki page:

• mlx_qwen3_asr/decoder.py
• mlx_qwen3_asr/attention.py
• mlx_qwen3_asr/generate.py
• mlx_qwen3_asr/config.py
• mlx_qwen3_asr/model.py
• CLAUDE.md

The Text Decoder and Attention Mechanisms form the autoregressive core of the mlx-qwen3-asr project. This system is responsible for generating text transcriptions by processing audio features injected into text embedding sequences. The architecture utilizes a Qwen3-style decoder with 28 layers, featuring specialized Interleaved Multi-RoPE (MRoPE) for positional encoding and Grouped Query Attention (GQA) for efficient inference.

The implementation is a ground-up MLX rewrite designed specifically for Apple Silicon (Metal). It prioritizes zero-dependency runtime execution by excluding transformers and torch from the core inference path. Key components include specialized attention helpers for Scaled Dot Product Attention (SDPA), SwiGLU activation functions, and KVCache management for linear-complexity incremental decoding.

Sources: CLAUDE.md:1-25, mlx_qwen3_asr/decoder.py:1-20

Decoder Architecture

The text decoder follows a transformer-based architecture optimized for autoregressive generation. In the 1.7B model variant, the decoder consists of 28 layers using RMSNorm and no bias in the linear layers, distinguishing it from the encoder which uses LayerNorm with bias.

Core Components

• Qwen3Decoder: The top-level module containing an embedding layer, 28 transformer blocks, and a final RMSNorm layer.
• Qwen3DecoderLayer: A single transformer block composed of an RMSNorm layer, a Qwen3Attention module, and an MLP (Multi-Layer Perceptron) using SwiGLU.
• RMSNorm: Used for layer normalization without learnable bias, ensuring stability during high-throughput inference.
• MLP: Implements the Gated Linear Unit with Swish activation (SwiGLU).

Sources: mlx_qwen3_asr/decoder.py:126-193, CLAUDE.md:65-68

The diagram shows the vertical stack of 28 decoder layers and the internal residual structure of a single Qwen3DecoderLayer.
Sources: mlx_qwen3_asr/decoder.py:126-150

Attention Mechanisms

The project utilizes Grouped Query Attention (GQA) to balance performance and memory efficiency. The attention mechanism is enhanced by Interleaved Multi-RoPE (MRoPE), which applies rotational positional embeddings across three dimensions (temporal, height, and width) using a stride-3 frequency assignment.

Multi-RoPE (MRoPE) Implementation

MRoPE is a critical correctness invariant. Unlike standard RoPE, it is interleaved rather than chunked. It divides the embedding dimensions into sections of [24, 20, 20] for the temporal, height, and width components respectively.

Scaled Dot Product Attention (SDPA)

The attention.py module provides a fused SDPA helper with a fallback mechanism. It prefers mx.fast.scaled_dot_product_attention when available on Apple Silicon but includes a manual implementation for compatibility and specific mask requirements.

Sources: mlx_qwen3_asr/attention.py:1-35, CLAUDE.md:65-72

Parameter	1.7B Model Value	0.6B Model Value
Attention Heads	16	16
KV Heads	8	8
Head Dimension	128	64
Intermediate Size	6144	3072
Hidden Size	2048	1024

Sources: mlx_qwen3_asr/config.py:75-120, tests/test_config.py:100-140

Generation and KV-Caching

The decoder performs autoregressive generation through two phases: Prefill and Decode.

1. Prefill Phase: Processes the initial input sequence (text prompt combined with audio features at <|audio_pad|> positions) and populates the KVCache.
2. Decode Phase: Generates tokens one-by-one, using the KVCache to store and retrieve previously computed Key and Value tensors, reducing complexity from $O(N^2)$ to $O(N)$.

The sequence shows how the KVCache is initialized during prefill and updated incrementally during each decoding step.
Sources: mlx_qwen3_asr/generate.py:110-165, mlx_qwen3_asr/model.py:100-150

Decode Termination Logic

The generation process continues until one of three conditions is met:

• EOS: An End-of-Sequence token is emitted (151643 or 151645).
• Length: The token budget defined by max_new_tokens is reached.
• Repetition: The system detects repetitive token patterns using a two-stage detector (consecutive single tokens or repeating patterns of 2-10 tokens).

Sources: mlx_qwen3_asr/generate.py:15-30, mlx_qwen3_asr/generate.py:280-315

Speculative Decoding

For optimized performance, the project supports Greedy Speculative Decoding. This involves a smaller "draft model" (e.g., 0.6B) that proposes a batch of tokens, which are then verified in parallel by the larger "target model" (e.g., 1.7B).

• Draft Width: Configured via num_draft_tokens (default: 4).
• Verification: The target model uses step_many to evaluate the draft tokens. If the target model disagrees with a draft token, the sequence is truncated at the mismatch, and the caches are rewound.

Sources: mlx_qwen3_asr/generate.py:195-260, mlx_qwen3_asr/cli.py:270-280

Configuration Parameters

The TextDecoderConfig dataclass manages the architectural hyperparameters of the decoder.

Key	Type	Default (1.7B)	Description
vocab_size	int	151936	Size of the vocabulary.
hidden_size	int	2048	Dimension of the hidden layers.
num_hidden_layers	int	28	Total number of transformer blocks.
num_attention_heads	int	16	Number of query heads.
num_key_value_heads	int	8	Number of key/value heads (GQA).
rope_theta	float	1000000.0	Base for RoPE frequency calculation.
tie_word_embeddings	bool	True	Whether to share weights between input and output embeddings.

Sources: mlx_qwen3_asr/config.py:78-105, tests/test_config.py:85-115

Conclusion

The Text Decoder and Attention Mechanisms in mlx-qwen3-asr provide a high-performance, Metal-optimized implementation of the Qwen3 architecture. By combining Grouped Query Attention with Interleaved Multi-RoPE and efficient KVCache management, the system achieves sub-200ms latency for speech-to-text tasks on Apple Silicon. The inclusion of speculative decoding further enhances throughput for larger model variants while maintaining zero-dependency runtime environments.

Sources: CLAUDE.md:5-20, mlx_qwen3_asr/decoder.py:1-15