[Performance][MiniCPM-o] Whole-Euler CUDA graphs for Code2Wav with a shared attention arena (#8007)

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -15,7 +15,16 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
 
-from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper
+from .cuda_graph_wrapper import (
+    CFMGraphWrapper,
+    HiFTGraphWrapper,
+    WholeEulerCFMGraphWrapper,
+    _att_keep_ranges,
+    _copy_frame_segments,
+    _euler_step,
+    _euler_timeline,
+    _zero_padded_cnn_cache,
+)
 
 logger = init_logger(__name__)
 
@@ -25,6 +34,17 @@ _SILENCE_TOKEN = 4218
 # than this many codec tokens overflows the PE slice (matrix_ac vs matrix_bd).
 _DEFAULT_RELPOS_MAX_POS = 5000
 _MAX_ENCODE_TOKEN_CAP = 1024
+# Streaming cache trim (stepaudio2 ``Token2wav.stream``): past ``prompt_len +
+# 100`` frames, keep the first ``prompt_len`` frames and the last 100.
+_CACHE_TRIM_SUFFIX = 100
+
+
+def _trim_streaming_cache(cache: torch.Tensor, prompt_len: int) -> torch.Tensor:
+    """The streaming trim of a cache whose frame axis is ``-2``."""
+    ranges = _att_keep_ranges(int(cache.shape[-2]), (prompt_len, _CACHE_TRIM_SUFFIX))
+    if len(ranges) == 1:
+        return cache
+    return torch.cat([cache[..., start : start + length, :] for start, length in ranges], dim=-2)
 
 
 def relpos_encode_token_budget(
@@ -47,6 +67,26 @@ def relpos_encode_token_budget(
     return max(lookahead + 1, min(int(cap), room))
 
 
+def _supports_ragged_kernel(estimator: nn.Module) -> bool:
+    """Whether ``_blocks_forward_chunk_ragged`` can run this estimator's DiT blocks."""
+    blocks = getattr(estimator, "blocks", None)
+    if not blocks or not hasattr(estimator, "in_proj") or not hasattr(estimator, "final_layer"):
+        return False
+    block = blocks[0]
+    attn = getattr(block, "attn", None)
+    conv_blocks = getattr(getattr(block, "conv", None), "block", None)
+    return (
+        hasattr(block, "adaLN_modulation")
+        and all(
+            hasattr(attn, name)
+            for name in ("forward_chunk", "to_q", "to_k", "to_v", "to_heads", "q_norm", "k_norm", "proj")
+        )
+        and conv_blocks is not None
+        and len(conv_blocks) >= 8
+        and all(isinstance(conv_blocks[index], nn.Conv1d) for index in (1, 6))
+    )
+
+
 def _cfm_pad_frames(
     *,
     mel_frames: int,
@@ -98,38 +138,6 @@ def _zero_padded_frames(tensor: torch.Tensor, valid_frames: int | None) -> None:
     tensor[..., valid_frames:] = 0.0
 
 
-def _zero_padded_cnn_cache(
-    cnn_cache: torch.Tensor,
-    estimator: nn.Module,
-    pad_frames: int,
-) -> None:
-    """Clear the cache positions that come from padded frames, in place.
-
-    Each block's CNN cache holds the tail of its convolution output
-    (``new_cnn_cache = x[..., -causal_padding[0]:]``, inside ``stepaudio2``),
-    so when the padding sits at the chunk tail those positions are
-    padding-derived and would otherwise become the next chunk's left context.
-    Padding is at most ``pad_frames`` wide, so only the trailing positions that
-    can come from it are cleared; the valid part of the window is kept.
-    """
-    blocks = estimator.blocks
-    if pad_frames > 0 and blocks and len(blocks) == cnn_cache.shape[0]:
-        width = int(cnn_cache.shape[-1])
-        if width > 0 and all(int(block.conv.block[1].causal_padding[0]) == width for block in blocks):
-            # _estimator_buffers packs equal-width blocks into one tensor.
-            # Clear their shared tail with one write instead of one per block.
-            cnn_cache[..., max(0, width - pad_frames) :] = 0.0
-            return
-
-    for index, block in enumerate(estimator.blocks):
-        width = int(block.conv.block[1].causal_padding[0])
-        if width <= 0:
-            continue
-        zero_from = max(0, width - pad_frames)
-        if zero_from < width:
-            cnn_cache[index][..., zero_from:] = 0.0
-
-
 def plan_token2wav_encode_slices(
     num_frames: int,
     *,
@@ -297,6 +305,7 @@ class BatchedToken2Wav(nn.Module):
             token2wav.speech_window.detach().clone(),
             persistent=False,
         )
+        self._timeline_cache: dict[tuple[torch.device, torch.dtype], tuple[torch.Tensor, list[float]]] = {}
         self.hift_graph_wrapper: HiFTGraphWrapper | None = None
         graph_config = dict(hift_graph_config or {})
         if bool(graph_config.get("enabled", False)):
@@ -312,26 +321,79 @@ class BatchedToken2Wav(nn.Module):
                         "MiniCPM-o HiFT CUDA Graph requires source_cache_len to be divisible by mel_cache_len"
                     )
                 capture_batch_sizes = graph_config.get("capture_batch_sizes", [1])
-                logger.info("Enabling HiFT CUDA Graph with batch sizes %s", capture_batch_sizes)
+                max_serial_batch = int(graph_config.get("max_serial_batch", 4))
+                logger.info(
+                    "Enabling HiFT CUDA Graph with batch sizes %s (max_serial_batch=%d)",
+                    capture_batch_sizes,
+                    max_serial_batch,
+                )
                 self.hift_graph_wrapper = HiFTGraphWrapper(
                     token2wav=token2wav,
                     connector_config=dict(connector_config),
                     capture_batch_sizes=capture_batch_sizes,
+                    max_serial_batch=max_serial_batch,
                 )
                 with torch.inference_mode(), _autocast_disabled(hift_parameter.device):
                     self.hift_graph_wrapper.capture()
                 logger.info("HiFT CUDA Graph captured successfully")
         self._cfm_graph_wrapper: CFMGraphWrapper | None = None
+        self._whole_euler_graph_wrapper: WholeEulerCFMGraphWrapper | None = None
         cfm_graph_cfg = dict(cfm_graph_config or {})
         if bool(cfm_graph_cfg.get("enabled", False)):
             flow_parameter = next(self.flow.parameters(), None)
             if flow_parameter is not None and flow_parameter.device.type == "cuda":
                 estimator = self.flow.decoder.estimator
+                max_graphs = int(cfm_graph_cfg.get("max_graphs", 32))
+                max_serial_batch = int(cfm_graph_cfg.get("max_serial_batch", 4))
+                max_graph_batch_cfg = cfm_graph_cfg.get("max_graph_batch")
+                max_graph_batch = int(max_graph_batch_cfg) if max_graph_batch_cfg is not None else None
+                micro_batch_size_cfg = cfm_graph_cfg.get("micro_batch_size")
+                micro_batch_size = int(micro_batch_size_cfg) if micro_batch_size_cfg is not None else 4
+                if bool(cfm_graph_cfg.get("enable_whole_euler", True)) and self._trt_stepper is None:
+                    capture_frames = cfm_graph_cfg.get("capture_frames")
+                    if capture_frames is None and connector_config is not None:
+                        # The steady chunk's mel width, derived like HiFTGraphWrapper's
+                        # capture bucket: steady chunks run unpadded and shorter
+                        # chunks share their graph under a mask.
+                        capture_frames = int(connector_config["codec_chunk_frames"]) * int(
+                            getattr(self.flow, "token_mel_ratio", 2)
+                        )
+                    query_bucket_frames = int(capture_frames if capture_frames is not None else 64)
+                    self._whole_euler_graph_wrapper = WholeEulerCFMGraphWrapper(
+                        estimator=estimator,
+                        n_timesteps=self.n_timesteps,
+                        inference_cfg_rate=getattr(self.flow.decoder, "inference_cfg_rate", 0.7),
+                        att_cache_dtype=self._estimator_att_cache_dtype,
+                        max_graphs=max_graphs,
+                        max_serial_batch=max_serial_batch,
+                        max_graph_batch=max_graph_batch,
+                        micro_batch_size=micro_batch_size,
+                        query_bucket_frames=query_bucket_frames,
+                        pad_max_rows=cfm_graph_cfg.get("pad_max_rows"),
+                        # Looked up at capture time, so a replaced ragged kernel is the one captured.
+                        ragged_body=(
+                            (lambda *args: self._blocks_forward_chunk_ragged(*args))
+                            if _supports_ragged_kernel(estimator)
+                            else None
+                        ),
+                    )
+                    logger.info(
+                        "Whole-Euler CFM CUDA Graph enabled "
+                        "(max_graphs=%d, max_serial_batch=%d, max_graph_batch=%s, "
+                        "micro_batch_size=%d, query_bucket_frames=%d)",
+                        max_graphs,
+                        max_serial_batch,
+                        str(max_graph_batch),
+                        micro_batch_size,
+                        query_bucket_frames,
+                    )
+                elif self._trt_stepper is not None and bool(cfm_graph_cfg.get("enable_whole_euler", True)):
+                    logger.info("Whole-Euler CFM CUDA Graph disabled because TensorRT stepper is configured")
                 self._cfm_graph_wrapper = CFMGraphWrapper(
                     graph_fn=estimator.blocks_forward_chunk,
-                    max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)),
+                    max_graphs=max_graphs,
                 )
-                logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32)))
+                logger.info("CFM CUDA Graph enabled (max_graphs=%d)", max_graphs)
             else:
                 logger.info(
                     "CFM CUDA Graph is disabled on device type %s",
@@ -341,7 +403,9 @@ class BatchedToken2Wav(nn.Module):
         # chunk up to a multiple of this many frames so the graph cache key
         # space stays small (0 disables bucketing, e.g. when graphs are off).
         self._cfm_graph_bucket_frames = (
-            int(cfm_graph_cfg.get("bucket_frames", 0)) if self._cfm_graph_wrapper is not None else 0
+            int(cfm_graph_cfg.get("bucket_frames", 0))
+            if (self._cfm_graph_wrapper is not None or self._whole_euler_graph_wrapper is not None)
+            else 0
         )
         if self._cfm_graph_bucket_frames > 1:
             logger.info(
@@ -350,9 +414,7 @@ class BatchedToken2Wav(nn.Module):
             )
         self._prompt_features: dict[tuple[str, str], PromptFeatures] = {}
         self._setup_cache_size = setup_cache_size
-        self._setup_cache: OrderedDict[tuple[tuple[str, str], int, int], tuple[BatchedToken2WavState, ...]] = (
-            OrderedDict()
-        )
+        self._setup_cache: OrderedDict[tuple[tuple[str, str], int], BatchedToken2WavState] = OrderedDict()
 
     def _hift_inference(
         self,
@@ -507,6 +569,7 @@ class BatchedToken2Wav(nn.Module):
         attn_mask: torch.Tensor | None = None,
         valid_lengths: list[int] | None = None,
         valid_frames: int | None = None,
+        time_embedding: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if self._trt_stepper is not None and valid_lengths is None:
             out, new_cnn, new_att = self._trt_stepper.step(
@@ -519,7 +582,8 @@ class BatchedToken2Wav(nn.Module):
                 att_cache=att_cache,
             )
             return out.to(mu.dtype), new_cnn, new_att
-        time_embedding = estimator.t_embedder(time).unsqueeze(1)
+        if time_embedding is None:
+            time_embedding = estimator.t_embedder(time).unsqueeze(1)
         width = int(x.shape[-1])
         speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
         if valid_frames is not None and valid_frames < width:
@@ -583,14 +647,95 @@ class BatchedToken2Wav(nn.Module):
         valid_lengths: torch.Tensor,
         width: int,
     ) -> torch.Tensor:
+        """The ``width`` frames of ``history`` (b, width + t, c) ending each row's valid frames, as (b, c, width)."""
         indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :]
-        return history.gather(
-            2,
-            indices[:, None, :].expand(-1, int(history.shape[1]), -1),
-        )
+        return history.gather(1, indices[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2)
 
+    @staticmethod
+    def _causal_conv_frames_major(conv: nn.Conv1d, history: torch.Tensor) -> torch.Tensor:
+        """``conv`` over ``history`` (b, width + t, c_in) as (b, t, c_out), without a layout change.
+
+        The DiT keeps activations as (b, t, c), which already is NHWC for a
+        height-1 image. Upstream transposes them to (b, c, t) for ``Conv1d``,
+        and cuDNN converts that to NHWC and back around every convolution. As
+        a channels-last conv2d it reads the activations where they are. The
+        channels-last weight is converted once and reused. Other devices run
+        upstream's ``Conv1d`` layout (Ascend has no channels-last format).
+        """
+        if not history.is_cuda:
+            output = F.conv1d(
+                history.transpose(1, 2),
+                conv.weight,
+                conv.bias,
+                conv.stride,
+                conv.padding,
+                conv.dilation,
+                conv.groups,
+            )
+            return output.transpose(1, 2)
+        weight = conv.weight
+        version = (weight.data_ptr(), weight._version, weight.dtype)
+        cached = conv.__dict__.get("_channels_last_weight")
+        if cached is None or cached[0] != version:
+            cached = (version, weight.detach().unsqueeze(2).contiguous(memory_format=torch.channels_last))
+            conv.__dict__["_channels_last_weight"] = cached
+        output = F.conv2d(
+            history.transpose(1, 2).unsqueeze(2),
+            cached[1],
+            conv.bias,
+            (1, conv.stride[0]),
+            (0, conv.padding[0]),
+            (1, conv.dilation[0]),
+            conv.groups,
+        )
+        return output.squeeze(2).transpose(1, 2)
+
+    @staticmethod
+    def _attend_in_place(
+        attn: nn.Module,
+        x: torch.Tensor,
+        att_cache: torch.Tensor | None,
+        attn_mask: torch.Tensor | None,
+        kv: torch.Tensor,
+    ) -> torch.Tensor:
+        """``attn.forward_chunk`` that keeps the new keys and values in ``kv``.
+
+        ``kv`` (b, heads, >= t + cache frames, 2 * head_dim) is the block's new
+        cache, ``[current | cache]`` like upstream's. Upstream concatenates the
+        whole cache behind the current keys and values, then both into the new
+        cache, which the caller copied into ``kv``: four passes over Stage 2's
+        largest tensor. Here only the current chunk is written and attention reads ``kv`` in
+        place. A cache that already sits behind the current chunk in ``kv``
+        (Whole-Euler) is not moved; any other is copied there once. Other
+        devices keep upstream's attention and copy its cache into ``kv``.
+        """
+        if kv.dtype != x.dtype or not x.is_cuda:
+            # A reduced-precision cache attends in the activation dtype.
+            x_att, new_att = attn.forward_chunk(x, att_cache, attn_mask)
+            kv[:, :, : int(new_att.shape[2])].copy_(new_att)
+            return x_att
+        batch, frames, _ = x.shape
+        q = attn.q_norm(attn.to_heads(attn.to_q(x)))
+        k = attn.k_norm(attn.to_heads(attn.to_k(x)))
+        v = attn.to_heads(attn.to_v(x))
+        cached = 0 if att_cache is None else int(att_cache.shape[2])
+        kv = kv[:, :, : frames + cached]
+        behind = kv[:, :, frames:]
+        if cached and (att_cache.data_ptr() != behind.data_ptr() or att_cache.stride() != behind.stride()):
+            behind.copy_(att_cache)
+        head_dim = int(k.shape[3])
+        kv[..., :frames, :head_dim].copy_(k)
+        kv[..., :frames, head_dim:].copy_(v)
+        out = F.scaled_dot_product_attention(
+            q,
+            kv[..., :head_dim],
+            kv[..., head_dim:],
+            attn_mask=None if attn_mask is None else attn_mask.unsqueeze(1),
+        )
+        return attn.proj(out.transpose(1, 2).reshape(batch, frames, -1))
+
+    @staticmethod
     def _blocks_forward_chunk_ragged(
-        self,
         estimator: nn.Module,
         estimator_input: torch.Tensor,
         time_embedding: torch.Tensor,
@@ -599,14 +744,24 @@ class BatchedToken2Wav(nn.Module):
         att_cache: Any,
         cnn_cache_buffer: torch.Tensor,
         att_cache_buffer: torch.Tensor,
-        valid_lengths: list[int],
+        valid_lengths: list[int] | torch.Tensor,
     ) -> torch.Tensor:
-        """Run one padded DiT batch while capturing exact per-row CNN state."""
-        lengths = torch.tensor(
-            (*valid_lengths, *valid_lengths),
-            device=estimator_input.device,
-            dtype=torch.long,
-        )
+        """Run one padded DiT batch while capturing exact per-row CNN state.
+
+        ``valid_lengths`` is one length per request, or the ``(2B,)`` CFG
+        lengths tensor a CUDA graph captures as a static input. Upstream's
+        ``DiTBlock.forward_chunk``, with attention over the output cache in
+        place (``_attend_in_place``) and the causal convolutions on the
+        frames-major activations (``_causal_conv_frames_major``).
+        """
+        if isinstance(valid_lengths, torch.Tensor):
+            lengths = valid_lengths
+        else:
+            lengths = torch.tensor(
+                (*valid_lengths, *valid_lengths),
+                device=estimator_input.device,
+                dtype=torch.long,
+            )
         x = estimator.in_proj(estimator_input.transpose(1, 2))
         for block_index, block in enumerate(estimator.blocks):
             (
@@ -622,49 +777,55 @@ class BatchedToken2Wav(nn.Module):
             ) = block.adaLN_modulation(time_embedding).chunk(9, dim=-1)
 
             normalized = block.norm1(x) * (1 + scale_msa) + shift_msa
-            x_att, new_att = block.attn.forward_chunk(
+            x_att = BatchedToken2Wav._attend_in_place(
+                block.attn,
                 normalized,
                 att_cache[block_index],
                 attn_mask,
+                att_cache_buffer[block_index],
             )
             x = x + gate_msa * x_att
 
+            # CausalConvBlock: [T, conv1, T, LayerNorm, Mish, T, conv2, T]
+            # with the transposes (T) dropped.
+            conv = block.conv
             conv_input = block.norm3(x) * (1 + scale_conv) + shift_conv
             old_cnn = cnn_cache[block_index]
             if old_cnn is None:
-                width = int(block.conv.block[1].causal_padding[0])
                 old_cnn = conv_input.new_zeros(
                     (
                         int(conv_input.shape[0]),
-                        int(block.conv.in_channels + block.conv.out_channels),
-                        width,
+                        int(conv.in_channels + conv.out_channels),
+                        int(conv.block[1].causal_padding[0]),
                     )
                 )
-            old_cnn1, old_cnn2 = old_cnn.split((block.conv.in_channels, block.conv.out_channels), dim=1)
-
-            conv1_input = block.conv.block[0](conv_input)
-            conv1_output, _ = block.conv.block[1].forward_chunk(conv1_input, old_cnn1)
-            new_cnn1 = self._gather_causal_cache(
-                torch.cat((old_cnn1, conv1_input), dim=2),
-                lengths,
-                int(old_cnn1.shape[2]),
-            )
-
-            conv2_input = block.conv.block[2:6](conv1_output)
-            conv2_output, _ = block.conv.block[6].forward_chunk(conv2_input, old_cnn2)
-            new_cnn2 = self._gather_causal_cache(
-                torch.cat((old_cnn2, conv2_input), dim=2),
-                lengths,
-                int(old_cnn2.shape[2]),
-            )
-            x = x + gate_conv * block.conv.block[7](conv2_output)
+            width = int(old_cnn.shape[2])
+            history = torch.cat((old_cnn[:, : conv.in_channels].transpose(1, 2), conv_input), dim=1)
+            new_cnn1 = BatchedToken2Wav._gather_causal_cache(history, lengths, width)
+            hidden = conv.block[4](conv.block[3](BatchedToken2Wav._causal_conv_frames_major(conv.block[1], history)))
+            history = torch.cat((old_cnn[:, conv.in_channels :].transpose(1, 2), hidden), dim=1)
+            new_cnn2 = BatchedToken2Wav._gather_causal_cache(history, lengths, width)
+            x = x + gate_conv * BatchedToken2Wav._causal_conv_frames_major(conv.block[6], history)
             x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp)
 
-            cnn_cache_buffer[block_index].copy_(torch.cat((new_cnn1, new_cnn2), dim=1))
-            att_cache_buffer[block_index][:, :, : int(new_att.shape[2]), :].copy_(new_att)
+            cnn_cache_buffer[block_index][:, : conv.in_channels].copy_(new_cnn1)
+            cnn_cache_buffer[block_index][:, conv.in_channels :].copy_(new_cnn2)
 
         return estimator.final_layer(x, time_embedding).transpose(1, 2)
 
+    def _get_timeline(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, list[float]]:
+        """The Euler timeline and its step sizes, read to the host once.
+
+        ``_euler_step`` takes ``dt`` as a Python float, so reading it
+        from the device timeline on every decode would sync once per step.
+        """
+        key = (device, dtype)
+        cached = self._timeline_cache.get(key)
+        if cached is None:
+            cached = _euler_timeline(self.n_timesteps, device, dtype)
+            self._timeline_cache[key] = cached
+        return cached
+
     def _decode_cfm(
         self,
         mu: torch.Tensor,
@@ -672,30 +833,44 @@ class BatchedToken2Wav(nn.Module):
         cond: torch.Tensor,
         *,
         cnn_cache: torch.Tensor | None,
-        att_cache: torch.Tensor | None,
+        att_cache: torch.Tensor | list[torch.Tensor] | None,
         valid_lengths: list[int] | None = None,
+        att_keep: tuple[int, int] | None = None,
     ) -> tuple[
         torch.Tensor,
         torch.Tensor,
         torch.Tensor | list[torch.Tensor],
     ]:
+        """Run the CFM Euler solve for one chunk.
+
+        ``att_cache`` may be the stacked estimator cache or one tensor per
+        request (the Whole-Euler path). A Whole-Euler solve returns the new
+        cache in the same form, already trimmed by ``att_keep``; the eager
+        path returns it stacked and untrimmed (ragged: per request, untrimmed).
+        """
         decoder = self.flow.decoder
         estimator = decoder.estimator
         batch_size = int(mu.shape[0])
-        offset = int(att_cache.shape[4]) if att_cache is not None else 0
+        att_rows = att_cache if isinstance(att_cache, list) else None
+        if att_rows is not None:
+            offset = int(att_rows[0].shape[4])
+        else:
+            offset = int(att_cache.shape[4]) if att_cache is not None else 0
         mel_frames = int(mu.shape[2])
+        graphs_active = valid_lengths is None and (
+            self._whole_euler_active()
+            or (
+                self._cfm_graph_wrapper is not None
+                and getattr(self._cfm_graph_wrapper, "enabled", True)
+                and self._trt_stepper is None
+            )
+        )
         pad_frames = _cfm_pad_frames(
             mel_frames=mel_frames,
             offset=offset,
             noise_capacity=int(decoder.rand_noise.shape[2]),
             bucket_frames=self._cfm_graph_bucket_frames,
-            disabled=(
-                valid_lengths is not None
-                or self._cfm_graph_wrapper is None
-                # `_disable` keeps the wrapper object alive, so check the flag
-                # too: padding under a disabled wrapper is pure overhead.
-                or not self._cfm_graph_wrapper.enabled
-            ),
+            disabled=(valid_lengths is not None or not graphs_active),
         )
         if pad_frames:
             # Replicate rather than zero: a repeated last frame is a closer
@@ -718,15 +893,7 @@ class BatchedToken2Wav(nn.Module):
             # is what removes them; zeroing alone would not, since a zero row
             # still occupies part of the softmax denominator.
             x[:, :, mel_frames:] = 0.0
-        timeline = torch.linspace(
-            0,
-            1,
-            self.n_timesteps + 1,
-            device=mu.device,
-            dtype=mu.dtype,
-        )
-        timeline = 1 - torch.cos(timeline * 0.5 * torch.pi)
-        time = timeline[0].expand(batch_size)
+        timeline, dt_steps = self._get_timeline(mu.device, mu.dtype)
         mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0)
         speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0)
         cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0)
@@ -761,19 +928,42 @@ class BatchedToken2Wav(nn.Module):
                 device=mu.device,
             )
             attn_mask[:, :, mel_frames : mel_frames + pad_frames] = False
+
+        if self._whole_euler_active() and (valid_lengths is None or self._whole_euler_ragged_active()):
+            whole_euler_result = self._whole_euler_graph_wrapper.replay(
+                x=x,
+                mu_cfg=mu_cfg,
+                speakers_cfg=speakers_cfg,
+                cond_cfg=cond_cfg,
+                cnn_cache=cnn_cache,
+                att_cache=att_cache,
+                attn_mask=attn_mask,
+                mel_frames=mel_frames,
+                pad_frames=pad_frames,
+                valid_lengths=valid_lengths,
+                att_keep=att_keep,
+            )
+            if whole_euler_result is not None:
+                return whole_euler_result
+        if att_rows is not None:
+            att_cache = self._stack_estimator_att(att_rows)
+
         next_cnn: list[torch.Tensor] = []
         next_att_cache: torch.Tensor | None = None
         ragged_att_cache: list[torch.Tensor] | None = None
-        dt = timeline[1] - timeline[0]
+        time_embeddings = [
+            estimator.t_embedder(timeline[s].expand(2 * batch_size)).unsqueeze(1) for s in range(self.n_timesteps)
+        ]
         with _token2wav_sdpa_context(mu.device):
             for step in range(self.n_timesteps):
                 old_cnn = cnn_cache[step] if cnn_cache is not None else None
                 old_att = att_cache[step] if att_cache is not None else None
+                step_time = timeline[step].expand(2 * batch_size)
                 estimate, step_cnn, step_att = self._estimator_step(
                     estimator,
                     x=torch.cat((x, x), dim=0),
                     mu=mu_cfg,
-                    time=torch.cat((time, time), dim=0),
+                    time=step_time,
                     speakers=speakers_cfg,
                     cond=cond_cfg,
                     cnn_cache=old_cnn,
@@ -781,16 +971,13 @@ class BatchedToken2Wav(nn.Module):
                     attn_mask=attn_mask,
                     valid_lengths=valid_lengths,
                     valid_frames=mel_frames if pad_frames else None,
+                    time_embedding=time_embeddings[step],
                 )
                 if pad_frames:
                     _zero_padded_cnn_cache(step_cnn, estimator, pad_frames)
-                conditional, unconditional = estimate.split(batch_size, dim=0)
-                velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional
-                x = x + dt * velocity
-                _zero_padded_frames(x, mel_frames if pad_frames else None)
-                time = time + dt
-                if step + 1 < self.n_timesteps:
-                    dt = timeline[step + 2] - time[0]
+                x = _euler_step(x, estimate, dt_steps[step], decoder.inference_cfg_rate, batch_size)
+                if pad_frames:
+                    _zero_padded_frames(x, mel_frames)
                 next_cnn.append(step_cnn)
                 if valid_lengths is not None:
                     if ragged_att_cache is None:
@@ -842,28 +1029,40 @@ class BatchedToken2Wav(nn.Module):
         assert next_att_cache is not None
         return x, torch.stack(next_cnn), next_att_cache
 
-    def _split_flow_cache(self, cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]:
+    def _split_flow_cache(
+        self,
+        cache: dict[str, Any],
+        batch_size: int,
+        *,
+        estimator_att_keep: tuple[int, int] | None = None,
+    ) -> list[dict[str, torch.Tensor]]:
+        """Per-request flow caches from a batch's stacked caches.
+
+        ``estimator_att_cache`` is stacked ``(n_t, depth, 2B, heads, L, d)``
+        with CFG rows ``[cond x B | uncond x B]``, or already one tensor per
+        request (Whole-Euler). The streaming trim ``estimator_att_keep`` is
+        applied while each request is copied out, so no trimmed copy of the
+        whole batch is built first: at 16 requests that copy alone was ~7 GiB.
+        """
+        estimator_att = cache["estimator_att_cache"]
+        stacked_att = isinstance(estimator_att, torch.Tensor)
+        if stacked_att:
+            keep_ranges = _att_keep_ranges(int(estimator_att.shape[4]), estimator_att_keep)
+            request_shape = list(estimator_att.shape)
+            request_shape[2] = 2
+            request_shape[4] = sum(length for _, length in keep_ranges)
+            att_rows = estimator_att.unflatten(2, (2, batch_size))
         result: list[dict[str, torch.Tensor]] = []
         for row in range(batch_size):
-            estimator_att = cache["estimator_att_cache"]
-            if estimator_att.dtype == self._estimator_att_cache_dtype:
-                request_att = torch.cat(
-                    (
-                        estimator_att[:, :, row : row + 1],
-                        estimator_att[:, :, batch_size + row : batch_size + row + 1],
-                    ),
-                    dim=2,
-                ).detach()
-            else:
-                request_shape = list(estimator_att.shape)
-                request_shape[2] = 2
+            if stacked_att:
                 request_att = torch.empty(
                     request_shape,
                     device=estimator_att.device,
                     dtype=self._estimator_att_cache_dtype,
                 )
-                request_att[:, :, 0:1].copy_(estimator_att[:, :, row : row + 1])
-                request_att[:, :, 1:2].copy_(estimator_att[:, :, batch_size + row : batch_size + row + 1])
+                _copy_frame_segments(request_att, att_rows[:, :, :, row], keep_ranges)
+            else:
+                request_att = estimator_att[row]
             result.append(
                 {
                     "conformer_cnn_cache": cache["conformer_cnn_cache"][row : row + 1].detach().clone(),
@@ -880,20 +1079,45 @@ class BatchedToken2Wav(nn.Module):
             )
         return result
 
-    def _stack_flow_cache(self, states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]:
+    @staticmethod
+    def _stack_estimator_att(rows: list[torch.Tensor]) -> torch.Tensor:
+        """Stack per-request ``(n_t, depth, 2, ...)`` caches into CFG order ``[cond x B | uncond x B]``."""
+        return torch.cat((*[row[:, :, 0:1] for row in rows], *[row[:, :, 1:2] for row in rows]), dim=2)
+
+    def _stack_flow_cache(
+        self,
+        states: list[BatchedToken2WavState],
+        *,
+        include_estimator_att: bool = True,
+    ) -> dict[str, torch.Tensor | None]:
+        """Stack per-request flow caches for one batched call.
+
+        The estimator attention cache is by far the largest (~0.5 GiB per
+        request in fp32). The Whole-Euler path reads it per request, so it
+        skips stacking it (``include_estimator_att=False``).
+        """
         flows = [state.flow_cache for state in states]
         conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows]
         unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows]
-        conditional_att = [flow["estimator_att_cache"][:, :, 0:1] for flow in flows]
-        unconditional_att = [flow["estimator_att_cache"][:, :, 1:2] for flow in flows]
-        estimator_att = torch.cat((*conditional_att, *unconditional_att), dim=2)
         return {
             "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0),
             "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1),
             "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2),
-            "estimator_att_cache": estimator_att,
+            "estimator_att_cache": (
+                self._stack_estimator_att([flow["estimator_att_cache"] for flow in flows])
+                if include_estimator_att
+                else None
+            ),
         }
 
+    def _whole_euler_active(self) -> bool:
+        wrapper = self._whole_euler_graph_wrapper
+        return wrapper is not None and wrapper.enabled and self._trt_stepper is None
+
+    def _whole_euler_ragged_active(self) -> bool:
+        """Whether Whole-Euler graphs also solve ragged (per-row length) batches."""
+        return self._whole_euler_active() and self._whole_euler_graph_wrapper.ragged_body is not None
+
     def _create_initial_states(
         self,
         features: PromptFeatures,
@@ -913,12 +1137,15 @@ class BatchedToken2Wav(nn.Module):
                 att_cache=None,
             )
             projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1))
+            # The prompt cache never reaches the trim length; ``att_keep`` only
+            # tells the Whole-Euler arena the steady cache length to reserve.
             _, estimator_cnn, estimator_att = self._decode_cfm(
                 hidden.transpose(1, 2).contiguous(),
                 projected_speakers,
                 prompt_mels.transpose(1, 2).contiguous(),
                 cnn_cache=None,
                 att_cache=None,
+                att_keep=(int(prompt_mels.shape[1]), _CACHE_TRIM_SUFFIX),
             )
         flow_cache = {
             "conformer_cnn_cache": conformer_cnn,
@@ -941,7 +1168,17 @@ class BatchedToken2Wav(nn.Module):
         ]
 
     def setup_batch(self, features: PromptFeatures, batch_size: int) -> list[BatchedToken2WavState]:
-        """Reuse read-only prompt conditioning for the same shape and padding."""
+        """Initial states for ``batch_size`` requests sharing one prompt.
+
+        Every request starts from the same prompt tokens, speaker, mels and
+        noise, so the prompt is solved once and all requests share that one
+        read-only state. Decoding writes a state only where Whole-Euler updates
+        an estimator cache in place, which needs a cache held by one request
+        and allocated with room for its new length; the prompt cache is shared
+        and allocated exactly, and the first chunk always grows it. Solving it
+        per batch size recomputed identical rows and cached one ~0.3 GiB
+        estimator cache copy per request.
+        """
         bucket_frames = (
             self._cfm_graph_bucket_frames
             if self._cfm_graph_wrapper is not None and self._cfm_graph_wrapper.enabled
@@ -949,17 +1186,17 @@ class BatchedToken2Wav(nn.Module):
         )
         # Capture the policy before setup: graph capture may disable the wrapper
         # after padding has already been chosen for these initial states.
-        cache_key = (features.cache_key, batch_size, bucket_frames)
-        cached = self._setup_cache.get(cache_key)
-        if cached is not None:
+        cache_key = (features.cache_key, bucket_frames)
+        state = self._setup_cache.get(cache_key)
+        if state is not None:
             self._setup_cache.move_to_end(cache_key)
-            return list(cached)
-        states = self._create_initial_states(features, batch_size)
-        if self._setup_cache_size > 0:
-            self._setup_cache[cache_key] = tuple(states)
-            while len(self._setup_cache) > self._setup_cache_size:
-                self._setup_cache.popitem(last=False)
-        return states
+        else:
+            (state,) = self._create_initial_states(features, 1)
+            if self._setup_cache_size > 0:
+                self._setup_cache[cache_key] = state
+                while len(self._setup_cache) > self._setup_cache_size:
+                    self._setup_cache.popitem(last=False)
+        return [state] * batch_size
 
     @staticmethod
     def _fade_in_out(
@@ -1058,7 +1295,12 @@ class BatchedToken2Wav(nn.Module):
         flush_encoder: bool = False,
     ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]:
         batch_size = int(tokens.shape[0])
-        flow_cache = self._stack_flow_cache(states)
+        # Whole-Euler reads and writes each request's estimator cache directly,
+        # so the batch never holds a stacked copy of it.
+        per_request_att = self._whole_euler_active()
+        flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att)
+        prompt_len = int(features.mels.shape[1])
+        att_keep = (prompt_len, _CACHE_TRIM_SUFFIX)
         speakers = features.speaker_embedding.expand(batch_size, -1)
         with self._autocast(tokens.device):
             hidden, conformer_cnn, conformer_att = self._encode_chunk(
@@ -1074,20 +1316,15 @@ class BatchedToken2Wav(nn.Module):
                 projected_speakers,
                 cond,
                 cnn_cache=flow_cache["estimator_cnn_cache"],
-                att_cache=flow_cache["estimator_att_cache"],
+                att_cache=(
+                    [state.flow_cache["estimator_att_cache"] for state in states]
+                    if per_request_att
+                    else flow_cache["estimator_att_cache"]
+                ),
+                att_keep=att_keep,
             )
 
-        prompt_len = int(features.mels.shape[1])
-        if estimator_att.shape[4] > prompt_len + 100:
-            estimator_att = torch.cat(
-                (estimator_att[..., :prompt_len, :], estimator_att[..., -100:, :]),
-                dim=4,
-            )
-        if conformer_att.shape[3] > prompt_len + 100:
-            conformer_att = torch.cat(
-                (conformer_att[..., :prompt_len, :], conformer_att[..., -100:, :]),
-                dim=3,
-            )
+        conformer_att = _trim_streaming_cache(conformer_att, prompt_len)
         new_flow = self._split_flow_cache(
             {
                 "conformer_cnn_cache": conformer_cnn,
@@ -1096,6 +1333,7 @@ class BatchedToken2Wav(nn.Module):
                 "estimator_att_cache": estimator_att,
             },
             batch_size,
+            estimator_att_keep=att_keep,
         )
         old_mel = torch.cat([state.hift_cache["mel"] for state in states], dim=0)
         old_source = torch.cat([state.hift_cache["source"] for state in states], dim=0)
@@ -1164,10 +1402,17 @@ class BatchedToken2Wav(nn.Module):
                 )
 
         max_frames = self._max_encode_token_frames(states)
-        if any(int(row_tokens.numel()) > max_frames for row_tokens in tokens):
+        per_request_att = self._whole_euler_ragged_active()
+        if (self._whole_euler_active() and not per_request_att) or any(
+            int(row_tokens.numel()) > max_frames for row_tokens in tokens
+        ):
             # decode_batch owns the RelPos-safe encoder slicing used by
-            # non-async Talker dumps. Preserve that path for overlong rows;
-            # ordinary short duplex chunks continue through one padded DiT.
+            # non-async Talker dumps. Preserve that path for overlong rows.
+            # Whole-Euler graphs that cannot run the ragged kernel also take
+            # this path: under concurrency most batches hold some request's
+            # short final chunk, and one padded eager DiT would run the whole
+            # batch ungraphed on a stacked copy of every cache. Otherwise short
+            # duplex chunks share one padded DiT solve, graphed when possible.
             exact_groups: dict[tuple[int, bool], list[int]] = {}
             for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
                 exact_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row)
@@ -1193,7 +1438,8 @@ class BatchedToken2Wav(nn.Module):
         conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size
         for (_, last_chunk), rows in encoder_groups.items():
             group_states = [states[row] for row in rows]
-            group_cache = self._stack_flow_cache(group_states)
+            # The encoder reads only the conformer caches.
+            group_cache = self._stack_flow_cache(group_states, include_estimator_att=False)
             group_tokens = torch.stack([tokens[row] for row in rows], dim=0)
             with self._autocast(group_tokens.device):
                 hidden, conformer_cnn, conformer_att = self._encode_chunk(
@@ -1219,7 +1465,10 @@ class BatchedToken2Wav(nn.Module):
         for row, hidden in enumerate(resolved_hidden):
             padded_hidden[row, : int(hidden.shape[1])].copy_(hidden[0])
 
-        flow_cache = self._stack_flow_cache(states)
+        # A graphed ragged solve reads and writes each request's estimator
+        # cache directly (trimmed on the way out), like _decode_batch_once.
+        flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att)
+        prompt_len = int(features.mels.shape[1])
         speakers = features.speaker_embedding.expand(batch_size, -1)
         with self._autocast(padded_hidden.device):
             projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1))
@@ -1229,33 +1478,24 @@ class BatchedToken2Wav(nn.Module):
                 projected_speakers,
                 cond,
                 cnn_cache=flow_cache["estimator_cnn_cache"],
-                att_cache=flow_cache["estimator_att_cache"],
+                att_cache=(
+                    [state.flow_cache["estimator_att_cache"] for state in states]
+                    if per_request_att
+                    else flow_cache["estimator_att_cache"]
+                ),
                 valid_lengths=hidden_lengths,
+                att_keep=(prompt_len, _CACHE_TRIM_SUFFIX),
             )
 
-        prompt_len = int(features.mels.shape[1])
         assert isinstance(estimator_att, list)
         new_flow: list[dict[str, torch.Tensor]] = []
         for row, row_estimator_att in enumerate(estimator_att):
             conformer_cnn = conformer_cnn_rows[row]
             conformer_att = conformer_att_rows[row]
             assert conformer_cnn is not None and conformer_att is not None
-            if row_estimator_att.shape[4] > prompt_len + 100:
-                row_estimator_att = torch.cat(
-                    (
-                        row_estimator_att[..., :prompt_len, :],
-                        row_estimator_att[..., -100:, :],
-                    ),
-                    dim=4,
-                )
-            if conformer_att.shape[3] > prompt_len + 100:
-                conformer_att = torch.cat(
-                    (
-                        conformer_att[..., :prompt_len, :],
-                        conformer_att[..., -100:, :],
-                    ),
-                    dim=3,
-                )
+            # A graphed solve already trimmed the estimator cache; trimming is idempotent.
+            row_estimator_att = _trim_streaming_cache(row_estimator_att, prompt_len)
+            conformer_att = _trim_streaming_cache(conformer_att, prompt_len)
             new_flow.append(
                 {
                     "conformer_cnn_cache": conformer_cnn.detach().clone(),

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -268,14 +268,34 @@ class MiniCPMO45Code2Wav(nn.Module):
             raise ValueError(f"Invalid MiniCPM-o connector chunk configuration: {self._connector_config}")
         raw_capture_batch_sizes = extra.get("hift_graph_capture_batch_sizes")
         capture_batch_sizes = [1] if raw_capture_batch_sizes is None else raw_capture_batch_sizes
+        max_serial_batch = extra.get("max_serial_batch")
+        max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch)
         self._hift_graph_config = {
             "enabled": bool(extra.get("enable_hift_graph", False)),
             "capture_batch_sizes": capture_batch_sizes,
+            "max_serial_batch": max_serial_batch,
         }
+        enable_whole_euler = extra.get("enable_whole_euler")
+        max_graph_batch_raw = extra.get("max_graph_batch")
+        max_graph_batch = int(max_graph_batch_raw) if max_graph_batch_raw is not None else None
+        micro_batch_size_raw = extra.get("micro_batch_size")
+        if micro_batch_size_raw is not None:
+            micro_batch_size = int(micro_batch_size_raw)
+        else:
+            # The Whole-Euler arena reserves one attention cache per micro-batch
+            # row, so size it for the most requests this stage ever batches.
+            max_num_seqs = getattr(getattr(vllm_config, "scheduler_config", None), "max_num_seqs", None)
+            micro_batch_size = min(int(max_num_seqs), max_graph_batch or 16) if max_num_seqs else None
         self._cfm_graph_config = {
             "enabled": bool(extra.get("enable_cfm_graph", False)),
             "max_graphs": int(extra.get("cfm_max_graphs", 32)),
             "bucket_frames": int(extra.get("cfm_graph_bucket_frames", 0)),
+            "capture_frames": extra.get("cfm_graph_capture_frames"),
+            "enable_whole_euler": enable_whole_euler is None or bool(enable_whole_euler),
+            "max_serial_batch": max_serial_batch,
+            "max_graph_batch": max_graph_batch,
+            "micro_batch_size": micro_batch_size,
+            "pad_max_rows": extra.get("whole_euler_pad_max_rows"),
         }
         self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS))
         if self._ref_max_seconds <= 0: