[Feature] Add FlashAttention3 as a backend for VisionAttention (#5764)

Co-authored-by: othame <chenzhu_912@zju.edu.cn> Co-authored-by: Mick <mickjagger19@icloud.com> Co-authored-by: Yi Zhang <1109276519@qq.com>
2025-05-09 01:01:19 +08:00
parent f1ff736d68
commit fa7d7fd9e5
11 changed files with 328 additions and 186 deletions
--- a/python/sglang/srt/models/mllama.py
+++ b/python/sglang/srt/models/mllama.py
@@ -203,7 +203,7 @@ class MllamaVisionEncoderLayer(nn.Module):
            use_qkv_parallel=True,
            quant_config=quant_config,
            dropout=0.0,
-            use_context_forward=False,
+            qkv_backend="sdpa",
            softmax_in_single_precision=False,
            flatten_batch=False,
            prefix=add_prefix("self_attn", prefix),