[Bugfix/Feature] Remove Hardcoded Flash Attention in Bagel & Support GQA in SDPA Backend (#3728)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -56,10 +56,7 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
         packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long),
         packed_seqlens=torch.tensor([seq_len], dtype=torch.int),
         packed_position_ids=torch.arange(seq_len, dtype=torch.long),
-        packed_indexes=torch.arange(seq_len, dtype=torch.long),
         past_key_values=NaiveCache(1),
-        key_values_lens=torch.tensor([0], dtype=torch.int),
-        packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
         num_timesteps=NUM_TIMESTEPS,
         timestep_shift=1.0,
         cfg_text_scale=1.0,
@@ -68,11 +65,8 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
     if cfg:
         base |= dict(
             cfg_text_scale=4.0,
-            cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long),
             cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long),
             cfg_text_past_key_values=NaiveCache(1),
-            cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int),
-            cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
         )
     return base
 

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -85,6 +85,7 @@ class SDPAImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.requires_gqa = num_heads != num_kv_heads
         if backend_kwargs:
             logger.warning("SDPAImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
@@ -111,6 +112,7 @@ class SDPAImpl(AttentionImpl):
             dropout_p=0.0,
             is_causal=self.causal,
             scale=self.softmax_scale,
+            enable_gqa=self.requires_gqa,
         )
         out = output.permute(0, 2, 1, 3)
         return out

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -307,12 +307,9 @@ def extract_bagel_context(
     packed_vae_position_ids: torch.LongTensor,
     packed_text_ids: torch.LongTensor,
     packed_text_indexes: torch.LongTensor,
-    packed_indexes: torch.LongTensor,
     packed_position_ids: torch.LongTensor,
     packed_seqlens: torch.IntTensor,
-    key_values_lens: torch.IntTensor,
     past_key_values: Any,
-    packed_key_value_indexes: torch.LongTensor,
     **kwargs: Any,
 ) -> CacheContext:
     """
@@ -326,12 +323,9 @@ def extract_bagel_context(
         packed_vae_position_ids: Position IDs for VAE tokens
         packed_text_ids: Text token IDs
         packed_text_indexes: Indexes for text tokens in packed sequence
-        packed_indexes: Global indexes
         packed_position_ids: Global position IDs
         packed_seqlens: Sequence lengths
-        key_values_lens: KV cache lengths
         past_key_values: KV cache
-        packed_key_value_indexes: KV cache indexes
         **kwargs: Additional keyword arguments
 
     Returns:
@@ -375,10 +369,7 @@ def extract_bagel_context(
             packed_query_sequence=packed_sequence,
             query_lens=packed_seqlens,
             packed_query_position_ids=packed_position_ids,
-            packed_query_indexes=packed_indexes,
             past_key_values=past_key_values,
-            key_values_lens=key_values_lens,
-            packed_key_value_indexes=packed_key_value_indexes,
             update_past_key_values=False,
             is_causal=False,
             **extra_inputs,

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -793,13 +793,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 cfg_renorm_type=gen_params.cfg_renorm_type,
                 **generation_input,
                 cfg_text_packed_position_ids=generation_input_cfg_text["cfg_packed_position_ids"],
-                cfg_text_packed_query_indexes=generation_input_cfg_text["cfg_packed_query_indexes"],
-                cfg_text_key_values_lens=generation_input_cfg_text["cfg_key_values_lens"],
-                cfg_text_packed_key_value_indexes=generation_input_cfg_text["cfg_packed_key_value_indexes"],
                 cfg_img_packed_position_ids=generation_input_cfg_img["cfg_packed_position_ids"],
-                cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"],
-                cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"],
-                cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"],
                 return_trajectory_latents=req.sampling_params.return_trajectory_latents,
                 scheduler=self.scheduler,
                 scheduler_kwargs=self.scheduler_kwargs,