[Enhancement] Fix header level warnings and adapt flash_attn of Ming audio VAE (#3146)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>

vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py

@@ -10,18 +10,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Model
+from transformers.utils import is_flash_attn_2_available
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
-try:
-    import flash_attn  # noqa: F401
-except (ImportError, ModuleNotFoundError):
-    flash_attn = None
-    logger.warning(
-        "flash_attn is not available, the model may not yield the "
-        "exactly same result as the transformers implementation "
-        "in the audio tower part."
-    )
 
 
 class AudioVAEConfig(PretrainedConfig):
@@ -200,8 +192,13 @@ class Decoder(nn.Module):
     def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1):
         super().__init__()
         config = Qwen2Config.from_dict(config_dict=decoder_args)
-        if flash_attn is None:
+        if is_flash_attn_2_available():
+            config._attn_implementation_autoset = True
+            config._attn_implementation = "flash_attention_2"
+        else:
             config._attn_implementation = "sdpa"
+
+        logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation)
         self.decoder = Qwen2Model(config)
         self.output_dim = output_dim
         self.latent_dim = latent_dim
@@ -276,8 +273,13 @@ class Encoder(nn.Module):
     def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1):
         super().__init__()
         config = Qwen2Config.from_dict(config_dict=encoder_args)
-        if flash_attn is None:
+        if is_flash_attn_2_available():
+            config._attn_implementation_autoset = True
+            config._attn_implementation = "flash_attention_2"
+        else:
             config._attn_implementation = "sdpa"
+
+        logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation)
         self.encoder = Qwen2Model(config)
         self.input_dim = input_dim
         self.hop_size = hop_size

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -885,6 +885,10 @@ class MingAudioGenerator:
 
         self._use_cuda_graphs = use_cuda_graphs
 
+        # For FA2, let it see a full-length seq Q
+        # trailing latent frames prepended on each decode call
+        self._vae_decode_pad_frames = 32
+
     @cached_property
     def _sampler_pool(self) -> CFMGraphExecutorPool | None:
         device = next(self._model.parameters()).device
@@ -1087,21 +1091,33 @@ class MingAudioGenerator:
     # VAE streaming decode
     def _stream_decode(self, latents: list[torch.Tensor]) -> torch.Tensor:
         sr = int(self._audio_vae.config.sample_rate)
-        vae_cache = {"past_key_values": None, "stream_state": (None, None, None)}
+        decode_pad: torch.Tensor | None = None
         sil_cache: dict | None = None
         wav_chunks: list[torch.Tensor] = []
 
         for i, lat in enumerate(latents):
             last_chunk = i == (len(latents) - 1)
-            speech, stream_state, past_key_values = self._audio_vae.decode(
-                lat,
-                past_key_values=vae_cache["past_key_values"],
-                use_cache=True,
-                stream_state=vae_cache["stream_state"],
-                last_chunk=last_chunk,
+
+            if decode_pad is not None:
+                vae_input = torch.cat([decode_pad, lat], dim=1)
+                pad_frames = decode_pad.shape[1]
+            else:
+                vae_input = lat
+                pad_frames = 0
+
+            # Stateless, no KV cache accum intentionally.
+            speech, _, _ = self._audio_vae.decode(
+                vae_input,
+                use_cache=False,
+                stream_state=(None, None, None),
+                last_chunk=True,
             )
-            vae_cache = {"past_key_values": past_key_values, "stream_state": stream_state}
-            speech_chunk = speech[0].detach().float()
+
+            total_frames = vae_input.shape[1]
+            dcs = speech.shape[-1] // total_frames
+
+            # keep only the new audio.
+            speech_chunk = speech[:, :, pad_frames * dcs :][0].detach().float()
             speech_chunk, sil_cache = silence_holder(
                 speech_chunk,
                 sr,
@@ -1111,6 +1127,9 @@ class MingAudioGenerator:
             if speech_chunk.numel() > 0:
                 wav_chunks.append(speech_chunk)
 
+            # Advance the sliding buffer
+            decode_pad = vae_input[:, -self._vae_decode_pad_frames :, :].detach()
+
         if not wav_chunks:
             device = next(self._model.parameters()).device
             dtype = next(self._model.parameters()).dtype