[Perf] Ming-TTS Stage-0 flow-head speedup: RoPE/QKV fusion, opt-in torch.compile, PIECEWISE cudagraph (#4942)

vllm_omni/deploy/ming_tts_moe.yaml

@@ -31,7 +31,9 @@ stages:
     max_num_seqs: 1
     tensor_parallel_size: 1
     gpu_memory_utilization: 0.60
-    enforce_eager: true
+    enforce_eager: false
+    compilation_config:
+      cudagraph_mode: PIECEWISE
     enable_prefix_caching: false
     async_scheduling: false
     max_num_batched_tokens: 8192

vllm_omni/model_executor/models/common/ming/dit.py

@@ -6,6 +6,19 @@ from torch import nn
 from x_transformers.x_transformers import apply_rotary_pos_emb
 
 
+def _apply_rope_cached(t, cos, sin):
+    """Bit-exact RoPE for ``scale==1`` / interleaved-half.
+    Uses precomputed ``cos``/``sin`` per seq_len to skip per-step trig and cat."""
+    rot_dim = cos.shape[-1]
+    tr = t[..., :rot_dim]
+    x = tr.reshape(*tr.shape[:-1], -1, 2)
+    rot_half = torch.stack((-x[..., 1], x[..., 0]), dim=-1).reshape_as(tr)
+    out = tr * cos + rot_half * sin
+    if rot_dim < t.shape[-1]:  # partial rotary (unused by ming heads); keep general
+        out = torch.cat((out, t[..., rot_dim:]), dim=-1)
+    return out.type(t.dtype)
+
+
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
         super().__init__()
@@ -66,12 +79,33 @@ class Attention(nn.Module):
         self.to_out.append(nn.Dropout(dropout))
         self.pe_attn_head = pe_attn_head
         self.attn_mask_enabled = attn_mask_enabled
+        # cos/sin cache keyed by seq_len (freqs are deterministic in seq_len).
+        self._rope_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+        # TODO: materialize fused weights eagerly outside any torch.compile
+        # scope (e.g. post-weight-load);
+        self._qkv_w: torch.Tensor | None = None
+        self._qkv_b: torch.Tensor | None = None
+
+    def _qkv(self, x):
+        if self._qkv_w is None:
+            self._qkv_w = torch.cat([self.to_q.weight, self.to_k.weight, self.to_v.weight], dim=0)
+            if self.to_q.bias is not None:
+                self._qkv_b = torch.cat([self.to_q.bias, self.to_k.bias, self.to_v.bias], dim=0)
+        qkv = F.linear(x, self._qkv_w, self._qkv_b)
+        return qkv.split(self.inner_dim, dim=-1)
+
+    def _rope_cos_sin(self, freqs):
+        seq_len = freqs.shape[-2]
+        cached = self._rope_cache.get(seq_len)
+        if cached is None:
+            f = freqs.unsqueeze(1) if freqs.ndim == 3 else freqs  # b n d -> b 1 n d for 4D q/k
+            cached = (f.cos(), f.sin())
+            self._rope_cache[seq_len] = cached
+        return cached
 
     def forward(self, x, mask=None, rope=None):
         batch_size = x.shape[0]
-        query = self.to_q(x)
-        key = self.to_k(x)
-        value = self.to_v(x)
+        query, key, value = self._qkv(x)
 
         inner_dim = key.shape[-1]
         head_dim = inner_dim // self.heads
@@ -86,14 +120,21 @@ class Attention(nn.Module):
 
         if rope is not None:
             freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-            if self.pe_attn_head is not None:
-                on = self.pe_attn_head
-                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
-                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+            scale_is_one = xpos_scale is None or (not torch.is_tensor(xpos_scale) and xpos_scale == 1.0)
+            if scale_is_one and self.pe_attn_head is None:
+                # Fast path: cached, cat-free RoPE (bit-exact for scale==1, full rotary).
+                cos, sin = self._rope_cos_sin(freqs)
+                query = _apply_rope_cached(query, cos, sin)
+                key = _apply_rope_cached(key, cos, sin)
             else:
-                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+                q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+                if self.pe_attn_head is not None:
+                    on = self.pe_attn_head
+                    query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
+                    key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+                else:
+                    query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                    key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
 
         if self.attn_mask_enabled and mask is not None:
             valid_sample_indices = mask.any(dim=1)

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -326,6 +326,28 @@ class MingLLMModel(nn.Module):
         try:
             from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
 
+            # TODO(perf):
+            #   (1) compile per-block (dit_model.blocks) instead of the
+            #   whole forward for tighter fusion / fewer graph breaks;
+            #   (2) move the compile + a synthetic warmup forward to a post-weight-load hook
+            dit_model = self.flowloss.cfm.model
+            if not getattr(dit_model, "_ming_compiled", False):
+                try:
+                    # Scoped options (avoid mutating global inductor config).
+                    dit_model.forward = torch.compile(
+                        dit_model.forward,
+                        fullgraph=False,
+                        dynamic=False,
+                        options={
+                            "triton.cudagraphs": False,
+                            "triton.cudagraph_trees": False,
+                        },
+                    )
+                    dit_model._ming_compiled = True
+                    logger.info("Ming CFM DiT torch.compile enabled (inductor cudagraphs off).")
+                except Exception as exc:
+                    logger.warning("Ming CFM torch.compile failed (%s); using uncompiled DiT.", exc)
+
             sampler = CFMSampler(self.flowloss.cfm.model, steps=_CFM_STEPS)
             self._cfm_graph = CFMGraphExecutor(
                 sampler,