# StreamingTransformer with RoPE: How Pocket-TTS Enables Real-Time Speech Synthesis

> Discover how Pocket-TTS uses StreamingTransformer with RoPE for real-time speech synthesis. This custom architecture enables low-latency, incremental audio frame generation for efficient inference.

- Repository: [kyutai/pocket-tts](https://github.com/kyutai-labs/pocket-tts)
- Tags: deep-dive
- Published: 2026-07-11

---

**Pocket-TTS implements a custom StreamingTransformer with RoPE (Rotary Positional Embedding) that combines stateful key-value caching and causal attention masking to generate audio frames incrementally, enabling low-latency text-to-speech inference.**

The kyutai-labs/pocket-tts repository introduces a transformer architecture specifically engineered for streaming audio generation. Unlike standard transformers that process entire sequences simultaneously, the StreamingTransformer with RoPE consumes text embeddings frame-by-frame while maintaining contextual coherence through rotary positional embeddings and a sophisticated caching mechanism. This design targets real-time synthesis scenarios where latency and memory efficiency are critical constraints.

## Core Architecture Components

### StreamingMultiheadAttention with KV Caching

The attention mechanism in Pocket-TTS is implemented in [`pocket_tts/modules/transformer.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/transformer.py) as `StreamingMultiheadAttention`. This layer maintains a **stateful linear key-value cache** across generation steps rather than recomputing attention for the full sequence history at each forward pass.

The cache logic utilizes a `_LinearKVCacheBackend` that stores past keys and values in a contiguous tensor along with an offset tracking the current position. During each forward pass, the module retrieves cached tensors (or initializes fresh ones for the first step), concatenates them with newly computed `k` and `v` tensors, and updates the cache for subsequent calls. This approach reduces computational complexity from quadratic to linear relative to sequence length during generation.

### Rotary Positional Embedding (RoPE)

Positional information is injected via Rotary Positional Embedding (RoPE) defined in [`pocket_tts/modules/rope.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/rope.py). Rather than using absolute positional encodings that require fixed sequence lengths, the `RotaryEmbedding` class rotates query and key vectors using sinusoidal frequencies that depend on their relative positions.

The `apply_rope` function takes the current streaming offset into account when computing rotation angles, ensuring that positional representations remain consistent even when processing variable-length chunks. This is crucial for the `StreamingTransformer` because it processes audio one frame at a time while needing to maintain accurate relative positional relationships between distant tokens.

### StatefulModule Base Class

All stateful components inherit from `StatefulModule`, defined in [`pocket_tts/modules/stateful_module.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/stateful_module.py). This abstract base class provides a unified interface for modules requiring persistent state across generation steps through three key methods:

- **`init_state(batch_size, sequence_length)`** – Allocates cache tensors with appropriate dimensions
- **`get_state(model_state)`** – Retrieves the current cache from a shared state dictionary
- **`increment_step(model_state)`** – Advances the internal position offset after each generation step

The transformer uses these methods to coordinate cache initialization and timestep advancement across multiple layers without manual state management in the inference loop.

## Step-by-Step Streaming Mechanism

During inference, the `StreamingTransformer` executes the following operations for each audio frame:

1. **Linear Projection**: The input embedding is projected simultaneously into query, key, and value tensors using a single linear layer.
2. **RoPE Application**: The `rope_offset` retrieved from the state is passed to `RotaryEmbedding` to rotate the query and key tensors according to their absolute positions in the stream.
3. **Cache Retrieval**: The `_LinearKVCacheBackend` extracts previously cached keys and values or creates fresh tensors if processing the first frame.
4. **Causal Masking**: The `_build_attention_mask` function constructs a mask that restricts attention to current and past positions only. An optional `context` parameter can limit the attention window to a fixed number of past frames.
5. **Attention Computation**: The module calls `torch.nn.functional.scaled_dot_product_attention` with the prepared query, cached keys/values, and causal mask.
6. **Output Projection**: The attention output is projected back to the model dimension and returned for the next processing stage.

## Implementation Example

The following example demonstrates initializing a single streaming attention layer with RoPE and running one generation step:

```python
import torch
from pocket_tts.modules.transformer import StreamingMultiheadAttention
from pocket_tts.modules.rope import RotaryEmbedding

# Model hyperparameters

embed_dim = 256
num_heads = 8

# Instantiate RoPE and the attention layer

rope = RotaryEmbedding(max_period=10000.0)
attn = StreamingMultiheadAttention(embed_dim, num_heads, rope)

# Dummy input: batch-size 1, 1 time-step, embed_dim features

x = torch.randn(1, 1, embed_dim)

# Initialise per-layer state for a maximum sequence length of 1024

model_state = {"": attn.init_state(batch_size=1, sequence_length=1024)}

# Forward pass – the model_state argument provides the KV cache

out = attn(x, model_state)

# Advance the internal offset for the next step

attn.increment_step(model_state)

```

For a complete `StreamingTransformer` setup across multiple layers, you initialize the full model state using the `StatefulModule` interface and iterate through text embeddings:

```python
from pocket_tts.modules.transformer import StreamingTransformer
from pocket_tts.modules.stateful_module import init_states, increment_steps

# Configure the transformer

transformer = StreamingTransformer(
    d_model=256,
    n_head=8,
    n_layer=6,
    rope=RotaryEmbedding(),
    context=None,  # Use full causal context

)

# Initialise states for all layers

full_state = init_states(transformer, batch_size=1, sequence_length=1024)

# Generate audio frame-by-frame

for frame in text_embeddings:
    audio_chunk = transformer(frame, full_state)
    increment_steps(transformer, full_state)

```

## Summary

- **StreamingMultiheadAttention** in [`transformer.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/transformer.py) implements a stateful KV cache that enables linear-time incremental generation by reusing past key-value computations.
- **RoPE** in [`rope.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/rope.py) provides relative positional encoding through rotary embeddings that respect the current streaming offset, eliminating the need for absolute position vectors.
- **StatefulModule** in [`stateful_module.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/stateful_module.py) abstracts cache management through `init_state`, `get_state`, and `increment_step` methods, standardizing state handling across transformer layers.
- The architecture supports optional context windows through causal masking, allowing developers to balance computational cost against receptive field size during real-time synthesis.

## Frequently Asked Questions

### How does StreamingMultiheadAttention differ from standard multi-head attention?

Standard multi-head attention recomputes attention weights over the entire sequence history at every step, resulting in quadratic time complexity. `StreamingMultiheadAttention` uses a `_LinearKVCacheBackend` to store keys and values from previous steps, allowing each forward pass to process only the current frame while retrieving cached tensors from `model_state`. This reduces inference to linear complexity relative to sequence length.

### Why does Pocket-TTS use RoPE instead of absolute positional encodings?

Absolute positional encodings require fixed sequence lengths and can degrade when processing variable-length streams. The `RotaryEmbedding` implementation applies rotations based on relative positions calculated using the current `rope_offset`, making it naturally suited for the incremental, frame-by-frame processing pattern of the `StreamingTransformer` without requiring padding or recalculation of position vectors.

### What is the purpose of the StatefulModule base class?

`StatefulModule` provides a protocol for layers that must persist state across generation steps. It defines `init_state` for cache allocation, `get_state` for retrieval, and `increment_step` for position tracking. This abstraction allows the `StreamingTransformer` to manage complex multi-layer cache states through helper functions like `init_states` and `increment_steps` without exposing implementation details to the inference loop.

### How does the context parameter affect the transformer’s memory usage?

The `context` parameter in `StreamingMultiheadAttention` limits the attention window to a fixed number of past positions when set to an integer value. When `None`, the model attends to all previous positions (full causal context). Limiting context reduces the size of the KV cache and computational overhead for very long sequences, trading off some receptive field for improved memory efficiency during streaming generation.