[Perf][Qwen3-TTS]Cached Incremental Decode (#5202)

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -24,7 +24,7 @@ class _FakeDecoder(nn.Module):
         super().__init__()
         self.total_upsample = total_upsample
         self.decode_calls: list[dict[str, int]] = []
-        self.batched_decode_calls: list[dict[str, int]] = []
+        self.batched_decode_calls: list[dict[str, object]] = []
         self.decode_codes: list[torch.Tensor] = []
         self.cudagraph_calls: list[dict[str, int | torch.device]] = []
 
@@ -35,6 +35,7 @@ class _FakeDecoder(nn.Module):
         self,
         codes: torch.Tensor,
         *,
+        caches=None,
         chunk_size: int = 300,
         left_context_size: int = 25,
     ) -> torch.Tensor:
@@ -58,6 +59,7 @@ class _FakeDecoder(nn.Module):
         codes: torch.Tensor,
         lengths: list[int],
         *,
+        caches=None,
         chunk_size: int = 300,
         left_context_size: int = 25,
         max_batch_size: int = 0,
@@ -69,14 +71,28 @@ class _FakeDecoder(nn.Module):
                 "max_batch_size": max_batch_size,
                 "codes_shape": tuple(codes.shape),
                 "lengths": tuple(lengths),
+                "caches": caches,
             }
         )
+        if caches is not None:
+            outputs = []
+            for row, (length, cache) in enumerate(zip(lengths, caches, strict=True)):
+                output = self.chunked_decode(
+                    codes[row : row + 1, :, :length],
+                    caches=cache,
+                    chunk_size=chunk_size,
+                    left_context_size=left_context_size,
+                )
+                outputs.append(output[0])
+            return outputs
+
         batch = codes.shape[0]
         frames = codes.shape[-1]
         wav_len = frames * self.total_upsample + 6
         wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1)
         offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000
-        return wav.expand(batch, 1, wav_len) + offsets
+        outputs = wav.expand(batch, 1, wav_len) + offsets
+        return [outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)]
 
     def enable_cudagraph(self, **kwargs):
         self.cudagraph_calls.append(kwargs)
@@ -157,7 +173,7 @@ def _load_weights_noop(model: Qwen3TTSCode2Wav) -> set[str]:
         return model.load_weights(iter(()))
 
 
-def test_forward_trims_context_on_exact_frame_boundaries():
+def test_forward_uses_decoder_audio_contract_on_exact_frame_boundaries():
     model = _make_model()
 
     out = model.forward(
@@ -170,7 +186,7 @@ def test_forward_trims_context_on_exact_frame_boundaries():
     torch.testing.assert_close(audio, expected)
 
 
-def test_forward_trims_trailing_padding_without_context():
+def test_forward_uses_decoder_audio_contract_without_context():
     model = _make_model()
 
     out = model.forward(
@@ -183,24 +199,40 @@ def test_forward_trims_trailing_padding_without_context():
     torch.testing.assert_close(audio, expected)
 
 
-def test_forward_reuses_cached_ref_context_for_followup_chunk():
+@pytest.mark.parametrize("skipped_length", [0, 1], ids=["empty", "malformed"])
+def test_stateless_batch_uses_original_request_index_for_left_context(skipped_length):
     model = _make_model()
 
-    first_codes = torch.tensor(
-        [
-            9,
-            8,
-            1,
-            2,
-            9,
-            8,
-            3,
-            4,
+    skipped_request = torch.zeros(skipped_length, dtype=torch.long)
+    valid_request = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
+    out = model.forward(
+        input_ids=torch.cat([skipped_request, valid_request]),
+        seq_token_counts=[skipped_length, valid_request.numel()],
+        runtime_additional_information=[
+            {
+                "codes": {"audio": torch.zeros(skipped_length, dtype=torch.float32)},
+                "meta": {"left_context_size": 0},
+            },
+            {
+                "codes": {"audio": torch.zeros(8, dtype=torch.float32)},
+                "meta": {"left_context_size": 2},
+            },
         ],
-        dtype=torch.long,
     )
-    model.forward(
-        input_ids=first_codes,
+    wav_output = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(wav_output[0], torch.empty(0))
+    torch.testing.assert_close(wav_output[1], torch.arange(8, 16, dtype=torch.float32))
+
+
+def test_stateless_forward_trims_reference_prefix_from_decoder_output():
+    model = _make_model()
+
+    def _incremental_decode(codes, lengths, *, caches, **_kwargs):
+        return torch.arange(16, dtype=torch.float32).view(1, 1, -1)
+
+    model.decoder.batched_chunked_decode = _incremental_decode
+    out = model.forward(
+        input_ids=torch.arange(8, dtype=torch.long),
         runtime_additional_information=[
             {
                 "meta": {
@@ -213,104 +245,48 @@ def test_forward_reuses_cached_ref_context_for_followup_chunk():
         ],
     )
 
-    followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
+    audio = out.multimodal_outputs["model_outputs"][0]
+    torch.testing.assert_close(audio, torch.arange(8, 16, dtype=torch.float32))
+
+
+def test_request_aware_forward_passes_only_delta_frames_to_decoder():
+    model = _make_model(async_chunk=True)
+
     model.forward(
-        input_ids=followup_codes,
-        runtime_additional_information=[
-            {
-                "meta": {
-                    "left_context_size": 3,
-                    "ref_context_size": 2,
-                    "ref_context_request_id": "rid",
-                    "ref_context_included": False,
-                    "finished": torch.tensor(True),
-                }
+        input_ids=torch.arange(6, dtype=torch.long),
+        runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}],
+    )
+    model.forward(
+        input_ids=torch.arange(4, dtype=torch.long),
+        runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}],
+    )
+
+    assert [tuple(codes.shape) for codes in model.decoder.decode_codes[-2:]] == [
+        (1, _NUM_QUANTIZERS, 3),
+        (1, _NUM_QUANTIZERS, 2),
+    ]
+
+
+def test_four_frame_first_chunk_uses_configured_initial_size_without_ramp():
+    model = _make_model(
+        async_chunk=True,
+        stage_connector_config={
+            "extra": {
+                "codec_chunk_frames": 25,
+                "initial_codec_chunk_frames": 4,
             }
-        ],
+        },
+    )
+    _load_weights_noop(model)
+
+    model.forward(
+        input_ids=torch.arange(8, dtype=torch.long),
+        runtime_additional_information=[{"meta": {"request_id": "four-frame-request", "left_context_size": 0}}],
     )
 
-    torch.testing.assert_close(
-        model.decoder.decode_codes[-1],
-        torch.tensor(
-            [
-                [
-                    [9, 8, 5, 6, 7],
-                    [9, 8, 15, 16, 17],
-                ]
-            ],
-            dtype=torch.long,
-        ),
-    )
-    assert "rid" not in model._ref_context_cache
-
-
-def test_forward_fails_fast_when_ref_context_cache_is_missing():
-    model = _make_model()
-
-    followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
-    with pytest.raises(ValueError, match="Missing Qwen3-TTS ref context cache"):
-        model.forward(
-            input_ids=followup_codes,
-            runtime_additional_information=[
-                {
-                    "meta": {
-                        "left_context_size": 3,
-                        "ref_context_size": 2,
-                        "ref_context_request_id": "rid",
-                        "ref_context_included": False,
-                    }
-                }
-            ],
-        )
-
-
-def test_ref_context_cache_evicts_lru_entries_when_requests_abort():
-    model = _make_model()
-    model._ref_context_cache_max_entries = 2
-
-    first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
-    for rid in ("a", "b", "c"):
-        model.forward(
-            input_ids=first_codes,
-            runtime_additional_information=[
-                {
-                    "meta": {
-                        "left_context_size": 2,
-                        "ref_context_size": 2,
-                        "ref_context_request_id": rid,
-                        "ref_context_included": True,
-                    }
-                }
-            ],
-        )
-
-    assert list(model._ref_context_cache) == ["b", "c"]
-    assert model._ref_context_cache_bytes == sum(model._tensor_nbytes(t) for t in model._ref_context_cache.values())
-
-
-def test_ref_context_cache_evicts_to_byte_cap():
-    model = _make_model()
-    model._ref_context_cache_max_entries = 100
-    model._ref_context_cache_max_bytes = 33
-
-    first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
-    for rid in ("a", "b"):
-        model.forward(
-            input_ids=first_codes,
-            runtime_additional_information=[
-                {
-                    "meta": {
-                        "left_context_size": 2,
-                        "ref_context_size": 2,
-                        "ref_context_request_id": rid,
-                        "ref_context_included": True,
-                    }
-                }
-            ],
-        )
-
-    assert list(model._ref_context_cache) == ["b"]
-    assert model._ref_context_cache_bytes == 32
+    assert model.decoder._initial_codec_chunk_frames == 4
+    assert model.decoder._incremental_chunk_ramp == []
+    assert tuple(model.decoder.decode_codes[-1].shape) == (1, _NUM_QUANTIZERS, 4)
 
 
 def test_connector_codec_chunking_does_not_override_decode_chunking():
@@ -332,7 +308,7 @@ def test_connector_codec_chunking_does_not_override_decode_chunking():
 
     model.forward(
         input_ids=torch.arange(12, dtype=torch.long),
-        runtime_additional_information=[{"meta": {"left_context_size": 0}}],
+        runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}],
     )
 
     assert model.decoder.decode_calls[-1] == {
@@ -391,6 +367,66 @@ def test_forward_emits_zero_samples_for_empty_codec_payload():
     assert out.multimodal_outputs["model_outputs"][0].numel() == 0
 
 
+def test_dummy_runtime_information_provides_finished_request_state():
+    model = _make_model(async_chunk=True)
+
+    runtime_info = model.get_dummy_runtime_additional_information(2)
+
+    assert len(runtime_info) == 2
+    assert [info["meta"]["request_id"] for info in runtime_info] == [
+        "__qwen3_tts_dummy_run__0",
+        "__qwen3_tts_dummy_run__1",
+    ]
+    assert all(info["meta"]["finished"] for info in runtime_info)
+
+
+def test_dummy_request_state_is_marked_for_graph_stats_suppression():
+    model = _make_model(async_chunk=True)
+    runtime_info = model.get_dummy_runtime_additional_information(1)
+
+    model.forward(
+        input_ids=torch.arange(4, dtype=torch.long),
+        runtime_additional_information=runtime_info,
+    )
+
+    cache = model.decoder.batched_decode_calls[0]["caches"][0]
+    assert cache["_is_dummy_run"] is True
+    assert model._decoder_state_cache == {}
+
+
+def test_decoder_state_cache_capacity_warns_without_evicting_or_failing():
+    model = _make_model(async_chunk=True)
+    assert model._decoder_state_cache_warn_entries == 512
+    model._decoder_state_cache.update({f"active-{index}": {"prefix_frames": 0} for index in range(512)})
+
+    with patch("vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.logger.warning_once") as warning:
+        model.forward(
+            input_ids=torch.arange(4, dtype=torch.long),
+            runtime_additional_information=[{"meta": {"request_id": "new-request"}}],
+        )
+
+    warning.assert_called_once()
+    assert len(model._decoder_state_cache) == 513
+    assert "new-request" in model._decoder_state_cache
+
+
+def test_decoder_state_cache_uses_scheduler_request_id_for_finished_cleanup():
+    model = _make_model(async_chunk=True)
+
+    model.forward(
+        input_ids=torch.arange(4, dtype=torch.long),
+        request_ids=["internal-request-id"],
+        runtime_additional_information=[{"meta": {"request_id": "external-request-id"}}],
+    )
+
+    assert "internal-request-id" in model._decoder_state_cache
+    assert "external-request-id" not in model._decoder_state_cache
+
+    model.on_requests_finished(["internal-request-id"])
+
+    assert model._decoder_state_cache == {}
+
+
 def test_forward_batches_equal_length_requests_in_one_decoder_call():
     model = _make_model()
 
@@ -403,11 +439,15 @@ def test_forward_batches_equal_length_requests_in_one_decoder_call():
         ],
     )
 
-    assert model.decoder.decode_calls == [
+    assert model.decoder.decode_calls == []
+    assert model.decoder.batched_decode_calls == [
         {
             "chunk_size": 300,
             "left_context_size": 25,
+            "max_batch_size": 0,
             "codes_shape": (2, _NUM_QUANTIZERS, 3),
+            "lengths": (3, 3),
+            "caches": None,
         }
     ]
     audios = out.multimodal_outputs["model_outputs"]
@@ -415,10 +455,81 @@ def test_forward_batches_equal_length_requests_in_one_decoder_call():
     torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32))
 
 
-def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
+def test_forward_batches_request_aware_decoder_states():
+    model = _make_model(async_chunk=True)
+    calls = []
+
+    def _batched_chunked_decode(codes, lengths, caches, **kwargs):
+        calls.append((codes, lengths, caches))
+        return [torch.arange(16, dtype=torch.float32).view(1, -1) + row * 100 for row in range(2)]
+
+    model.decoder.batched_chunked_decode = _batched_chunked_decode
+    per_request = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
+    out = model.forward(
+        input_ids=torch.cat([per_request, per_request]),
+        seq_token_counts=[8, 8],
+        runtime_additional_information=[
+            {
+                "meta": {
+                    "left_context_size": 2,
+                    "ref_context_size": 2,
+                    "ref_context_request_id": request_id,
+                    "request_id": request_id,
+                    "ref_context_included": True,
+                }
+            }
+            for request_id in ("a", "b")
+        ],
+    )
+
+    assert len(calls) == 1
+    assert tuple(calls[0][0].shape) == (2, 2, 4)
+    assert calls[0][1] == [4, 4]
+    assert [cache["prefix_frames"] for cache in calls[0][2]] == [2, 2]
+    audios = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(audios[0], torch.arange(16, dtype=torch.float32))
+    torch.testing.assert_close(audios[1], torch.arange(100, 116, dtype=torch.float32))
+
+
+def test_forward_allows_followup_delta_without_reference_prefix():
+    model = _make_model(async_chunk=True)
+    first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
+    model.forward(
+        input_ids=first_codes,
+        runtime_additional_information=[
+            {
+                "meta": {
+                    "request_id": "rid",
+                    "ref_context_request_id": "rid",
+                    "ref_context_size": 2,
+                    "ref_context_included": True,
+                }
+            }
+        ],
+    )
+
+    model.forward(
+        input_ids=torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long),
+        runtime_additional_information=[
+            {
+                "meta": {
+                    "request_id": "rid",
+                    "ref_context_request_id": "rid",
+                    "ref_context_size": 2,
+                    "ref_context_included": False,
+                }
+            }
+        ],
+    )
+
+    assert [tuple(codes.shape) for codes in model.decoder.decode_codes[-2:]] == [
+        (1, _NUM_QUANTIZERS, 4),
+        (1, _NUM_QUANTIZERS, 3),
+    ]
+
+
+def test_forward_passes_variable_length_padded_batch_to_decoder():
     model = _make_model()
-    model._decode_batch_bucket_frames = [4]
-    model._decode_variable_chunk_batch_min_frames = 1
 
     out = model.forward(
         input_ids=torch.arange(12, dtype=torch.long),
@@ -437,6 +548,7 @@ def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
             "max_batch_size": 0,
             "codes_shape": (2, _NUM_QUANTIZERS, 4),
             "lengths": (2, 4),
+            "caches": None,
         }
     ]
     audios = out.multimodal_outputs["model_outputs"]
@@ -444,108 +556,6 @@ def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
     torch.testing.assert_close(audios[1], torch.arange(1008, 1016, dtype=torch.float32))
 
 
-def test_forward_bucket_batches_different_length_requests_and_trims_rows():
-    model = _make_model()
-    model._decode_batch_bucket_frames = [4]
-
-    out = model.forward(
-        input_ids=torch.arange(18, dtype=torch.long),
-        seq_token_counts=[4, 6, 8],
-        runtime_additional_information=[
-            {"meta": {"left_context_size": 0}},
-            {"meta": {"left_context_size": 1}},
-            {"meta": {"left_context_size": 2}},
-        ],
-    )
-
-    assert model.decoder.decode_calls == [
-        {
-            "chunk_size": 300,
-            "left_context_size": 25,
-            "codes_shape": (3, _NUM_QUANTIZERS, 4),
-        }
-    ]
-    assert model.decoder.batched_decode_calls == []
-    audios = out.multimodal_outputs["model_outputs"]
-    torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
-    torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32))
-    torch.testing.assert_close(audios[2], torch.arange(2008, 2016, dtype=torch.float32))
-
-
-def test_forward_bucket_pads_only_to_group_max_frame_length():
-    model = _make_model()
-    model._decode_batch_bucket_frames = [8]
-
-    out = model.forward(
-        input_ids=torch.arange(10, dtype=torch.long),
-        seq_token_counts=[4, 6],
-        runtime_additional_information=[
-            {"meta": {"left_context_size": 0}},
-            {"meta": {"left_context_size": 0}},
-        ],
-    )
-
-    assert model.decoder.decode_calls == [
-        {
-            "chunk_size": 300,
-            "left_context_size": 25,
-            "codes_shape": (2, _NUM_QUANTIZERS, 3),
-        }
-    ]
-    assert model.decoder.batched_decode_calls == []
-    audios = out.multimodal_outputs["model_outputs"]
-    torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
-    torch.testing.assert_close(audios[1], torch.arange(1000, 1012, dtype=torch.float32))
-
-
-def test_forward_splits_bucket_groups_by_configured_max_batch_size():
-    model = _make_model()
-    model._decode_batch_bucket_frames = [4]
-    model._decode_batch_max_size = 2
-
-    model.forward(
-        input_ids=torch.arange(18, dtype=torch.long),
-        seq_token_counts=[4, 6, 8],
-        runtime_additional_information=[
-            {"meta": {"left_context_size": 0}},
-            {"meta": {"left_context_size": 0}},
-            {"meta": {"left_context_size": 0}},
-        ],
-    )
-
-    assert model.decoder.decode_calls == [
-        {
-            "chunk_size": 300,
-            "left_context_size": 25,
-            "codes_shape": (2, _NUM_QUANTIZERS, 3),
-        },
-        {
-            "chunk_size": 300,
-            "left_context_size": 25,
-            "codes_shape": (1, _NUM_QUANTIZERS, 4),
-        },
-    ]
-    assert model.decoder.batched_decode_calls == []
-
-
-def test_forward_does_not_pad_singleton_bucket_group():
-    model = _make_model()
-    model._decode_batch_bucket_frames = [4]
-
-    model.forward(
-        input_ids=torch.arange(4, dtype=torch.long),
-        runtime_additional_information=[{"meta": {"left_context_size": 0}}],
-    )
-
-    assert model.decoder.decode_calls == [
-        {
-            "chunk_size": 300,
-            "left_context_size": 25,
-            "codes_shape": (1, _NUM_QUANTIZERS, 2),
-        }
-    ]
-
-
 def test_decode_chunking_override_is_passed_to_cudagraph():
     model = _make_model(
         async_chunk=True,
@@ -563,27 +573,26 @@ def test_decode_chunking_override_is_passed_to_cudagraph():
     _load_weights_noop(model)
 
     assert model.decoder.cudagraph_calls[-1] == {
-        "capture_sizes": None,
         "capture_batch_sizes": None,
-        "extra_capture_shapes": None,
-        "compile_shapes": None,
+        "stateless_capture_sizes": None,
         "device": torch.device("cuda"),
         "codec_chunk_frames": 25,
         "codec_left_context_frames": 72,
+        "initial_codec_chunk_frames": 1,
+        "codec_chunk_ramp": None,
+        "async_chunk": True,
         "decode_chunk_size": 400,
         "decode_left_context": 17,
     }
 
 
-def test_cudagraph_capture_shapes_can_be_configured():
+def test_cudagraph_batch_config_is_preserved():
     model = _make_model(
         async_chunk=True,
         device=torch.device("cuda"),
         stage_connector_config={
             "extra": {
-                "decode_cudagraph_capture_sizes": "97,325",
                 "decode_cudagraph_batch_sizes": [1, 2, 4, 8],
-                "decode_cudagraph_extra_capture_shapes": ["3:325", [5, 325]],
             }
         },
     )
@@ -591,26 +600,39 @@ def test_cudagraph_capture_shapes_can_be_configured():
     _load_weights_noop(model)
 
     call = model.decoder.cudagraph_calls[-1]
-    assert call["capture_sizes"] == [97, 325]
     assert call["capture_batch_sizes"] == [1, 2, 4, 8]
-    assert call["extra_capture_shapes"] == [(3, 325), (5, 325)]
 
 
-def test_decode_compile_shapes_can_be_configured():
+def test_stateless_cudagraph_capture_sizes_are_passed_from_config():
+    model = _make_model(
+        async_chunk=False,
+        device=torch.device("cuda"),
+        stage_connector_config={
+            "extra": {
+                "decode_cudagraph_capture_sizes": [97, 400],
+            }
+        },
+    )
+
+    _load_weights_noop(model)
+
+    assert model.decoder.cudagraph_calls[-1]["stateless_capture_sizes"] == [97, 400]
+
+
+def test_async_cudagraph_ignores_stateless_capture_sizes():
     model = _make_model(
         async_chunk=True,
         device=torch.device("cuda"),
         stage_connector_config={
             "extra": {
-                "decode_compile_shapes": ["1:325", [1, 73]],
+                "decode_cudagraph_capture_sizes": [97, 400],
             }
         },
     )
 
     _load_weights_noop(model)
 
-    call = model.decoder.cudagraph_calls[-1]
-    assert call["compile_shapes"] == [(1, 73), (1, 325)]
+    assert model.decoder.cudagraph_calls[-1]["stateless_capture_sizes"] is None
 
 
 def test_decode_tf32_can_be_configured():
@@ -642,23 +664,19 @@ def test_decode_tf32_can_be_configured():
         torch.set_float32_matmul_precision(old_matmul_precision)
 
 
-def test_decode_batch_bucket_frames_can_be_configured():
+def test_decode_batch_max_size_can_be_configured():
     model = _make_model(
         async_chunk=True,
         stage_connector_config={
             "extra": {
-                "decode_batch_bucket_frames": "73,169",
                 "decode_batch_max_size": 10,
-                "decode_variable_chunk_batch_min_frames": 512,
             }
         },
     )
 
     _load_weights_noop(model)
 
-    assert model._decode_batch_bucket_frames == [73, 169]
     assert model._decode_batch_max_size == 10
-    assert model._decode_variable_chunk_batch_min_frames == 512
 
 
 def test_invalid_decode_batch_max_size_is_rejected():

tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py

@@ -1,125 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import inspect
-
-import pytest
-import torch
-from transformers.masking_utils import create_sliding_window_causal_mask
-
-from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
-    Qwen3TTSTokenizerV2DecoderConfig,
-)
-from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
-    Qwen3TTSTokenizerV2DecoderTransformerModel,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_config(attention_implementation: str) -> Qwen3TTSTokenizerV2DecoderConfig:
-    config = Qwen3TTSTokenizerV2DecoderConfig(
-        hidden_size=32,
-        latent_dim=16,
-        max_position_embeddings=512,
-        num_attention_heads=4,
-        num_key_value_heads=4,
-        intermediate_size=64,
-        num_hidden_layers=1,
-        num_quantizers=2,
-        decoder_dim=32,
-        upsample_rates=(2,),
-        upsampling_ratios=(2,),
-    )
-    config._attn_implementation = attention_implementation
-    return config
-
-
-def _make_model(
-    attention_implementation: str,
-    dtype: torch.dtype = torch.float32,
-) -> Qwen3TTSTokenizerV2DecoderTransformerModel:
-    model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(_make_config(attention_implementation))
-    return model.to(dtype=dtype).eval()
-
-
-def _create_reference_mask(
-    model: Qwen3TTSTokenizerV2DecoderTransformerModel,
-    inputs_embeds: torch.Tensor,
-) -> torch.Tensor | None:
-    sequence_length = inputs_embeds.shape[1]
-    mask_kwargs = {
-        "config": model.config,
-        "attention_mask": None,
-        "past_key_values": None,
-        "position_ids": None,
-    }
-    sig = inspect.signature(create_sliding_window_causal_mask)
-    if "input_embeds" in sig.parameters:
-        mask_kwargs["input_embeds"] = inputs_embeds[:1]
-        mask_kwargs["cache_position"] = torch.arange(sequence_length)
-    else:
-        mask_kwargs["inputs_embeds"] = inputs_embeds[:1]
-    return create_sliding_window_causal_mask(**mask_kwargs)
-
-
-@pytest.mark.parametrize("sequence_length", [17, 72, 73, 169, 325])
-@pytest.mark.parametrize(
-    ("attention_implementation", "dtype"),
-    [("sdpa", torch.float32), ("eager", torch.float32), ("eager", torch.float16)],
-)
-def test_cached_sliding_mask_matches_transformers(
-    sequence_length: int,
-    attention_implementation: str,
-    dtype: torch.dtype,
-):
-    model = _make_model(attention_implementation, dtype)
-    inputs_embeds = torch.empty(2, sequence_length, model.config.hidden_size, dtype=dtype)
-
-    actual = model._get_sliding_attention_mask(inputs_embeds)
-    expected = _create_reference_mask(model, inputs_embeds)
-
-    if expected is None:
-        assert actual is None
-        assert model._get_sliding_attention_mask(inputs_embeds) is None
-        return
-
-    assert actual is not None
-    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
-    assert actual.shape == (1, 1, sequence_length, sequence_length)
-    assert actual.is_contiguous()
-    assert model._get_sliding_attention_mask(inputs_embeds) is actual
-
-
-def test_cached_eager_mask_tracks_compute_dtype():
-    model = _make_model("eager")
-
-    fp32_inputs = torch.empty(1, 73, model.config.hidden_size)
-    fp32_mask = model._get_sliding_attention_mask(fp32_inputs)
-    assert fp32_mask is not None
-    assert fp32_mask.dtype == torch.float32
-
-    model.half()
-    fp16_inputs = fp32_inputs.half()
-    fp16_mask = model._get_sliding_attention_mask(fp16_inputs)
-    assert fp16_mask is not None
-    assert fp16_mask.dtype == torch.float16
-    assert torch.isfinite(fp16_mask).all()
-    assert fp16_mask.min() == torch.finfo(torch.float16).min
-
-
-@pytest.mark.parametrize(
-    ("kwargs", "message"),
-    [
-        ({"use_cache": True}, "use_cache=True"),
-        ({"past_key_values": object()}, "past_key_values"),
-        ({"cache_position": torch.arange(1, 18)}, "zero-based cache_position"),
-        ({"position_ids": torch.arange(1, 18).unsqueeze(0)}, "zero-based position_ids"),
-    ],
-)
-def test_cached_mask_rejects_unsupported_cache_and_positions(kwargs: dict, message: str):
-    model = _make_model("sdpa")
-    inputs_embeds = torch.randn(1, 17, model.config.latent_dim)
-
-    with pytest.raises(ValueError, match=message):
-        model(inputs_embeds=inputs_embeds, **kwargs)

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -788,10 +788,16 @@ def test_qwen3_tts_code2wav_forward_decodes_connector_payload() -> None:
             self.last_codes = codes.detach().clone()
             return codes.sum(dim=1).to(torch.float32)
 
+        def batched_chunked_decode(self, codes, lengths, caches, **kwargs):
+            assert lengths == [codes.shape[-1]]
+            assert caches is None
+            return self.chunked_decode(codes, **kwargs)
+
     decoder = _Decoder()
     model = Qwen3TTSCode2Wav.__new__(Qwen3TTSCode2Wav)
     torch.nn.Module.__init__(model)
     model.decoder = decoder
+    model._async_chunk = False
     model._num_quantizers = 2
     model._total_upsample = 1
     model._output_sample_rate = 24000

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -156,8 +156,9 @@ def test_streaming_phases(config, n_frames, finished, expected):
     else:
         exp_ctx, exp_window = expected
         assert payload is not None
-        assert payload.meta.left_context_size == exp_ctx
-        assert len(payload.codes.audio) == _Q * exp_window
+        assert payload.meta.left_context_size == 0
+        expected_delta = exp_window if finished else exp_window - exp_ctx
+        assert len(payload.codes.audio) == _Q * expected_delta
 
 
 def test_dynamic_ic_adapts_to_load():
@@ -202,7 +203,8 @@ def test_ic_load_change_mid_request():
     assert _call(tm, "r", n_frames=25) is None
     p3 = _call(tm, "r", n_frames=27)
     assert p3 is not None
-    assert p3.meta.left_context_size == 2
+    assert p3.meta.left_context_size == 0
+    assert len(p3.codes.audio) == _Q * 25
 
     # A *new* request under high load gets IC=16 (not IC=2).
     # Frame 2 would emit under IC=2 but must hold under IC=16.
@@ -227,7 +229,8 @@ def test_connector_initial_chunk_config_overrides_dynamic_ic():
     assert _call(tm, "r", n_frames=25) is None
     p2 = _call(tm, "r", n_frames=29)
     assert p2 is not None
-    assert p2.meta.left_context_size == 4
+    assert p2.meta.left_context_size == 0
+    assert len(p2.codes.audio) == _Q * 25
 
 
 @pytest.mark.parametrize(
@@ -282,8 +285,8 @@ def test_first_streaming_chunk_prepends_ref_code_context():
     assert len(payload.codes.audio) == _Q * 12
 
 
-def test_followup_ref_code_context_is_sent_as_metadata_handle():
-    """Follow-up chunks keep full ref context semantically without resending it."""
+def test_followup_sends_only_codec_delta_without_ref_metadata():
+    """Follow-up chunks rely on decoder state and resend neither reference codes nor metadata."""
     tm = _tm()
     rid = "r-ref2"
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(35)]
@@ -299,13 +302,11 @@ def test_followup_ref_code_context_is_sent_as_metadata_handle():
     )
 
     assert payload is not None
-    # ref_code (2 frames) is represented in metadata so Code2Wav can restore it
-    # from its request-local cache. It must not be resent in codes.audio.
-    assert payload.meta.left_context_size == 10 + 2
-    assert payload.meta.ref_context_size == 2
-    assert payload.meta.ref_context_request_id == rid
-    assert payload.meta.ref_context_included is False
-    assert len(payload.codes.audio) == _Q * 35
+    assert payload.meta.left_context_size == 0
+    assert payload.meta.ref_context_size is None
+    assert payload.meta.ref_context_request_id is None
+    assert payload.meta.ref_context_included is None
+    assert len(payload.codes.audio) == _Q * 25
 
 
 def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
@@ -332,8 +333,8 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
     )
 
     assert payload is not None
-    assert payload.meta.left_context_size == 3 + 3
-    assert len(payload.codes.audio) == _Q * (3 + 3 + 4)
+    assert payload.meta.left_context_size == 3
+    assert len(payload.codes.audio) == _Q * (3 + 4)
     frames = payload.codes.audio.reshape(_Q, -1).transpose(0, 1)
     torch.testing.assert_close(frames[:3], ref_code[-3:])
 
@@ -700,36 +701,40 @@ class TestChunkRampEmission:
         assert self._emit(tm, rid, 4) is None
         p1 = self._emit(tm, rid, 5)
         assert p1 is not None
-        assert len(p1.codes.audio) == _Q * 5
-        assert p1.meta.left_context_size == 1
+        assert len(p1.codes.audio) == _Q * 4
+        assert p1.meta.left_context_size == 0
         tm.ramp_chunk_count[rid] = 2
 
         for n in range(6, 13):
             assert self._emit(tm, rid, n) is None
         p2 = self._emit(tm, rid, 13)
         assert p2 is not None
-        assert p2.meta.left_context_size == 5
+        assert p2.meta.left_context_size == 0
+        assert len(p2.codes.audio) == _Q * 8
         tm.ramp_chunk_count[rid] = 3
 
         for n in range(14, 29):
             assert self._emit(tm, rid, n) is None
         p3 = self._emit(tm, rid, 29)
         assert p3 is not None
-        assert p3.meta.left_context_size == 13
+        assert p3.meta.left_context_size == 0
+        assert len(p3.codes.audio) == _Q * 16
         tm.ramp_chunk_count[rid] = 4
 
         for n in range(30, 54):
             assert self._emit(tm, rid, n) is None
         p4 = self._emit(tm, rid, 54)
         assert p4 is not None
-        assert p4.meta.left_context_size == 25
+        assert p4.meta.left_context_size == 0
+        assert len(p4.codes.audio) == _Q * 25
         tm.ramp_chunk_count[rid] = 5
 
         for n in range(55, 79):
             assert self._emit(tm, rid, n) is None
         p5 = self._emit(tm, rid, 79)
         assert p5 is not None
-        assert p5.meta.left_context_size == 25
+        assert p5.meta.left_context_size == 0
+        assert len(p5.codes.audio) == _Q * 25
 
     def test_ramp_finished_flush_mid_chunk(self):
         tm = _tm(chunk_ramp=self.RAMP)
@@ -742,8 +747,8 @@ class TestChunkRampEmission:
         p_fin = self._emit(tm, rid, 3, finished=True)
         assert p_fin is not None
         assert p_fin.meta.finished.item() is True
-        assert len(p_fin.codes.audio) == _Q * (1 + 2)
-        assert p_fin.meta.left_context_size == 1
+        assert len(p_fin.codes.audio) == _Q * 2
+        assert p_fin.meta.left_context_size == 0
 
     def test_ramp_finished_no_new_frames(self):
         tm = _tm(chunk_ramp=self.RAMP)
@@ -783,7 +788,8 @@ class TestChunkRampEmission:
             assert self._emit(tm, rid, n) is None
         p1 = self._emit(tm, rid, 26)
         assert p1 is not None
-        assert p1.meta.left_context_size == 1
+        assert p1.meta.left_context_size == 0
+        assert len(p1.codes.audio) == _Q * 25
 
     def test_ramp_with_ref_code_first_chunk(self):
         tm = _tm(chunk_ramp=self.RAMP, left_context=25)
@@ -820,14 +826,16 @@ class TestChunkRampEmission:
             assert self._emit(tm, rid, n) is None
         p2 = self._emit(tm, rid, 30)
         assert p2 is not None
-        assert p2.meta.left_context_size == 5
+        assert p2.meta.left_context_size == 0
+        assert len(p2.codes.audio) == _Q * 25
         tm.ramp_chunk_count[rid] = 3
 
         for n in range(31, 55):
             assert self._emit(tm, rid, n) is None
         p3 = self._emit(tm, rid, 55)
         assert p3 is not None
-        assert p3.meta.left_context_size == 25
+        assert p3.meta.left_context_size == 0
+        assert len(p3.codes.audio) == _Q * 25
 
     def test_ramp_profile_4_4_8_16_25(self):
         """Ramp [4,4,8,16,25]: chunk 0=4 frames (320ms audio covers chunk 1
@@ -846,35 +854,40 @@ class TestChunkRampEmission:
             assert self._emit(tm, rid, n) is None
         p1 = self._emit(tm, rid, 8)
         assert p1 is not None
-        assert p1.meta.left_context_size == 4
+        assert p1.meta.left_context_size == 0
+        assert len(p1.codes.audio) == _Q * 4
         tm.ramp_chunk_count[rid] = 2
 
         for n in range(9, 16):
             assert self._emit(tm, rid, n) is None
         p2 = self._emit(tm, rid, 16)
         assert p2 is not None
-        assert p2.meta.left_context_size == 8
+        assert p2.meta.left_context_size == 0
+        assert len(p2.codes.audio) == _Q * 8
         tm.ramp_chunk_count[rid] = 3
 
         for n in range(17, 32):
             assert self._emit(tm, rid, n) is None
         p3 = self._emit(tm, rid, 32)
         assert p3 is not None
-        assert p3.meta.left_context_size == 16
+        assert p3.meta.left_context_size == 0
+        assert len(p3.codes.audio) == _Q * 16
         tm.ramp_chunk_count[rid] = 4
 
         for n in range(33, 57):
             assert self._emit(tm, rid, n) is None
         p4 = self._emit(tm, rid, 57)
         assert p4 is not None
-        assert p4.meta.left_context_size == 25
+        assert p4.meta.left_context_size == 0
+        assert len(p4.codes.audio) == _Q * 25
         tm.ramp_chunk_count[rid] = 5
 
         for n in range(58, 82):
             assert self._emit(tm, rid, n) is None
         p5 = self._emit(tm, rid, 82)
         assert p5 is not None
-        assert p5.meta.left_context_size == 25
+        assert p5.meta.left_context_size == 0
+        assert len(p5.codes.audio) == _Q * 25
 
     def test_ramp_resets_on_segment_boundary(self):
         """After segment 1 emits chunks, a segment boundary clears

vllm_omni/deploy/qwen3_tts.yaml

@@ -33,11 +33,19 @@ connectors:
       # initial_codec_chunk_frames (including the per-request API field).
       # Example:
       #   codec_chunk_ramp: [4, 4, 8, 16, 25]
-      # Common Stage1 decode buckets:
-          #   no-ref first/steady chunks: 25 / 97 frames
-          #   Base ref-context first/steady chunks: 73 / 169 frames
-          #   decoder internal non-streaming chunks: 325 frames
-      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]
+      # Decoder CUDA Graph sizing:
+      #   * async_chunk=true derives all stateful ICL/xvec shapes from
+      #     codec_chunk_ramp when enabled, otherwise from
+      #     initial_codec_chunk_frames, codec_chunk_frames, and sliding_window;
+      #     decode_cudagraph_capture_sizes is ignored in this mode.
+      #   * async_chunk=false captures stateless shapes 150 and
+      #     decode_chunk_frames + decode_left_context_frames (325 by default).
+      #     Extra stateless shapes may be added (union, not intersection) with:
+      #       decode_cudagraph_capture_sizes: [97, 400]
+      # Graph batch size defaults to 1. Add larger buckets only when needed:
+      #   decode_cudagraph_batch_sizes: [1, 2, 4]
+      # Limit eager batches to reduce temporary stateful cache allocations; CUDA Graph batching uses the buckets above.
+      decode_batch_max_size: 4
 
 stages:
   - stage_id: 0

vllm_omni/deploy/qwen3_tts_high_concurrency.yaml

@@ -38,16 +38,12 @@ connectors:
       # initial_codec_chunk_frames (including the per-request API field).
       # Example:
       #   codec_chunk_ramp: [4, 4, 8, 16, 25]
-      # Common Stage1 decode buckets:
-      #   no-ref first/steady chunks: 25 / 97 frames
-      #   Base ref-context first/steady chunks: 73 / 169 frames
-      #   decoder internal non-streaming chunks: 325 frames
-      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]
-      # Keep B>1 captures opt-in; c64 e2e validation did not show a stable win.
+      # async_chunk derives stateful ICL/xvec graph shapes from the ramp when
+      # enabled, otherwise from the fixed chunk settings above, and ignores
+      # decode_cudagraph_capture_sizes. Keep B>1
+      # captures opt-in; c64 e2e validation did not show a stable win.
       decode_cudagraph_batch_sizes: [1]
-      decode_compile_shapes: []
       decode_batch_max_size: 1
-      decode_batch_bucket_frames: []
       decode_enable_tf32: false
 
 stages:

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from collections import Counter, OrderedDict
+from collections import Counter
 from collections.abc import Iterable
 from typing import Any
 
@@ -14,6 +14,7 @@ from vllm.model_executor.model_loader import DefaultModelLoader
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
 
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
     Qwen3TTSTokenizerV2Config,
@@ -24,8 +25,7 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
 
 logger = init_logger(__name__)
 
-_REF_CONTEXT_CACHE_MAX_ENTRIES = 4096
-_REF_CONTEXT_CACHE_MAX_BYTES = 64 * 1024 * 1024
+_DUMMY_REQUEST_ID = "__qwen3_tts_dummy_run__"
 
 
 def _codec_ids_from_payload_or_input(
@@ -56,10 +56,16 @@ class Qwen3TTSCode2Wav(nn.Module):
 
     input_modalities = "audio"
 
+    # Ask the model runner for the scheduler-side request IDs. Stateful
+    # decoder caches must use the same IDs delivered by on_requests_finished;
+    # payload metadata carries an external ID which may differ.
+    requires_request_ids = True
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.vllm_config = vllm_config
         self.model_path = vllm_config.model_config.model
+        self._async_chunk = bool(getattr(vllm_config.model_config, "async_chunk", False))
 
         self.have_multimodal_outputs = True
         self.has_preprocess = False
@@ -69,9 +75,7 @@ class Qwen3TTSCode2Wav(nn.Module):
 
         self._decode_chunk_frames = 300
         self._decode_left_context_frames = 25
-        self._decode_batch_bucket_frames: list[int] = []
         self._decode_batch_max_size = 0
-        self._decode_variable_chunk_batch_min_frames = self._decode_chunk_frames + self._decode_left_context_frames + 1
         self._logged_codec_stats = False
         self._logged_malformed_codec_lengths: set[tuple[int, int]] = set()
         self._batch_stats_enabled = os.environ.get("VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS", "").lower() in (
@@ -103,10 +107,8 @@ class Qwen3TTSCode2Wav(nn.Module):
         self._output_sample_rate = int(tok_config.output_sample_rate)
         self._total_upsample = int(self.decoder.total_upsample)
         self._decoder_sliding_window = int(getattr(dec_config, "sliding_window", 0) or 0)
-        self._ref_context_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
-        self._ref_context_cache_bytes = 0
-        self._ref_context_cache_max_entries = _REF_CONTEXT_CACHE_MAX_ENTRIES
-        self._ref_context_cache_max_bytes = _REF_CONTEXT_CACHE_MAX_BYTES
+        self._decoder_state_cache: dict[str, dict[str, Any]] = {}
+        self._decoder_state_cache_warn_entries = 512
 
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
         # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner.
@@ -114,6 +116,13 @@ class Qwen3TTSCode2Wav(nn.Module):
             return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
         return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
 
+    def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, Any]]:
+        """Provide request state metadata for vLLM's external graph dummy run."""
+        return [
+            {"meta": {"request_id": f"{_DUMMY_REQUEST_ID}{index}", "finished": True}}
+            for index in range(max(0, int(num_reqs)))
+        ]
+
     def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
         return None
 
@@ -145,10 +154,10 @@ class Qwen3TTSCode2Wav(nn.Module):
         device: torch.device,
         codec_chunk_frames: int,
         codec_left_context_frames: int,
-        decode_cudagraph_capture_sizes: list[int] | None,
+        initial_codec_chunk_frames: int,
+        codec_chunk_ramp: list[int] | None,
         decode_cudagraph_batch_sizes: list[int] | None,
-        decode_cudagraph_extra_capture_shapes: list[tuple[int, int]] | None,
-        decode_compile_shapes: list[tuple[int, int]] | None,
+        decode_cudagraph_capture_sizes: list[int] | None,
     ) -> None:
         """Enable inner Code2Wav CUDA graph unless stage is enforce_eager."""
         if not hasattr(self.decoder, "enable_cudagraph") or device.type != "cuda":
@@ -160,7 +169,8 @@ class Qwen3TTSCode2Wav(nn.Module):
             return
 
         if (
-            codec_chunk_frames > 0
+            self._async_chunk
+            and codec_chunk_frames > 0
             and codec_left_context_frames > 0
             and self._decoder_sliding_window
             and codec_left_context_frames < self._decoder_sliding_window
@@ -177,24 +187,19 @@ class Qwen3TTSCode2Wav(nn.Module):
             )
 
         self.decoder.enable_cudagraph(
-            capture_sizes=decode_cudagraph_capture_sizes,
             capture_batch_sizes=decode_cudagraph_batch_sizes,
-            extra_capture_shapes=decode_cudagraph_extra_capture_shapes,
-            compile_shapes=decode_compile_shapes,
+            stateless_capture_sizes=decode_cudagraph_capture_sizes,
             device=device,
             codec_chunk_frames=codec_chunk_frames,
             codec_left_context_frames=codec_left_context_frames,
+            initial_codec_chunk_frames=initial_codec_chunk_frames,
+            codec_chunk_ramp=codec_chunk_ramp,
+            async_chunk=self._async_chunk,
             decode_chunk_size=self._decode_chunk_frames,
             decode_left_context=self._decode_left_context_frames,
         )
         logger.info("Code2Wav decoder CUDA Graph enabled")
 
-    def _get_decode_batch_bucket_frames(self, actual_frames: int) -> int:
-        for bucket_frames in self._decode_batch_bucket_frames:
-            if actual_frames <= bucket_frames:
-                return bucket_frames
-        return actual_frames
-
     def _record_decode_batch_stats(
         self,
         *,
@@ -212,48 +217,9 @@ class Qwen3TTSCode2Wav(nn.Module):
         self._batch_stats_actual_frames.update(actual_frames)
         self._batch_stats_bucket_groups[(group_size, bucket_frames)] += 1
 
-    @staticmethod
-    def _tensor_nbytes(tensor: torch.Tensor) -> int:
-        return int(tensor.numel() * tensor.element_size())
-
-    def _evict_ref_context_cache_if_needed(self) -> None:
-        evicted = 0
-        while len(self._ref_context_cache) > self._ref_context_cache_max_entries:
-            _, cached = self._ref_context_cache.popitem(last=False)
-            self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
-            evicted += 1
-        while self._ref_context_cache_bytes > self._ref_context_cache_max_bytes and len(self._ref_context_cache) > 1:
-            _, cached = self._ref_context_cache.popitem(last=False)
-            self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
-            evicted += 1
-        if evicted:
-            logger.debug(
-                "Evicted %d Qwen3-TTS ref context cache entries; entries=%d bytes=%d",
-                evicted,
-                len(self._ref_context_cache),
-                self._ref_context_cache_bytes,
-            )
-
-    def _cache_ref_context(self, request_id: str, tensor: torch.Tensor) -> None:
-        previous = self._ref_context_cache.pop(request_id, None)
-        if previous is not None:
-            self._ref_context_cache_bytes -= self._tensor_nbytes(previous)
-        cached = tensor.detach().contiguous()
-        self._ref_context_cache[request_id] = cached
-        self._ref_context_cache.move_to_end(request_id)
-        self._ref_context_cache_bytes += self._tensor_nbytes(cached)
-        self._evict_ref_context_cache_if_needed()
-
-    def _get_ref_context(self, request_id: str) -> torch.Tensor | None:
-        cached = self._ref_context_cache.get(request_id)
-        if cached is not None:
-            self._ref_context_cache.move_to_end(request_id)
-        return cached
-
-    def _pop_ref_context(self, request_id: str) -> None:
-        cached = self._ref_context_cache.pop(request_id, None)
-        if cached is not None:
-            self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
+    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+        for req_id in finished_req_ids:
+            self._decoder_state_cache.pop(req_id, None)
 
     def log_decode_batch_stats(self) -> None:
         if not self._batch_stats_enabled or self._batch_stats_requests == 0:
@@ -299,7 +265,6 @@ class Qwen3TTSCode2Wav(nn.Module):
         self._batch_stats_forwards += 1
         decoder = self.decoder
         q = int(self._num_quantizers)
-        upsample = int(self._total_upsample)
         sr_val = int(self._output_sample_rate)
         sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
         empty = torch.zeros((0,), dtype=torch.float32)
@@ -314,11 +279,13 @@ class Qwen3TTSCode2Wav(nn.Module):
         ids = input_ids.reshape(-1).to(dtype=torch.long)
         request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
 
-        parsed: list[tuple[int, int]] = []
-        valid_codes_qf: list[torch.Tensor] = []
+        internal_request_ids = kwargs.get("request_ids")
+        valid_codes_qf: list[tuple[str | None, torch.Tensor]] = []
         valid_indices: list[int] = []
         left_context_size = [0] * len(request_ids_list)
         ref_context_size = [0] * len(request_ids_list)
+        segment_finished_flags = [False] * len(request_ids_list)
+        request_state_ids: list[str | None] = [None] * len(request_ids_list)
         ref_context_request_ids: list[str | None] = [None] * len(request_ids_list)
         ref_context_included = [False] * len(request_ids_list)
         finished_flags = [False] * len(request_ids_list)
@@ -346,28 +313,42 @@ class Qwen3TTSCode2Wav(nn.Module):
 
         if runtime_infos:
             for i, info in enumerate(runtime_infos):
-                if i >= len(left_context_size):
+                if i >= len(ref_context_size):
                     break
                 if not isinstance(info, dict):
                     continue
                 meta = info.get("meta", {})
+                if "is_segment_finished" in meta:
+                    segment_finished_flags[i] = _meta_bool(meta["is_segment_finished"])
                 if "left_context_size" in meta:
                     left_context_size[i] = _meta_int(meta["left_context_size"])
                 if "ref_context_size" in meta:
                     ref_context_size[i] = _meta_int(meta["ref_context_size"])
                 if "ref_context_request_id" in meta:
                     ref_context_request_ids[i] = _meta_str(meta["ref_context_request_id"])
+                if "request_id" in meta:
+                    request_state_ids[i] = _meta_str(meta["request_id"])
                 if "ref_context_included" in meta:
                     ref_context_included[i] = _meta_bool(meta["ref_context_included"])
                 if "finished" in meta:
                     finished_flags[i] = _meta_bool(meta["finished"])
+
+        # Normal runner calls provide scheduler-side IDs, which are also used
+        # by scheduler_output.finished_req_ids. Direct forward calls and CUDA
+        # Graph dummy runs have no runner IDs, so retain the payload ID as a
+        # fallback for those paths.
+        cache_request_ids = request_state_ids.copy()
+        if internal_request_ids is not None:
+            for i, request_id in enumerate(internal_request_ids):
+                if i >= len(cache_request_ids):
+                    break
+                cache_request_ids[i] = str(request_id)
+
         for i, req_ids in enumerate(request_ids_list):
             runtime_info = runtime_infos[i] if i < len(runtime_infos) else None
             req_ids = _codec_ids_from_payload_or_input(req_ids, runtime_info)
             if req_ids.numel() < 1:
-                parsed.append((0, 0))
                 continue
-            ctx_frames = left_context_size[i]
             ref_ctx_frames = ref_context_size[i]
             flat = req_ids
             n = flat.numel()
@@ -382,36 +363,54 @@ class Qwen3TTSCode2Wav(nn.Module):
                             n,
                             q,
                         )
-                parsed.append((0, 0))
                 continue
             frames = n // q
             # [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F])
             codes_qf = flat.reshape(q, frames)
             ref_req_id = ref_context_request_ids[i]
-            if ref_req_id is not None and ref_ctx_frames > 0:
-                if ref_context_included[i]:
-                    if frames < ref_ctx_frames:
-                        raise ValueError(
-                            "Qwen3-TTS ref context metadata says ref prefix is included, "
-                            f"but frames={frames} < ref_context_size={ref_ctx_frames}"
+            state_req_id = cache_request_ids[i] if self._async_chunk else None
+            is_new_state = state_req_id is not None and state_req_id not in self._decoder_state_cache
+            if is_new_state and ref_req_id is not None and ref_ctx_frames > 0:
+                if not ref_context_included[i] or frames < ref_ctx_frames:
+                    raise ValueError("Qwen3-TTS async_chunk first ICL chunk must include its declared reference prefix")
+            valid_codes_qf.append((state_req_id, codes_qf))
+            if state_req_id is not None:
+                state = self._decoder_state_cache.get(state_req_id)
+                if state is None:
+                    if len(self._decoder_state_cache) >= self._decoder_state_cache_warn_entries:
+                        logger.warning_once(
+                            "Qwen3-TTS decoder state cache exceeded the expected active-request envelope: "
+                            "entries=%d threshold=%d. Keeping all active states; check request cleanup paths.",
+                            len(self._decoder_state_cache),
+                            self._decoder_state_cache_warn_entries,
                         )
-                    self._cache_ref_context(ref_req_id, codes_qf[:, :ref_ctx_frames])
-                else:
-                    cached_ref = self._get_ref_context(ref_req_id)
-                    if cached_ref is None:
+                    state = {}
+                    self._decoder_state_cache[state_req_id] = state
+                state.setdefault("prefix_frames", 0)
+                if state_req_id.startswith(_DUMMY_REQUEST_ID):
+                    state["_is_dummy_run"] = True
+                if ref_ctx_frames > 0:
+                    cached_prefix_frames = state.setdefault("prefix_frames", ref_ctx_frames)
+                    if cached_prefix_frames == 0:
+                        state["prefix_frames"] = ref_ctx_frames
+                        cached_prefix_frames = ref_ctx_frames
+                    if cached_prefix_frames != ref_ctx_frames:
                         raise ValueError(
-                            "Missing Qwen3-TTS ref context cache for "
-                            f"request {ref_req_id!r}; first chunk must include ref_code"
+                            "Qwen3-TTS ref context size changed within request "
+                            f"{ref_req_id!r}: cached={cached_prefix_frames}, current={ref_ctx_frames}"
                         )
-                    cached_ref = cached_ref.to(device=codes_qf.device, dtype=codes_qf.dtype)
-                    codes_qf = torch.cat((cached_ref, codes_qf), dim=1)
-                    frames = int(codes_qf.shape[1])
-            parsed.append((ctx_frames, frames))
-            valid_codes_qf.append(codes_qf)
             valid_indices.append(i)
 
         num_req = len(request_ids_list)
         if not valid_codes_qf:
+            for req_id, finished, segment_finished in zip(
+                cache_request_ids,
+                finished_flags,
+                segment_finished_flags,
+                strict=False,
+            ):
+                if req_id is not None and (finished or segment_finished):
+                    self._decoder_state_cache.pop(req_id, None)
             return OmniOutput(
                 text_hidden_states=None,
                 multimodal_outputs={
@@ -423,7 +422,7 @@ class Qwen3TTSCode2Wav(nn.Module):
         if not self._logged_codec_stats:
             self._logged_codec_stats = True
             try:
-                c = valid_codes_qf[0]
+                _, c = valid_codes_qf[0]
                 logger.info(
                     "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d",
                     c.shape[1],
@@ -436,92 +435,54 @@ class Qwen3TTSCode2Wav(nn.Module):
             except Exception:
                 pass
 
-        wav_tensors: list[torch.Tensor | None] = [None] * len(valid_codes_qf)
-
-        def _decode_group_chunks(group_chunks: list[list[tuple[int, torch.Tensor]]]) -> None:
-            for group_chunk in group_chunks:
-                actual_frames = [int(codes_qf.shape[1]) for _, codes_qf in group_chunk]
-                target_frames = max(actual_frames)
-                is_equal_length_batch = all(frames == target_frames for frames in actual_frames)
-                use_variable_length_batch = (
-                    len(group_chunk) > 1
-                    and not is_equal_length_batch
-                    and target_frames >= self._decode_variable_chunk_batch_min_frames
-                    and hasattr(decoder, "batched_chunked_decode")
+        request_states: list[dict[str, Any]] | None = None
+        if self._async_chunk:
+            missing_state = [index for index, (state_req_id, _) in enumerate(valid_codes_qf) if state_req_id is None]
+            if missing_state:
+                raise ValueError(
+                    "Qwen3-TTS async_chunk Code2Wav inputs require a request_id; "
+                    f"missing state for batch indices {missing_state}"
                 )
-                if len(group_chunk) == 1:
-                    codes_bqf = group_chunk[0][1].unsqueeze(0)
-                elif is_equal_length_batch:
-                    codes_bqf = torch.stack([codes_qf for _, codes_qf in group_chunk], dim=0)
-                else:
-                    first = group_chunk[0][1]
-                    codes_bqf = first.new_zeros((len(group_chunk), q, target_frames))
-                    for row, (_, codes_qf) in enumerate(group_chunk):
-                        codes_bqf[row, :, : codes_qf.shape[1]] = codes_qf
-                self._record_decode_batch_stats(
-                    group_size=len(group_chunk),
-                    bucket_frames=target_frames,
-                    actual_frames=actual_frames,
-                )
-                try:
-                    if use_variable_length_batch:
-                        wav_batch = decoder.batched_chunked_decode(
-                            codes_bqf,
-                            actual_frames,
-                            chunk_size=self._decode_chunk_frames,
-                            left_context_size=self._decode_left_context_frames,
-                            max_batch_size=self._decode_batch_max_size,
-                        )  # [B, 1, wav_len]
-                    else:
-                        wav_batch = decoder.chunked_decode(
-                            codes_bqf,
-                            chunk_size=self._decode_chunk_frames,
-                            left_context_size=self._decode_left_context_frames,
-                        )  # [B, 1, wav_len]
-                except TypeError:
-                    # Unit-test fakes and older decoder shims may not accept the
-                    # explicit chunk kwargs; production Qwen3-TTS decoders do.
-                    wav_batch = decoder.chunked_decode(codes_bqf)  # [B, 1, wav_len]
+            request_states = [
+                self._decoder_state_cache[state_req_id]
+                for state_req_id, _ in valid_codes_qf
+                if state_req_id is not None
+            ]
 
-                if wav_batch.dim() == 3 and wav_batch.shape[1] == 1:
-                    wav_rows = wav_batch[:, 0, :]
-                elif wav_batch.dim() == 2:
-                    wav_rows = wav_batch
-                else:
-                    raise ValueError(
-                        "Code2Wav decoder returned unexpected shape "
-                        f"{tuple(wav_batch.shape)} for batch size {len(group_chunk)}"
-                    )
-                if wav_rows.shape[0] != len(group_chunk):
-                    raise ValueError(
-                        f"Code2Wav decoder returned batch size {wav_rows.shape[0]} "
-                        f"for input batch size {len(group_chunk)}"
-                    )
-                for row, (j, _) in enumerate(group_chunk):
-                    wav_tensors[j] = wav_rows[row]
+        request_lengths = [int(codes_qf.shape[-1]) for _, codes_qf in valid_codes_qf]
+        max_request_length = max(request_lengths)
+        request_codes = valid_codes_qf[0][1].new_zeros((len(valid_codes_qf), q, max_request_length))
+        for row, (_, codes_qf) in enumerate(valid_codes_qf):
+            request_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf)
 
-        # Group by configured frame buckets instead of only exact lengths.
-        # For ordinary async streaming windows this is the real batching
-        # opportunity; decoder-internal variable chunk batching is gated to
-        # longer inputs where repeated full chunks can amortize its overhead.
-        grouped_codes: dict[int, list[tuple[int, torch.Tensor]]] = {}
-        for j, codes_qf in enumerate(valid_codes_qf):
-            frames = int(codes_qf.shape[1])
-            grouped_codes.setdefault(self._get_decode_batch_bucket_frames(frames), []).append((j, codes_qf))
-
-        for _bucket_frames, group in grouped_codes.items():
-            if self._decode_batch_max_size > 0 and len(group) > self._decode_batch_max_size:
-                # Keep each decoder call inside the configured CUDA graph batch
-                # envelope. Sorting by length lowers right-padding within each
-                # split while outputs are restored by original request index.
-                group = sorted(group, key=lambda item: int(item[1].shape[1]))
-                group_chunks = [
-                    group[start : start + self._decode_batch_max_size]
-                    for start in range(0, len(group), self._decode_batch_max_size)
-                ]
-            else:
-                group_chunks = [group]
-            _decode_group_chunks(group_chunks)
+        self._record_decode_batch_stats(
+            group_size=len(valid_codes_qf),
+            bucket_frames=max_request_length,
+            actual_frames=request_lengths,
+        )
+        request_wavs = decoder.batched_chunked_decode(
+            request_codes,
+            request_lengths,
+            caches=request_states,
+            chunk_size=self._decode_chunk_frames,
+            left_context_size=self._decode_left_context_frames,
+            max_batch_size=self._decode_batch_max_size,
+        )
+        if len(request_wavs) != len(valid_codes_qf):
+            raise ValueError(
+                f"Qwen3-TTS batched decoder returned {len(request_wavs)} outputs for {len(valid_codes_qf)} requests"
+            )
+        wav_tensors: list[torch.Tensor] = []
+        for row in range(len(valid_codes_qf)):
+            wav = request_wavs[row]
+            if wav.dim() == 2 and wav.shape[0] == 1:
+                wav = wav[0]
+            elif wav.dim() != 1:
+                raise ValueError(f"Qwen3-TTS batched decoder returned unexpected row shape {tuple(wav.shape)}")
+            if request_states is None:
+                start = left_context_size[valid_indices[row]] * self._total_upsample
+                wav = wav[start:]
+            wav_tensors.append(wav)
 
         if self._batch_stats_log_every > 0 and self._batch_stats_forwards % self._batch_stats_log_every == 0:
             self.log_decode_batch_stats()
@@ -530,27 +491,22 @@ class Qwen3TTSCode2Wav(nn.Module):
         srs = [sr_tensor] * num_req
 
         for j, idx in enumerate(valid_indices):
-            ctx_frames, actual_frames = parsed[idx]
             wav = wav_tensors[j]
             assert wav is not None
-            # Slice on exact codec-frame boundaries instead of proportionally.
-            start = max(0, ctx_frames * upsample)
-            end = max(start, actual_frames * upsample)
-            if start >= wav.shape[0]:
-                logger.warning(
-                    "Context trim start %d >= decoded length %d; returning empty audio.",
-                    start,
-                    wav.shape[0],
-                )
+            if wav.numel() == 0:
                 continue
-            wav = wav[start : min(end, wav.shape[0])]
             if wav.shape[0] > 0:
                 # Decoder already runs in fp32, so the .to(float32) is a redundant dispatch.
                 audios[idx] = (wav if wav.dtype == torch.float32 else wav.to(torch.float32)).reshape(-1)
 
-        for req_id, finished in zip(ref_context_request_ids, finished_flags, strict=False):
-            if req_id is not None and finished:
-                self._pop_ref_context(req_id)
+        for req_id, finished, segment_finished in zip(
+            cache_request_ids,
+            finished_flags,
+            segment_finished_flags,
+            strict=False,
+        ):
+            if req_id is not None and (finished or segment_finished):
+                self._decoder_state_cache.pop(req_id, None)
 
         return OmniOutput(
             text_hidden_states=None,
@@ -707,6 +663,8 @@ class Qwen3TTSCode2Wav(nn.Module):
         if isinstance(extra_cfg, dict):
             codec_chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0)
             codec_left_context_frames = int(extra_cfg.get("codec_left_context_frames") or 0)
+            initial_codec_chunk_frames = int(extra_cfg.get("initial_codec_chunk_frames") or 1)
+            codec_chunk_ramp = parse_chunk_ramp(extra_cfg, steady=codec_chunk_frames) if self._async_chunk else None
             decode_chunk_frames = _get_int_config("decode_chunk_frames", self._decode_chunk_frames)
             decode_left_context_frames = _get_int_config(
                 "decode_left_context_frames",
@@ -720,33 +678,22 @@ class Qwen3TTSCode2Wav(nn.Module):
                 )
             self._decode_chunk_frames = decode_chunk_frames
             self._decode_left_context_frames = decode_left_context_frames
-            decode_cudagraph_capture_sizes = _get_int_list_config("decode_cudagraph_capture_sizes")
             decode_cudagraph_batch_sizes = _get_int_list_config("decode_cudagraph_batch_sizes")
-            decode_cudagraph_extra_capture_shapes = _get_int_pair_list_config("decode_cudagraph_extra_capture_shapes")
-            decode_compile_shapes = _get_int_pair_list_config("decode_compile_shapes")
-            decode_batch_bucket_frames = _get_int_list_config("decode_batch_bucket_frames")
-            if decode_batch_bucket_frames is not None:
-                self._decode_batch_bucket_frames = decode_batch_bucket_frames
+            decode_cudagraph_capture_sizes = (
+                None if self._async_chunk else _get_int_list_config("decode_cudagraph_capture_sizes")
+            )
             decode_batch_max_size = _get_int_config("decode_batch_max_size", self._decode_batch_max_size)
             if decode_batch_max_size < 0:
                 raise ValueError(f"Invalid Qwen3-TTS Code2Wav config decode_batch_max_size={decode_batch_max_size}")
             self._decode_batch_max_size = decode_batch_max_size
-            decode_variable_chunk_batch_min_frames = _get_int_config(
-                "decode_variable_chunk_batch_min_frames",
-                self._decode_variable_chunk_batch_min_frames,
-            )
-            if decode_variable_chunk_batch_min_frames < 0:
-                raise ValueError(
-                    "Invalid Qwen3-TTS Code2Wav config "
-                    f"decode_variable_chunk_batch_min_frames={decode_variable_chunk_batch_min_frames}"
-                )
-            self._decode_variable_chunk_batch_min_frames = decode_variable_chunk_batch_min_frames
             decode_enable_tf32 = _get_bool_config("decode_enable_tf32", False)
         else:
-            decode_cudagraph_capture_sizes = None
+            codec_chunk_frames = 0
+            codec_left_context_frames = 0
+            initial_codec_chunk_frames = 1
+            codec_chunk_ramp = None
             decode_cudagraph_batch_sizes = None
-            decode_cudagraph_extra_capture_shapes = None
-            decode_compile_shapes = None
+            decode_cudagraph_capture_sizes = None
             decode_enable_tf32 = False
 
         if decode_enable_tf32 and device.type == "cuda":
@@ -764,16 +711,20 @@ class Qwen3TTSCode2Wav(nn.Module):
                 torch.get_float32_matmul_precision(),
             )
 
+        self.decoder._initial_codec_chunk_frames = initial_codec_chunk_frames
+        self.decoder._incremental_chunk_frames = codec_chunk_frames or 25
+        self.decoder._incremental_chunk_ramp = list(codec_chunk_ramp or ())
+
         if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda":
             try:
                 self._maybe_enable_decoder_cudagraph(
                     device=device,
                     codec_chunk_frames=codec_chunk_frames,
                     codec_left_context_frames=codec_left_context_frames,
-                    decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes,
+                    initial_codec_chunk_frames=initial_codec_chunk_frames,
+                    codec_chunk_ramp=codec_chunk_ramp,
                     decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes,
-                    decode_cudagraph_extra_capture_shapes=decode_cudagraph_extra_capture_shapes,
-                    decode_compile_shapes=decode_compile_shapes,
+                    decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes,
                 )
             except Exception:
                 logger.warning(

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """PyTorch Qwen3TTSTokenizerV2 model."""
 
+import copy
 import inspect
 import math
 from collections.abc import Callable
@@ -25,9 +26,9 @@ from torch import nn
 from torch.nn import functional as F
 from transformers import MimiConfig, MimiModel
 from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache
+from transformers.cache_utils import Cache, DynamicCache
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_sliding_window_causal_mask
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import BaseModelOutputWithPast
@@ -520,64 +521,15 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         self.norm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(config=config)
         self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
         self.window_size = config.sliding_window
 
         self.input_proj = nn.Linear(config.latent_dim, config.hidden_size)
         self.output_proj = nn.Linear(config.hidden_size, config.latent_dim)
 
-        # Code2Wav is cache-free and every decoder layer uses sliding attention.
-        # Cache an exact, compact mask per warmup shape instead of retaining a
-        # max_position_embeddings square buffer. CUDA Graph warmup populates
-        # these entries before capture, and dtype/device are part of the key.
-        self._sliding_attention_mask_cache: dict[tuple[int, torch.dtype, torch.device, str], torch.Tensor | None] = {}
-
         # Initialize weights and apply final processing
         self.post_init()
 
-    def _apply(self, fn, recurse=True):
-        # Cached masks are keyed by dtype/device and should not retain storage
-        # from the model's previous placement.
-        self._sliding_attention_mask_cache.clear()
-        return super()._apply(fn, recurse=recurse)
-
-    def _get_sliding_attention_mask(self, inputs_embeds: torch.Tensor) -> torch.Tensor | None:
-        sequence_length = int(inputs_embeds.shape[1])
-        max_position_embeddings = int(self.config.max_position_embeddings)
-        if sequence_length > max_position_embeddings:
-            raise ValueError(
-                f"Input length {sequence_length} exceeds max_position_embeddings {max_position_embeddings}"
-            )
-
-        attention_implementation = str(self.config._attn_implementation)
-        cache_key = (
-            sequence_length,
-            inputs_embeds.dtype,
-            inputs_embeds.device,
-            attention_implementation,
-        )
-        if cache_key not in self._sliding_attention_mask_cache:
-            # A batch-one mask broadcasts across requests and avoids caching a
-            # duplicate tensor for every batch size.
-            mask_inputs = inputs_embeds[:1]
-            mask_kwargs = {
-                "config": self.config,
-                "attention_mask": None,
-                "past_key_values": None,
-                "position_ids": None,
-            }
-            sig = inspect.signature(create_sliding_window_causal_mask)
-            if "input_embeds" in sig.parameters:
-                mask_kwargs["input_embeds"] = mask_inputs
-                mask_kwargs["cache_position"] = torch.arange(sequence_length, device=inputs_embeds.device)
-            else:
-                mask_kwargs["inputs_embeds"] = mask_inputs
-            sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs)
-            if sliding_attention_mask is not None:
-                sliding_attention_mask = sliding_attention_mask.contiguous()
-            self._sliding_attention_mask_cache[cache_key] = sliding_attention_mask
-
-        return self._sliding_attention_mask_cache[cache_key]
-
     # Note: @check_model_inputs decorator removed for vLLM compatibility
     # The decorator causes "unexpected keyword argument 'inputs_embeds'" error
     @auto_docstring
@@ -600,48 +552,44 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
             raise ValueError("input_ids is not expected")
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-        if use_cache:
-            raise ValueError("Code2Wav sliding-mask caching does not support use_cache=True")
-        if past_key_values is not None:
-            raise ValueError("Code2Wav sliding-mask caching does not support past_key_values")
-        if attention_mask is not None and not isinstance(attention_mask, dict):
-            raise ValueError("Code2Wav cached masks require attention_mask to be None or a prepared mask mapping")
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
         inputs_embeds = self.input_proj(inputs_embeds)
 
-        sequence_length = inputs_embeds.shape[1]
-        expected_cache_position = torch.arange(sequence_length, device=inputs_embeds.device)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
 
         if cache_position is None:
-            cache_position = expected_cache_position
-        elif attention_mask is None and (
-            cache_position.shape != expected_cache_position.shape
-            or cache_position.device != expected_cache_position.device
-            or not torch.equal(cache_position, expected_cache_position)
-        ):
-            raise ValueError("Cached Code2Wav masks require contiguous zero-based cache_position")
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
 
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
-        elif attention_mask is None:
-            expected_position_ids = expected_cache_position.unsqueeze(0)
-            if (
-                position_ids.ndim != 2
-                or position_ids.shape[-1] != sequence_length
-                or position_ids.device != expected_position_ids.device
-                or not torch.equal(position_ids, expected_position_ids.expand_as(position_ids))
-            ):
-                raise ValueError("Cached Code2Wav masks require contiguous zero-based position_ids")
 
         hidden_states = inputs_embeds
 
-        if isinstance(attention_mask, dict):
-            sliding_attention_mask = attention_mask["sliding_attention"]
-        else:
-            sliding_attention_mask = self._get_sliding_attention_mask(inputs_embeds)
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            sig = inspect.signature(create_causal_mask)
+            if "input_embeds" in sig.parameters:
+                mask_kwargs["input_embeds"] = inputs_embeds
+                mask_kwargs["cache_position"] = cache_position
+            else:
+                mask_kwargs["inputs_embeds"] = inputs_embeds
+            causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
+            if self.has_sliding_layers:
+                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
 
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -649,7 +597,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
             hidden_states = decoder_layer(
                 hidden_states,
-                attention_mask=sliding_attention_mask,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
                 position_ids=position_ids,
                 past_key_values=past_key_values,
                 use_cache=use_cache,
@@ -662,7 +610,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         hidden_states = self.output_proj(hidden_states)
         return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
-            past_key_values=None,
+            past_key_values=past_key_values if use_cache else None,
         )
 
 
@@ -852,10 +800,15 @@ class SplitResidualVectorQuantizer(nn.Module):
         return quantized
 
 
+_CONV_CONTEXT_FRAME = 2
+_DOWNSTREAM_CONTEXT_FRAME = 12
+
+
 class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
     def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig):
         super().__init__(config)
         self.total_upsample = np.prod(config.upsample_rates + config.upsampling_ratios)
+        self.upsample_factor = np.prod(config.upsampling_ratios)
         self.pre_transformer = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config)
 
         self.quantizer = SplitResidualVectorQuantizer(
@@ -913,17 +866,19 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
 
     def enable_cudagraph(
         self,
-        capture_sizes: list[int] | None = None,
         capture_batch_sizes: list[int] | None = None,
-        extra_capture_shapes: list[tuple[int, int]] | None = None,
-        compile_shapes: list[tuple[int, int]] | None = None,
+        stateless_capture_sizes: list[int] | None = None,
         device: torch.device | None = None,
         codec_chunk_frames: int = 0,
         codec_left_context_frames: int = 0,
+        initial_codec_chunk_frames: int = 1,
+        codec_chunk_ramp: list[int] | None = None,
+        async_chunk: bool = True,
         decode_chunk_size: int = 300,
         decode_left_context: int = 25,
     ):
-        from ..cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+        # from ..cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+        from ..segmented_graph_wrapper import CUDAGraphDecoderWrapper
 
         if device is None:
             device = next(self.parameters()).device
@@ -933,34 +888,35 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
 
         self._cudagraph_wrapper = CUDAGraphDecoderWrapper(
             decoder=self,
-            capture_sizes=capture_sizes,
             capture_batch_sizes=capture_batch_sizes,
-            extra_capture_shapes=extra_capture_shapes,
-            compile_shapes=compile_shapes,
+            stateless_capture_sizes=stateless_capture_sizes,
             num_quantizers=self.config.num_quantizers,
             enabled=True,
-        )
-        self._cudagraph_wrapper.warmup(
-            device,
-            dtype=torch.long,
-            codec_chunk_frames=codec_chunk_frames,
-            codec_left_context_frames=codec_left_context_frames,
+            async_chunk=async_chunk,
+            initial_chunk_frames=initial_codec_chunk_frames,
+            codec_chunk_frames=codec_chunk_frames or 25,
+            codec_chunk_ramp=codec_chunk_ramp,
             decode_chunk_size=decode_chunk_size,
             decode_left_context=decode_left_context,
         )
+        self._incremental_chunk_frames = codec_chunk_frames or 25
+        self._incremental_chunk_ramp = list(codec_chunk_ramp or ())
+        self._cudagraph_wrapper.warmup(
+            device,
+            dtype=torch.long,
+        )
         self._cudagraph_enabled = True
         logger.info(
-            "CUDA Graph enabled for decoder: batch_sizes=%s seq_lens=%s extra_shapes=%s compile_shapes=%s",
+            "CUDA Graph enabled for decoder: batch_sizes=%s seq_lens=%s",
             self._cudagraph_wrapper.capture_batch_sizes,
-            self._cudagraph_wrapper.capture_sizes,
-            self._cudagraph_wrapper.extra_capture_shapes,
-            self._cudagraph_wrapper.compile_shapes,
+            (
+                self._cudagraph_wrapper.icl_capture_sizes
+                if async_chunk
+                else self._cudagraph_wrapper.stateless_capture_sizes
+            ),
         )
 
-    def forward(self, codes):
-        if codes.shape[1] != self.config.num_quantizers:
-            raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
-
+    def _forward_exact(self, codes):
         hidden = self.quantizer.decode(codes)
         hidden = self.pre_conv(hidden).transpose(1, 2)
 
@@ -969,15 +925,267 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         for blocks in self.upsample:
             for block in blocks:
                 hidden = block(hidden)
+
         wav = hidden
         for block in self.decoder:
             wav = block(wav)
         return wav.clamp(min=-1, max=1)
 
-    def chunked_decode(self, codes, chunk_size=300, left_context_size=25):
+    def _decode_icl_first_chunk(
+        self,
+        codes,
+        caches,
+        prefix_frames,
+        prefix_cache=None,
+        prefix_hidden_mask=None,
+        prefix_attention_mask=None,
+        suffix_attention_mask=None,
+    ):
+        # The CUDA graph path may left-pad a shorter request prefix to the
+        # captured sliding-window length.  Keep the decoder's physical layout
+        # separate from the request metadata so eager fallbacks slice the
+        # cached tensors at the correct boundary.
+        caches["decoder_prefix_frames"] = prefix_frames
+        hidden = self.quantizer.decode(codes)
+        if prefix_hidden_mask is not None:
+            hidden = hidden * prefix_hidden_mask
+        caches.update({"ref_hidden": hidden[:, :, :prefix_frames]})
+        caches.update({"suffix_quantized": hidden[:, :, prefix_frames:]})
+
+        hidden = self.pre_conv(hidden).transpose(1, 2)
+        caches.update({"ref_conv": hidden[:, :prefix_frames, :]})
+        caches.update({"suffix_conv": hidden[:, prefix_frames:, :]})
+
+        if hidden.shape[1] < prefix_frames:
+            raise ValueError(
+                "Qwen3-TTS ICL prefix cache requires at least prefix_frames inputs, "
+                f"got prefix_frames={prefix_frames}, total_frames={hidden.shape[1]}"
+            )
+        prefix_result = self.pre_transformer(
+            inputs_embeds=hidden[:, :prefix_frames, :],
+            attention_mask=prefix_attention_mask,
+            past_key_values=prefix_cache,
+            use_cache=True,
+        )
+
+        prefix_cache = prefix_result.past_key_values
+        working_cache = copy.deepcopy(prefix_cache)
+
+        suffix_result = self.pre_transformer(
+            inputs_embeds=hidden[:, prefix_frames:, :],
+            attention_mask=suffix_attention_mask,
+            past_key_values=working_cache,
+            use_cache=True,
+        )
+
+        prefix_hidden = prefix_result.last_hidden_state
+        if prefix_attention_mask is not None:
+            prefix_hidden = prefix_hidden * prefix_attention_mask.unsqueeze(-1)
+        hidden = torch.cat(
+            [
+                prefix_hidden,
+                suffix_result.last_hidden_state,
+            ],
+            dim=1,
+        )
+        caches.update({"past_key_values": prefix_cache})
+        caches.update({"prefix_hidden": prefix_hidden})
+
+        hidden = hidden.permute(0, 2, 1)
+
+        for blocks in self.upsample:
+            for block in blocks:
+                hidden = block(hidden)
+        wav = hidden
+
+        for block in self.decoder:
+            wav = block(wav)
+        wav = wav.clamp(min=-1, max=1)
+        return wav
+
+    def _decode_xvec_first_chunk(self, codes: torch.Tensor, caches: dict) -> torch.Tensor:
+        """Initialize prefixless incremental state from the first codec chunk."""
+        hidden = self.quantizer.decode(codes)
+        caches["decoder_prefix_frames"] = 0
+        caches["ref_hidden"] = hidden[:, :, :0]
+        caches["suffix_quantized"] = hidden
+
+        hidden = self.pre_conv(hidden).transpose(1, 2)
+        caches["ref_conv"] = hidden[:, :0, :]
+        caches["suffix_conv"] = hidden
+
+        empty_prefix_cache = caches.get("past_key_values")
+        if empty_prefix_cache is None:
+            empty_prefix_cache = DynamicCache(config=self.config)
+            head_dim = getattr(
+                self.config,
+                "head_dim",
+                self.config.hidden_size // self.config.num_attention_heads,
+            )
+            empty_prefix_cache.early_initialization(
+                batch_size=int(codes.shape[0]),
+                num_heads=self.config.num_key_value_heads,
+                head_dim=head_dim,
+                dtype=hidden.dtype,
+                device=hidden.device,
+            )
+        working_cache = copy.deepcopy(empty_prefix_cache)
+        suffix_hidden = self.pre_transformer(
+            inputs_embeds=hidden,
+            past_key_values=working_cache,
+            use_cache=True,
+        ).last_hidden_state
+        caches["past_key_values"] = empty_prefix_cache
+        caches["prefix_hidden"] = suffix_hidden[:, :0, :]
+
+        hidden = suffix_hidden.permute(0, 2, 1)
+        for blocks in self.upsample:
+            for block in blocks:
+                hidden = block(hidden)
+        wav = hidden
+        for block in self.decoder:
+            wav = block(wav)
+        caches["suffix_frames"] = int(codes.shape[-1])
+        return wav.clamp(min=-1, max=1)
+
+    def _decode_suffix(
+        self,
+        new_codes: torch.Tensor,
+        old_quantized: torch.Tensor,
+        old_conv: torch.Tensor,
+        caches: dict,
+        prefix_frames: int,
+        new_frames: int,
+        rolling: bool,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        ref_hidden = caches["ref_hidden"]
+        conv_context = torch.cat([ref_hidden, old_quantized], dim=-1)[:, :, -_CONV_CONTEXT_FRAME:]
+        suffix_cache_length = int(self.config.sliding_window)
+
+        new_quantized = self.quantizer.decode(new_codes)
+        new_conv_input = torch.cat([conv_context, new_quantized], dim=-1)
+        new_conv = self.pre_conv(new_conv_input)
+        new_conv = new_conv[:, :, -new_codes.shape[-1] :].transpose(1, 2)
+
+        if rolling:
+            boundary_input = torch.cat(
+                [ref_hidden[:, :, -_CONV_CONTEXT_FRAME:], old_quantized[:, :, :_CONV_CONTEXT_FRAME]],
+                dim=-1,
+            )
+            boundary_conv = self.pre_conv(boundary_input)
+            boundary_conv = boundary_conv[:, :, -_CONV_CONTEXT_FRAME:].transpose(1, 2)
+            suffix_conv = torch.cat([boundary_conv, old_conv, new_conv], dim=1)
+        else:
+            suffix_conv = torch.cat([old_conv, new_conv], dim=1)
+
+        next_quantized = torch.cat([old_quantized, new_quantized], dim=-1)[:, :, -suffix_cache_length:]
+        next_conv = suffix_conv[:, -(suffix_cache_length - _CONV_CONTEXT_FRAME) :, :]
+
+        hidden = torch.cat([caches["ref_conv"], suffix_conv], dim=1)
+        hidden = hidden.transpose(1, 2).contiguous().transpose(1, 2)
+        working_cache = copy.deepcopy(caches["past_key_values"])
+        suffix_hidden = self.pre_transformer(
+            inputs_embeds=hidden[:, prefix_frames:, :],
+            attention_mask=attention_mask,
+            past_key_values=working_cache,
+            use_cache=True,
+        ).last_hidden_state
+        hidden = torch.cat([caches["prefix_hidden"], suffix_hidden], dim=1)
+
+        required_frames = min(hidden.shape[1], new_frames + _DOWNSTREAM_CONTEXT_FRAME)
+        hidden = hidden[:, -required_frames:, :].permute(0, 2, 1)
+        for blocks in self.upsample:
+            for block in blocks:
+                hidden = block(hidden)
+
+        wav = hidden
+        for block in self.decoder:
+            wav = block(wav)
+        wav = wav[..., -new_frames * self.total_upsample :].clamp(min=-1, max=1)
+        return wav, next_quantized, next_conv
+
+    def _is_suffix_cache_rolling(self, previous_suffix_frames: int, cached_frames: int) -> bool:
+        suffix_cache_length = int(self.config.sliding_window)
+        return previous_suffix_frames >= suffix_cache_length and cached_frames == suffix_cache_length
+
+    def decode_suffix(self, new_codes, caches, prefix_frames):
+        previous_suffix_frames = int(caches.get("suffix_frames", caches["suffix_quantized"].shape[-1]))
+        cached_frames = int(caches["suffix_quantized"].shape[-1])
+        suffix_cache_length = int(self.config.sliding_window)
+        rolling = self._is_suffix_cache_rolling(previous_suffix_frames, cached_frames)
+        retained_frames = suffix_cache_length if rolling else previous_suffix_frames
+        new_frames = int(new_codes.shape[-1])
+        max_new_frames = int(getattr(self, "_incremental_chunk_frames", 25))
+        if not 0 < new_frames <= max_new_frames:
+            raise ValueError(
+                f"Qwen3-TTS incremental decode expected 1..{max_new_frames} new frames, got new_frames={new_frames}"
+            )
+
+        attention_mask = torch.ones(
+            new_codes.shape[0],
+            prefix_frames + retained_frames + new_frames,
+            dtype=torch.bool,
+            device=new_codes.device,
+        )
+        attention_mask[:, : int(caches.get("prefix_pad_frames", 0))] = 0
+        output, next_quantized, next_conv = self._decode_suffix(
+            new_codes,
+            caches["suffix_quantized"],
+            caches["suffix_conv"],
+            caches,
+            prefix_frames,
+            new_frames,
+            rolling,
+            attention_mask=attention_mask,
+        )
+        caches["suffix_quantized"] = next_quantized
+        caches["suffix_conv"] = next_conv
+        caches["suffix_frames"] = retained_frames + new_frames
+        return output
+
+    def forward(self, codes, caches=None):
+        if codes.shape[1] != self.config.num_quantizers:
+            raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
+
+        if caches is None or caches.get("_is_dummy_run", False):
+            return self._forward_exact(codes)
+
+        prefix_frames = int(caches["prefix_frames"])
+        if prefix_frames < 0:
+            raise ValueError(
+                "Qwen3-TTS ICL prefix cache received invalid frame counts: "
+                f"prefix_frames={prefix_frames}, total_frames={codes.shape[-1]}"
+            )
+
+        if "suffix_quantized" not in caches:
+            if codes.shape[-1] < prefix_frames:
+                raise ValueError(
+                    "Qwen3-TTS ICL first chunk is shorter than its prefix: "
+                    f"prefix_frames={prefix_frames}, total_frames={codes.shape[-1]}"
+                )
+            if prefix_frames == 0:
+                return self._decode_xvec_first_chunk(codes, caches)
+            return self._decode_icl_first_chunk(codes, caches, prefix_frames)
+        else:
+            decoder_prefix_frames = int(caches.get("decoder_prefix_frames", prefix_frames))
+            return self.decode_suffix(codes, caches, decoder_prefix_frames)
+
+    def chunked_decode(
+        self,
+        codes,
+        caches=None,
+        chunk_size=300,
+        left_context_size=25,
+    ):
         # Use CUDA graph if enabled
         if self._cudagraph_enabled and self._cudagraph_wrapper is not None:
-            return self._cudagraph_wrapper.chunked_decode_with_cudagraph(codes, chunk_size, left_context_size)
+            return self._cudagraph_wrapper.chunked_decode_with_cudagraph(
+                codes,
+                caches=caches,
+                chunk_size=chunk_size,
+                left_context_size=left_context_size,
+            )
 
         # Original implementation (eager mode)
         wavs = []
@@ -986,31 +1194,385 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
             end_index = min(start_index + chunk_size, codes.shape[-1])
             context_size = left_context_size if start_index - left_context_size > 0 else start_index
             codes_chunk = codes[..., start_index - context_size : end_index]
-            wav_chunk = self(codes_chunk)
-            wavs.append(wav_chunk[..., context_size * self.total_upsample :])
+            is_first_stateful_chunk = caches is not None and "suffix_quantized" not in caches
+            wav_chunk = self(codes_chunk, caches)
+            if caches is not None:
+                if is_first_stateful_chunk:
+                    logical_prefix_frames = int(caches["prefix_frames"])
+                    suffix_frames = codes_chunk.shape[-1] - logical_prefix_frames
+                    wav_chunk = wav_chunk[..., -suffix_frames * self.total_upsample :]
+                wavs.append(wav_chunk)
+            else:
+                wavs.append(wav_chunk[..., context_size * self.total_upsample :])
             start_index = end_index
-        return torch.cat(wavs, dim=-1)
+        output = torch.cat(wavs, dim=-1)
+        return output
+
+    def batched_request_decode(
+        self,
+        codes_list: list[torch.Tensor],
+        request_caches: list[dict[str, Any]],
+        *,
+        prefix_length: int,
+        initial_chunk_frames: int,
+        codec_chunk_frames: int,
+        transitions_by_mode: dict[str, dict[int, int]],
+        decode_icl_prefix_batch: Callable[[list[torch.Tensor], list[dict[str, Any]]], list[torch.Tensor] | None] | None,
+        decode_xvec_prefix_batch: Callable[[list[torch.Tensor], list[dict[str, Any]]], list[torch.Tensor] | None]
+        | None,
+        decode_suffix_batch: Callable[
+            [str, int, list[torch.Tensor], list[dict[str, Any]], list[int]], list[torch.Tensor] | None
+        ]
+        | None,
+        decode_fallback: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor],
+        backend_max_batch_size: int = 0,
+    ) -> list[torch.Tensor]:
+        """Group stateful requests once, then execute them with a supplied backend.
+
+        The grouping policy is shared by CUDA Graph and eager execution.  A
+        backend may decline a grouped decode by returning ``None``; requests in
+        that group are then decoded individually through ``decode_fallback``.
+        """
+        if len(codes_list) != len(request_caches):
+            raise ValueError("codes_list and request_caches must have the same length")
+
+        if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
+            return [self._forward_exact(codes) for codes in codes_list]
+
+        outputs: list[torch.Tensor | None] = [None] * len(codes_list)
+        groups: dict[tuple[str, int], list[int]] = {}
+        for index, (codes, cache) in enumerate(zip(codes_list, request_caches, strict=True)):
+            if "suffix_quantized" not in cache:
+                prefix_frames = int(cache["prefix_frames"])
+                suffix_frames = int(codes.shape[-1]) - prefix_frames
+                if prefix_frames == 0:
+                    phase = "xvec_first"
+                elif 0 < prefix_frames <= prefix_length and suffix_frames == initial_chunk_frames:
+                    phase = "icl_prefix"
+                else:
+                    phase = "eager"
+                groups.setdefault((phase, 0), []).append(index)
+                continue
+
+            request_kv = cache["past_key_values"]
+            decoder_prefix = int(cache.get("decoder_prefix_frames", cache["prefix_frames"]))
+            mode = "xvec" if decoder_prefix == 0 else "icl"
+            expected_kv_length = 0 if mode == "xvec" else prefix_length
+            if decoder_prefix != expected_kv_length or request_kv.get_seq_length() != expected_kv_length:
+                groups.setdefault(("eager", 0), []).append(index)
+                continue
+
+            previous = int(cache.get("suffix_frames", cache["suffix_quantized"].shape[-1]))
+            rolling = self._is_suffix_cache_rolling(previous, int(cache["suffix_quantized"].shape[-1]))
+            target = (
+                prefix_length + codec_chunk_frames
+                if rolling
+                else next(
+                    (target for target, source in transitions_by_mode[mode].items() if source == previous),
+                    None,
+                )
+            )
+            new_frames = int(codes.shape[-1])
+            if target is None or not 0 < new_frames <= codec_chunk_frames:
+                groups.setdefault(("eager", 0), []).append(index)
+            else:
+                groups.setdefault((f"suffix:{mode}", target), []).append(index)
+
+        for (phase, target), indices in groups.items():
+            group_splits = (
+                [
+                    indices[start : start + backend_max_batch_size]
+                    for start in range(0, len(indices), backend_max_batch_size)
+                ]
+                if backend_max_batch_size > 0
+                else [indices]
+            )
+            for split_indices in group_splits:
+                group_codes = [codes_list[index] for index in split_indices]
+                group_caches = [request_caches[index] for index in split_indices]
+                group_outputs: list[torch.Tensor] | None = None
+                if phase == "icl_prefix" and decode_icl_prefix_batch is not None:
+                    group_outputs = decode_icl_prefix_batch(group_codes, group_caches)
+                elif phase == "xvec_first" and decode_xvec_prefix_batch is not None:
+                    group_outputs = decode_xvec_prefix_batch(group_codes, group_caches)
+                elif phase.startswith("suffix:") and decode_suffix_batch is not None:
+                    mode = phase.removeprefix("suffix:")
+                    new_frames = [int(codes.shape[-1]) for codes in group_codes]
+                    group_outputs = decode_suffix_batch(mode, target, group_codes, group_caches, new_frames)
+
+                if group_outputs is None:
+                    group_outputs = [
+                        decode_fallback(codes_list[index], request_caches[index]) for index in split_indices
+                    ]
+                for index, output in zip(split_indices, group_outputs, strict=True):
+                    outputs[index] = output
+
+        if any(output is None for output in outputs):
+            raise RuntimeError("stateful batched decode did not produce an output for every request")
+        return [output for output in outputs if output is not None]
+
+    @staticmethod
+    def _batch_dynamic_caches(request_caches: list[DynamicCache]) -> DynamicCache:
+        batched_cache = copy.deepcopy(request_caches[0])
+        for batched_layer, request_layers in zip(
+            batched_cache.layers,
+            zip(*(cache.layers for cache in request_caches), strict=True),
+            strict=True,
+        ):
+            if any(layer.keys is None or layer.values is None for layer in request_layers):
+                if not all(layer.keys is None and layer.values is None for layer in request_layers):
+                    raise ValueError("Cannot batch partially initialized request KV caches")
+                batched_layer.keys = None
+                batched_layer.values = None
+                continue
+            batched_layer.keys = torch.cat([layer.keys for layer in request_layers], dim=0)
+            batched_layer.values = torch.cat([layer.values for layer in request_layers], dim=0)
+        return batched_cache
+
+    @staticmethod
+    def _slice_dynamic_cache(cache: DynamicCache, row: int) -> DynamicCache:
+        request_cache = copy.deepcopy(cache)
+        for layer in request_cache.layers:
+            if layer.keys is not None:
+                layer.keys = layer.keys[row : row + 1].clone()
+            if layer.values is not None:
+                layer.values = layer.values[row : row + 1].clone()
+        return request_cache
+
+    @staticmethod
+    def _cache_tensors_are_batchable(request_caches: list[dict[str, Any]], keys: tuple[str, ...]) -> bool:
+        return all(
+            key in cache
+            and cache[key].shape[1:] == request_caches[0][key].shape[1:]
+            and cache[key].dtype == request_caches[0][key].dtype
+            and cache[key].device == request_caches[0][key].device
+            for cache in request_caches
+            for key in keys
+        )
+
+    def _decode_icl_prefix_eager_batch(
+        self,
+        codes_list: list[torch.Tensor],
+        request_caches: list[dict[str, Any]],
+        *,
+        prefix_length: int,
+        initial_chunk_frames: int,
+    ) -> list[torch.Tensor] | None:
+        batch_size = len(codes_list)
+        total_frames = prefix_length + initial_chunk_frames
+        batched_codes = codes_list[0].new_zeros((batch_size, codes_list[0].shape[1], total_frames))
+        hidden_mask = torch.ones((batch_size, 1, total_frames), dtype=self.dtype, device=batched_codes.device)
+        prefix_mask = torch.ones((batch_size, prefix_length), dtype=torch.bool, device=batched_codes.device)
+        suffix_mask = torch.ones((batch_size, total_frames), dtype=torch.bool, device=batched_codes.device)
+        prefix_pads: list[int] = []
+        for row, (codes, cache) in enumerate(zip(codes_list, request_caches, strict=True)):
+            prefix_frames = int(cache["prefix_frames"])
+            suffix_frames = int(codes.shape[-1]) - prefix_frames
+            if not 0 < prefix_frames <= prefix_length or suffix_frames != initial_chunk_frames:
+                return None
+            prefix_pad = prefix_length - prefix_frames
+            prefix_pads.append(prefix_pad)
+            batched_codes[row, :, prefix_pad:prefix_length].copy_(codes[0, :, :prefix_frames])
+            batched_codes[row, :, prefix_length:].copy_(codes[0, :, prefix_frames:])
+            hidden_mask[row, :, :prefix_pad] = 0
+            prefix_mask[row, :prefix_pad] = 0
+            suffix_mask[row, :prefix_pad] = 0
+
+        batched_cache: dict[str, Any] = {}
+        output = self._decode_icl_first_chunk(
+            batched_codes,
+            batched_cache,
+            prefix_length,
+            prefix_hidden_mask=hidden_mask,
+            prefix_attention_mask=prefix_mask,
+            suffix_attention_mask=suffix_mask,
+        )
+        outputs: list[torch.Tensor] = []
+        for row, (cache, prefix_pad) in enumerate(zip(request_caches, prefix_pads, strict=True)):
+            for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
+                cache[key] = batched_cache[key][row : row + 1].clone()
+            cache["past_key_values"] = self._slice_dynamic_cache(batched_cache["past_key_values"], row)
+            cache["decoder_prefix_frames"] = prefix_length
+            cache["prefix_pad_frames"] = prefix_pad
+            cache["suffix_frames"] = initial_chunk_frames
+            outputs.append(output[row : row + 1, :, -initial_chunk_frames * self.total_upsample :].clone())
+        return outputs
+
+    def _decode_xvec_prefix_eager_batch(
+        self,
+        codes_list: list[torch.Tensor],
+        request_caches: list[dict[str, Any]],
+        *,
+        initial_chunk_frames: int,
+    ) -> list[torch.Tensor] | None:
+        if any(int(codes.shape[-1]) != initial_chunk_frames for codes in codes_list):
+            return None
+        batched_codes = torch.cat(codes_list, dim=0)
+        batched_cache: dict[str, Any] = {}
+        output = self._decode_xvec_first_chunk(batched_codes, batched_cache)
+        outputs: list[torch.Tensor] = []
+        for row, cache in enumerate(request_caches):
+            for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
+                cache[key] = batched_cache[key][row : row + 1].clone()
+            cache["past_key_values"] = self._slice_dynamic_cache(batched_cache["past_key_values"], row)
+            cache["decoder_prefix_frames"] = 0
+            cache["suffix_frames"] = initial_chunk_frames
+            outputs.append(output[row : row + 1].clone())
+        return outputs
+
+    def _decode_suffix_eager_batch(
+        self,
+        mode: str,
+        target_frames: int,
+        codes_list: list[torch.Tensor],
+        request_caches: list[dict[str, Any]],
+        new_frames_list: list[int],
+        *,
+        prefix_length: int,
+        codec_chunk_frames: int,
+        transitions_by_mode: dict[str, dict[int, int]],
+    ) -> list[torch.Tensor] | None:
+        tensor_keys = ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv")
+        if not self._cache_tensors_are_batchable(request_caches, tensor_keys):
+            return None
+        previous_frames = transitions_by_mode[mode].get(target_frames)
+        if previous_frames is None:
+            return None
+        expected_new_frames = target_frames - previous_frames
+        if any(not 0 < new_frames <= expected_new_frames for new_frames in new_frames_list):
+            return None
+
+        batched_codes = codes_list[0].new_zeros((len(codes_list), codes_list[0].shape[1], expected_new_frames))
+        for row, (codes, new_frames) in enumerate(zip(codes_list, new_frames_list, strict=True)):
+            batched_codes[row, :, :new_frames].copy_(codes[0, :, -new_frames:])
+        batched_cache = {
+            key: torch.cat([cache[key] for cache in request_caches], dim=0)
+            for key in ("ref_hidden", "ref_conv", "prefix_hidden")
+        }
+        try:
+            batched_cache["past_key_values"] = self._batch_dynamic_caches(
+                [cache["past_key_values"] for cache in request_caches]
+            )
+        except ValueError:
+            return None
+        attention_mask = torch.ones(
+            len(codes_list),
+            (prefix_length if mode == "icl" else 0) + previous_frames + expected_new_frames,
+            dtype=torch.bool,
+            device=batched_codes.device,
+        )
+        if mode == "icl":
+            for row, cache in enumerate(request_caches):
+                attention_mask[row, : int(cache.get("prefix_pad_frames", 0))] = 0
+        old_quantized = torch.cat([cache["suffix_quantized"] for cache in request_caches], dim=0)
+        old_conv = torch.cat([cache["suffix_conv"] for cache in request_caches], dim=0)
+        rolling = self._is_suffix_cache_rolling(previous_frames, int(old_quantized.shape[-1]))
+        output, next_quantized, next_conv = self._decode_suffix(
+            batched_codes,
+            old_quantized,
+            old_conv,
+            batched_cache,
+            prefix_length if mode == "icl" else 0,
+            expected_new_frames,
+            rolling,
+            attention_mask=attention_mask,
+        )
+
+        outputs: list[torch.Tensor] = []
+        for row, (cache, new_frames) in enumerate(zip(request_caches, new_frames_list, strict=True)):
+            outputs.append(output[row : row + 1, :, : new_frames * self.total_upsample].clone())
+            if new_frames == expected_new_frames:
+                cache["suffix_quantized"] = next_quantized[row : row + 1].clone()
+                cache["suffix_conv"] = next_conv[row : row + 1].clone()
+                cache["suffix_frames"] = target_frames
+        return outputs
 
     def batched_chunked_decode(
         self,
         codes,
         lengths,
+        caches=None,
         chunk_size=300,
         left_context_size=25,
         max_batch_size=0,
     ):
-        if self._cudagraph_enabled and self._cudagraph_wrapper is not None:
+        if (
+            self._cudagraph_enabled
+            and self._cudagraph_wrapper is not None
+            and hasattr(self._cudagraph_wrapper, "batched_chunked_decode_with_cudagraph")
+        ):
             return self._cudagraph_wrapper.batched_chunked_decode_with_cudagraph(
                 codes,
                 lengths,
+                caches=caches,
                 chunk_size=chunk_size,
                 left_context_size=left_context_size,
                 max_batch_size=max_batch_size,
             )
 
+        if caches is not None:
+            prefix_length = int(getattr(self.config, "sliding_window", 0) or 0)
+            initial_codec_chunk_frames = int(getattr(self, "_initial_codec_chunk_frames", 1))
+            codec_chunk_frames = int(getattr(self, "_incremental_chunk_frames", 25))
+            chunk_ramp = list(getattr(self, "_incremental_chunk_ramp", ()) or ())
+            initial_chunk_frames = int(chunk_ramp[0]) if chunk_ramp else initial_codec_chunk_frames
+            transitions: dict[int, int] = {}
+            previous = initial_chunk_frames
+            for new_frames in chunk_ramp[1:]:
+                if previous >= prefix_length:
+                    break
+                target = previous + int(new_frames)
+                transitions[target] = previous
+                previous = target
+            while previous < prefix_length:
+                target = previous + codec_chunk_frames
+                transitions[target] = previous
+                previous = target
+            transitions[prefix_length + codec_chunk_frames] = prefix_length
+            codes_list = [codes[row : row + 1, :, :length] for row, length in enumerate(lengths)]
+            outputs = self.batched_request_decode(
+                codes_list,
+                caches,
+                prefix_length=prefix_length,
+                initial_chunk_frames=initial_chunk_frames,
+                codec_chunk_frames=codec_chunk_frames,
+                transitions_by_mode={"icl": transitions, "xvec": transitions},
+                decode_icl_prefix_batch=lambda group_codes, group_caches: self._decode_icl_prefix_eager_batch(
+                    group_codes,
+                    group_caches,
+                    prefix_length=prefix_length,
+                    initial_chunk_frames=initial_chunk_frames,
+                ),
+                decode_xvec_prefix_batch=lambda group_codes, group_caches: self._decode_xvec_prefix_eager_batch(
+                    group_codes,
+                    group_caches,
+                    initial_chunk_frames=initial_chunk_frames,
+                ),
+                decode_suffix_batch=lambda mode, target, group_codes, group_caches, new_frames: (
+                    self._decode_suffix_eager_batch(
+                        mode,
+                        target,
+                        group_codes,
+                        group_caches,
+                        new_frames,
+                        prefix_length=prefix_length,
+                        codec_chunk_frames=codec_chunk_frames,
+                        transitions_by_mode={"icl": transitions, "xvec": transitions},
+                    )
+                ),
+                decode_fallback=lambda request_codes, cache: self.chunked_decode(
+                    request_codes,
+                    caches=cache,
+                    chunk_size=chunk_size,
+                    left_context_size=left_context_size,
+                ),
+                backend_max_batch_size=max_batch_size,
+            )
+            return [output[0] for output in outputs]
+
         from ..cuda_graph_decoder_wrapper import _batched_chunked_decode
 
-        return _batched_chunked_decode(
+        padded = _batched_chunked_decode(
             codes,
             lengths,
             decode_fn=self,
@@ -1019,6 +1581,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
             left_context_size=left_context_size,
             max_batch_size=max_batch_size,
         )
+        return [padded[row, :, : length * self.total_upsample] for row, length in enumerate(lengths)]
 
 
 class Qwen3TTSTokenizerV2Encoder(MimiModel):

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -164,21 +164,31 @@ def talker2code2wav_async_chunk(
         if finished:
             return OmniPayloadStruct(
                 codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
-                meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+                meta=MetaStruct(
+                    request_id=request_id,
+                    left_context_size=0,
+                    finished=torch.tensor(True, dtype=torch.bool),
+                ),
             )
         return None
 
     if ramp is not None:
         chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0)
+        first_chunk = chunk_index <= 0
         emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size, finished)
         if not emit:
             return None
         if context_length == 0:
             return OmniPayloadStruct(
                 codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
-                meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+                meta=MetaStruct(
+                    request_id=request_id,
+                    left_context_size=0,
+                    finished=torch.tensor(True, dtype=torch.bool),
+                ),
             )
     else:
+        first_chunk = int(transfer_manager.put_req_chunk.get(request_id, 0)) <= 0
         use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < chunk_size
 
         if use_first_chunk and length <= initial_chunk_size:
@@ -193,15 +203,18 @@ def talker2code2wav_async_chunk(
             chunk_length = adjusted % chunk_size
             context_length = chunk_length if chunk_length != 0 else chunk_size
 
-    end_index = min(length, left_context_size_config + context_length)
-    left_context_size = max(0, end_index - context_length)
-    window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+    if finished and first_chunk:
+        context_length = length
 
-    # Prepend the bounded ref_code tail to the first emitted chunk so Code2Wav
-    # can cache the same reference context that mainline would otherwise send on
-    # every chunk. Follow-up chunks send only a metadata handle; Code2Wav restores
-    # this context locally before decode.
+    # Code2Wav keeps the quantizer/conv/Transformer context per request. Ship
+    # only the newly completed codec frames after the first chunk.
+    window_frames = transfer_manager.code_prompt_token_ids[request_id][-context_length:]
+    left_context_size = 0
+
+    # ICL reference codes are part of the first decoder initialization only.
+    # Follow-up chunks rely entirely on the request-local decoder cache.
     ref_code = request_payload.get(request_id)
+    emitted_chunks = int(transfer_manager.put_req_chunk.get(request_id, 0))
     ref_context_size = 0
     ref_context_request_id: str | None = None
     ref_context_included = False
@@ -231,13 +244,12 @@ def talker2code2wav_async_chunk(
             ref_context = ref_context[-ref_code_context_frames:]
         ref_context_size = int(ref_context.shape[0]) if ref_context.ndim > 1 else 0
         if ref_context_size > 0:
-            ref_context_request_id = request_id
-            emitted_chunks = int(transfer_manager.put_req_chunk.get(request_id, 0))
             if emitted_chunks <= 0:
+                ref_context_request_id = request_id
                 ref_frames = ref_context.tolist()
                 window_frames = ref_frames + window_frames
                 ref_context_included = True
-            left_context_size += ref_context_size
+                left_context_size = ref_context_size
 
     num_quantizers = len(window_frames[0])
     num_frames = len(window_frames)
@@ -247,6 +259,7 @@ def talker2code2wav_async_chunk(
     )
 
     meta = MetaStruct(
+        request_id=request_id,
         left_context_size=left_context_size,
         finished=torch.tensor(finished, dtype=torch.bool),
     )