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

[alex-jw-brooks]

Purpose

Fix: #3301
Follow-up on Bagel fixes to support Bagel on other attention backends.

• Removes the hard-coded Flash attention code from the Bagel to support other backends
• Removes dead args & refactors a bit to try to make the gen / und paths more readable
• Allow GQA for SDPA backend (needed for Bagel)

Outputs do change a bit, but I think this is due to the hardcoded flash_attn_varlen_func vs flash_attn_func, see details below for a minimal patch that'll flip the input to what this branch produces.

Minimal test script:

import argparse
import os

OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("prompt", nargs="?", default="A cute cat sitting on a windowsill")
    parser.add_argument("--ulysses-degree", type=int, default=1)
    parser.add_argument("--ring-degree", type=int, default=1)
    return parser.parse_args()


def main():
    args = parse_args()

    from vllm_omni.entrypoints.omni import Omni
    has_parallel = args.ulysses_degree > 1 or args.ring_degree > 1
    parallel_kwargs = {} if not has_parallel else {
        "parallel_config": {
            "ulysses_degree": args.ulysses_degree,
            "ring_degree": args.ring_degree,
        }
    }

    omni = Omni(
        model="ByteDance-Seed/BAGEL-7B-MoT",
        **parallel_kwargs,
    )

    formatted_prompt = {
        "prompt": f"<|im_start|>{args.prompt}<|im_end|>",
        "modalities": ["image"],
    }

    params_list = omni.default_sampling_params_list
    omni_outputs = list(omni.generate(prompts=[formatted_prompt], sampling_params_list=params_list))

    ar_output = omni_outputs[0]
    ro = getattr(ar_output, 'request_output', None)
    if ro and getattr(ro, 'outputs', None):
        token_ids = [tid for o in ro.outputs for tid in (o.token_ids or [])]
        print(f"AR stage tokens: {token_ids}")

    save_path = os.path.join(OUTPUT_DIR, "output_ulysses_{}_ring_{}.png".format(args.ulysses_degree, args.ring_degree))
    omni_outputs[1].images[0].save(save_path)
    print(f"Saved {save_path}")


if __name__ == "__main__":
    main()

On main with no SP & with SP you should get:

On this branch with no SP & with SP you should get:

You can also run it with DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA and make sure you get the same result.

To confirm the difference is due to flash_attn_varlen_func vs flash_attn_func and not something else, you can use this as a minimal patch on main :

--- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py
+++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py
@@ -616,14 +616,11 @@ class PackedAttentionMoT(nn.Module):
         cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(query_lens, dim=0), (1, 0))
         cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0))
 
-        packed_attn_output = flash_attn_varlen_func(
-            q=packed_query_states,
-            k=merged_key_states,
-            v=merged_value_states,
-            cu_seqlens_q=cu_seqlens_q.to(torch.int32),
-            cu_seqlens_k=cu_seqlens_k.to(torch.int32),
-            max_seqlen_q=max(query_lens).item(),
-            max_seqlen_k=max(key_values_lens).item(),
+        from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_func
+        packed_attn_output = flash_attn_func(
+            packed_query_states.unsqueeze(0),
+            merged_key_states.unsqueeze(0),
+            merged_value_states.unsqueeze(0),
             causal=is_causal,
         )
         packed_attn_output = packed_attn_output.reshape(-1, self.q_size)

You should see the output flip to what we get on this branch.

@princepride @Gaohan123 can you please take a look? Thanks!

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[alex-jw-brooks]

Also @gcanlin, can you please take a look at this to make sure you are ok with this change? GQA support is needed for Bagel, but this is the only change in the attn backend 🙂 otherwise it'll crash with a dims mismatch

[hsliuustc0106]

Review cannot proceed until docs build passes. Please fix the docs build failure.

[alex-jw-brooks]

Thanks @hsliuustc0106, resolved the docs issue!

[gcanlin]

@alex-jw-brooks Thanks! That's what I expect. I will check the details later.

[hsliuustc0106]

Review Summary

Reviewed PR #3728: Remove Hardcoded Flash Attention in Bagel & Support GQA in SDPA Backend.

What I validated

• Review gates: DCO, pre-commit, build (3.11/3.12) all passing. Docs build still pending but author confirmed docs issue resolved.
• Correctness of attention refactor: The refactored _forward_gen and _forward_und correctly replace flash_attn_varlen_func with the abstracted DiffusionAttention layer. The SP path preserves the joint mechanism. The non-SP gen path correctly concatenates text+vae tokens for bidirectional attention. The und path correctly handles KV cache appending and causal masking.
• Multi-request batching concern investigated: The Bagel pipeline (BagelPipeline.forward()) enforces single-request — it warns and truncates to 1 prompt if multiple are provided. All downstream calls (prepare_prompts, generate_text, generate_image) receive single-element inputs. The varlen→batched attention change is therefore safe in practice. The CFG batching in generate_image is multi-branch (gen/cfg_text/cfg_img), not multi-request.
• SDPA enable_gqa=True: Safe for all models — no-op when num_heads == num_kv_heads, correct expansion when num_kv_heads < num_heads. PyTorch >= 2.5.0 requirement is met. Other models that pre-expand K/V (e.g., HunyuanImage3) will see it as a no-op.
• Output difference: Author provides thorough evidence (visual comparison + minimal patch) attributing the change to flash_attn_varlen_func vs flash_attn_func numerical differences. This is a known artifact of switching between varlen and non-varlen Flash Attention APIs.
• Test update: test_trajectory_recording.py correctly updated to match new API (removed dead params).
• Dead arg cleanup: The removal of packed_query_indexes, key_values_lens, packed_key_value_indexes across the entire call chain (model, pipeline, teacache extractor, tests) is consistent and thorough.

Non-blocking observations

1. _forward_und docstring: Step numbering goes 1-6 then skips to 8 (line ~530 in new code). Trivial fix.

2. Internal API vs actual usage mismatch: The model's internal API still maintains packed-format infrastructure (packed_seqlens, packed_text_indexes, packed_vae_token_indexes) that suggests multi-request batching capability, but the pipeline enforces single-request. The author's TODO about removing packing is a good forward step — this would simplify the model significantly.

3. enable_gqa=True scope: This change applies to ALL diffusion models using the SDPA backend, not just Bagel. While it's safe (no-op for non-GQA models), it's worth being aware of the cross-model impact.

Verdict

No blocking issues found. The refactoring is a clear improvement in readability and backend flexibility. The evidence is solid. Approving from a correctness standpoint.

---

Reviewed by Claude Code with glm-5.1

[alex-jw-brooks]

Thanks @hsliuustc0106! Made the changes, should be ready for another look when you have a moment

[ZJY0516]

@codex review

[ZJY0516]

overall LGTM

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b7e654ac47

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve the cache offset when using causal SDPA

When past_key_values is populated and the selected backend is SDPA, this calls causal attention with query containing only the new tokens but key containing cache + new. PyTorch SDPA's is_causal=True uses an upper-left non-square mask, so a decode step with q_len=1 can only attend to the first cached key instead of all previous keys plus the current token. This breaks cached autoregressive text generation even for a single request; use an offset/lower-right causal bias or avoid SDPA's bare is_causal for cached decode.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Keep batched CFG branches isolated in attention

In non-SP image generation with CFG enabled, _generate_image_single builds cfg_batched by concatenating the main/text/img branches and passes batched_query_lens, but this path flattens all text and VAE tokens into a single SDPA sequence and never uses those lengths. As a result, branch 0/1/2 tokens attend to each other instead of remaining independent CFG branches, corrupting the guidance prediction whenever batched CFG is used.

Useful? React with 👍 / 👎.

[hsliuustc0106]

lgtm