Chunk Size and KV caching in the CAT (Audio Tokenizer)

[patrickltobing]

Thank you very much for the great work. The potential of this framework is awesome.

I have question about the context mechanism in the Audio Tokenizer, though.
It seems that there is a chunk size in both encoder and decoder of the causal audio transformer (CAT).

What is the effective number of chunk size used during the training?
I tried different durations and looks like the effect is quite sensitive, both in the encoding of the prompt, and the decoding side.

P.S. Just to confirm, that basically the KV state is reset for each chunk, right?

[gyt1145028706]

Thank you for your attention !

1. Training Phase:
During training, we do not use chunking. The chunk_size is set to None, meaning each audio sequence is processed in a single forward pass.

2. Inference Phase:
For inference, the chunk_size is flexible:

• None: You can set it to None, but this may lead to Out-Of-Memory (OOM) errors if the sequence is very long.
• Multiples of 0.08: If specified, the value must be a multiple of 0.08 and should be $\le$ 10 seconds.
• A larger chunk_size (e.g., 8 seconds) will result in faster inference speeds but will consume more peak VRAM.

3. KV Cache and State:
To clarify your question about the KV state: the KV cache has a fixed size of 10 seconds. This means the model maintains a sliding context window of up to 10s rather than simply resetting the state for every chunk.

I hope this helps! Let us know if you have any other questions."

[patrickltobing]

OK, thank you for your response.
I'm discussing for now about the moss_tts_realtime, not sure whether the other models are different or not.

So, from my understanding, either the _encode_frame or _decode_frame will iterate over the list of chunk during encoding or decoding, respectively, and simply call the transformer module list.

I don't see any caching activated during this call, which means that the attention of each transformer is limited to each chunk.
In the training, it means that it always see the whole sequence, as chunk size is None, or limited by 10 s. as you mentioned.

However, in the streaming mode, we cannot wait until whole audio tokens sequence is generated, i.e., chunk size is set to a certain value, e.g., 0.96 s. (12 audio tokens). So there is an inconsistency.

So, here in modeling_moss_audio_tokenizer.py, _decode_frame line 1532, the call is simply with the input sequence and its length. This is called for each chunk, and without any cache input, it means that the attention is limited to each chunk.

Similarly, for the _encode_frame, line 1503. These mean that the number of chunk size is actually meaningless with respect to the KV caching history to the previous chunks. It simply records the attention within each chunk only.

Please correct me if I'm wrong.

[gyt1145028706]

Thank you for your detailed observation! I understand your concern—it might look like the attention is limited to each chunk because no cache is explicitly passed as an argument. However, the KV cache is actually maintained internally through the StreamingModule state management.

Here is how it works:

1. Persistent State within Context:
   In the encode and decode methods, we use a streaming() context manager (e.g., line 1612 and line 1732). When exit_stack.enter_context(module.streaming(...)) is called, it initializes a _streaming_state for every submodule. This state persists as long as the code remains inside the with block.

2. Internal KV Cache Update:
   Inside MossAudioTokenizerMultiheadAttention.forward (line 627), the code checks for self._streaming_state.

   • If it exists, it calls self._complete_kv(k, v).
   • This method interacts with RingKVCache, which updates its internal self.cache tensor with the new keys and values from the current chunk.
   • Because the Transformer layers (submodules) are the same objects across different chunk iterations in the loop, they retain this _streaming_state.

3. Sliding Window & Offsets:
   The RingKVCache is initialized with a fixed capacity (derived from the 10s context). After each chunk is processed, the state.offset is incremented (line 464). This ensures that the next chunk correctly attends to the previous keys/values stored in the ring buffer.

Summary:
The reason you don't see caching "activated" as an input/output of _encode_frame is that it's side-effect based. The _encode_frame call triggers the forward pass of layers that already "own" their KV cache buffers. As long as you are within the streaming() context, the KV history is automatically accumulated and maintained across chunks.

[patrickltobing]

Ah, yes, there is the KVCacheRing thing. Thanks for confirming.

But, I think I found the bug in the example_llm_stream_to_tts.py. So, here in line 299, the chunk_duration in the decode_kwargs is set to -1.

Here, in streaming_mossttsrealtime.py line 761, the chunk_duration_arg is set to None because override_f is -1.

Then, here, in modeling_moss_audio_tokenizer.py line 1702, the _decode_frame is called without the "cache streaming context" because the chunk_duration is None.

So, the decoding is happening in streaming_mossttsrealtime.py's audio_chunks(), and in line 710, it produces the chunk based on the chunk_frames value, which is set here.

Also note in line 711, you use chunk_duration=0.32, but this will be overridden by the decode_kwargs that sets it to -1.

Hence, in that example_llm_stream_to_tts.py, you will get different effects when varying the args.decode_chunk_frames because technically, the _decode_frame is always called outside of the streaming KV cache context, and so the attention in the CAT is limited to each chunk.

I put some prints in crucial lines and traced the logs during decoding, and indeed that is the case.

Let me know if this is not what you observed.

P.S.

• Also note that in modeling_moss_audio_tokenizer.py line 1725, this needs to be careful because if the codes_length is less than the chunk_frame_length, and it is in streaming mode, i.e., iteratively sending chunk segments, it will never go into the KV cache state below it.
• Relating to the above explanation, the codes_length is basically the length of the chunk segment, which is set by args.decode_chunk_frames in that example_llm_stream_to_tts.py, while the chunk_frame_length is set by the chunk_duration which is None and never actually went into that script. (Technically, this section is not executed but the one on line 1702).

[patrickltobing]

Sorry for subsequent message.

Following the above, I tested to change the decode_kwargs's chunk_duration to be > 0 and less than the chunk_frames. For instance 0.32 and 0.64 (8 frames), respectively.

In this case, the ExitStack() section is executed, however, this would go in to error, such as following

Traceback (most recent call last):
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/example_llm_stream_to_tts.py", line 419, in <module>
    main()
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/example_llm_stream_to_tts.py", line 414, in main
    write_wav(out_path, args.sample_rate, wav_chunks)
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/example_llm_stream_to_tts.py", line 122, in write_wav
    for chunk in chunks:
                 ^^^^^^
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/example_llm_stream_to_tts.py", line 174, in run_streaming_tts
    yield from decode_audio_frames(
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/example_llm_stream_to_tts.py", line 84, in decode_audio_frames
    for wav in decoder.audio_chunks():
               ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/mossttsrealtime/streaming_mossttsrealtime.py", line 744, in audio_chunks
    wav = self._decode(chunk_tokens)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/vertox/MOSS-TTS/moss_tts_realtime/mossttsrealtime/streaming_mossttsrealtime.py", line 809, in _decode
    decoded = self.codec.decode(tokens_t, chunk_duration=chunk_duration_arg, **decode_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/modeling_moss_audio_tokenizer.py", line 1760, in decode
    exit_stack.enter_context(decoder_module.streaming(batch_size=B))
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/modeling_moss_audio_tokenizer.py", line 175, in streaming
    self._start_streaming(batch_size, exit_stack)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/modeling_moss_audio_tokenizer.py", line 160, in _start_streaming
    self._apply_named_streaming(_start_streaming_fn)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/modeling_moss_audio_tokenizer.py", line 150, in _apply_named_streaming
    fn(name, child)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/modeling_moss_audio_tokenizer.py", line 155, in _start_streaming_fn
    raise RuntimeError(f"{name} is already streaming!")
RuntimeError:  is already streaming!

This is because in example_llm_stream_to_tts.py line 161, the decoding is run under the codec.streaming(), which initializes the streaming_state of all MossAudioTokenizerProjectedTransformer of the codec.

And, in modeling_moss_audio_tokenizer.py line 155 it raises the error as the streaming_state is already initialized as above.

Basically, the example_llm_stream_to_tts.py is running under chunk independent attention due to the chunk_duration set to -1 in the decode_kwargs of the AudioStreamDecoder. Then, if you set it to proper value that is less than the chunk_frames so that it runs under the ExitStack context for the KV caching, it will fail as above due to the above related lines.