# Understanding StatefulModule in Pocket‑TTS Streaming: The Architecture Behind Real‑Time Voice Synthesis

> Discover the role of StatefulModule in pocket-tts streaming. This abstract base class standardizes temporal state management for real-time voice synthesis, ensuring coherent audio generation.

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

---

**`StatefulModule` is the abstract base class in `kyutai-labs/pocket-tts` that standardizes how streaming layers manage temporal state, enabling frame‑by‑frame CPU audio generation while maintaining coherence across chunks through unified state initialization, access, and increment methods.**

Pocket‑TTS generates audio frame‑by‑frame on the CPU, which requires every neural network component to preserve context from previous steps. The `StatefulModule` class defined in [`pocket_tts/modules/stateful_module.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/stateful_module.py) provides the contractual glue that makes this deterministic streaming possible, allowing the `TTSModel` to treat the entire architecture as a single stateful entity.

## What Is StatefulModule?

`StatefulModule` is an abstract base class that establishes a common interface for any layer needing to remember information across inference steps. Unlike standard PyTorch modules that process inputs independently, streaming components—such as causal convolutions and attention KV‑caches—must persist buffers between forward passes.

By inheriting from this class, components agree to implement specific methods for state lifecycle management. This abstraction allows the top‑level generation loop to initialize, access, and advance state for all streaming layers without knowing their internal implementation details.

## Core Responsibilities of the StatefulModule Interface

The abstraction handles four critical responsibilities that enable deterministic streaming:

### State Initialization

The `init_state(batch_size, sequence_length)` method returns a dictionary of tensors representing the module’s fresh state. In [`pocket_tts/modules/conv.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/conv.py), the `StreamingConv1d.init_state` implementation creates a buffer of previous input samples (stored as `previous`) and a “first‑frame” flag to handle causal padding correctly on the initial chunk.

### State Access

During forward passes, `get_state(model_state)` extracts the module‑specific sub‑dictionary from the global `model_state` dictionary. For example, the `StreamingMultiheadAttention` layer in [`pocket_tts/modules/transformer.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/transformer.py) calls `self.get_state(model_state)` to retrieve its KV‑cache without exposing the entire model’s state to the attention mechanism.

### Step Progression

As generation advances, `increment_step(state, increment=1)` updates temporal counters and buffer offsets. The `StreamingMultiheadAttention.increment_step` method forwards this call to its KV‑cache backend, moving the RoPE (Rotary Position Embedding) offset forward so that positional encodings remain correct for the next frame.

### Integration Helpers

The module provides two utility functions that orchestrate state across the entire model hierarchy:

- **`init_states(model, …)`** walks the whole model, calling each `StatefulModule.init_state` and building a nested state dictionary.
- **`increment_steps(module, model_state, …)`** walks the model, invoking `increment_step` on every `StatefulModule` to advance counters by the number of frames generated.

The `TTSModel` class in [`pocket_tts/models/tts_model.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/tts_model.py) relies on these helpers to manage complex hierarchies of convolutions, transformers, and codec layers without hard‑coding layer‑specific logic.

## The Streaming Lifecycle: How StatefulModule Powers Generation

Understanding `StatefulModule` requires seeing how it fits into the three‑phase generation cycle:

1. **At start of generation** – `init_states` creates a fresh state for each streaming layer, including conv buffers, attention KV‑caches, and Mimi codec buffers.

2. **During generation** – The model yields a chunk of audio. Individual layers access their specific state via `get_state` and update internal buffers (such as the `previous` sample cache in `StreamingConv1d`) accordingly.

3. **After producing a chunk** – `increment_steps` is called to advance internal counters (e.g., the KV‑cache offset) by the number of frames generated, ensuring the next chunk continues from the correct temporal position.

When a new audio prompt is supplied, the cached states are cleared and re‑initialized, guaranteeing that voice‑cloning works correctly without leaking information from prior prompts.

## Practical Code Example

The following pattern illustrates the typical workflow inside Pocket‑TTS:

```python

# 1️⃣ Initialise the whole model’s streaming state (called inside TTSModel.generate)

model_state = init_states(tts_model, batch_size=1, sequence_length=0)

# 2️⃣ Run a streaming convolution layer (internal call, shown for illustration)

conv = StreamingConv1d(in_channels=80, out_channels=256, kernel_size=3, stride=1)
output = conv(x_chunk, model_state)   # conv accesses its own state via get_state()

# 3️⃣ After producing a chunk, advance all modules by the number of frames generated

increment_steps(tts_model, model_state, increment=generated_frames)

```

This cycle repeats until the full audio sequence is synthesized.

## Key Files and Implementation Details

Several critical components inherit from `StatefulModule` to enable streaming:

- **[`pocket_tts/modules/stateful_module.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/stateful_module.py)** – Defines the abstract base class, the `init_states` utility, and the `increment_steps` orchestrator.

- **[`pocket_tts/modules/conv.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/conv.py)** – Implements `StreamingConv1d` and `StreamingConvTranspose1d`, both of which manage causal convolution buffers as `StatefulModule` subclasses.

- **[`pocket_tts/modules/transformer.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/transformer.py)** – Provides `StreamingMultiheadAttention`, which uses the abstraction to maintain its KV‑cache and RoPE offsets across frames.

- **[`pocket_tts/models/tts_model.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/tts_model.py)** – Orchestrates the full generation pipeline, calling `init_states` before generation begins and `increment_steps` after each generated chunk.

## Summary

- **`StatefulModule`** provides the abstract interface that enables frame‑wise audio generation in Pocket‑TTS without consuming excessive memory.

- It standardizes **state initialization**, **access**, and **incrementation** across convolutional layers, attention mechanisms, and audio codecs.

- The utility functions **`init_states`** and **`increment_steps`** automate state management across complex model hierarchies, allowing the `TTSModel` to treat the architecture as a unified stateful entity.

- By inheriting from this base class, components like `StreamingConv1d` and `StreamingMultiheadAttention` maintain temporal coherence—managing buffers like the `previous` sample cache and KV‑cache—while keeping their implementation details encapsulated.

- This architecture supports **deterministic streaming** on CPU, allowing real‑time voice synthesis with proper handling of new prompts via state re‑initialization.

## Frequently Asked Questions

### What is the difference between StatefulModule and standard PyTorch modules?

Standard PyTorch modules process each input independently and do not maintain internal state between forward passes. `StatefulModule` extends this pattern by requiring implementations to manage persistent buffers—such as convolutional history or attention KV‑caches—through standardized `init_state` and `increment_step` methods, making it suitable for streaming applications where temporal context is essential.

### How does StatefulModule handle the KV‑cache in streaming attention?

The `StreamingMultiheadAttention` class inherits from `StatefulModule` and implements `init_state` to create empty KV tensors. During inference, `get_state` retrieves these tensors from the global state dictionary, which are then updated in‑place during the forward pass. `increment_step` advances the position offset used for RoPE embeddings, ensuring that each new frame receives correct positional encodings relative to the stream position rather than the chunk position.

### When should increment_steps be called during audio generation?

You should call `increment_steps` immediately after the model produces an audio chunk and before processing the next chunk. The `TTSModel` class in [`pocket_tts/models/tts_model.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/tts_model.py) handles this automatically, advancing all stateful components by the number of frames generated (typically the chunk size) to ensure that buffers, caches, and positional encodings remain synchronized with the output timeline.

### Does StatefulModule support batched streaming generation?

Yes, the `init_state` method accepts a `batch_size` parameter, allowing state dictionaries to contain batched tensors. This enables Pocket‑TTS to process multiple audio streams simultaneously, provided that all streams advance at the same rate. The `increment_steps` function handles batching transparently, updating counters for all sequences in the batch concurrently.