Simplify flashinfer dispatch (#1552)

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -14,7 +14,10 @@ import torch.nn as nn
 
 from sglang.global_config import global_config
 from sglang.srt.layers.attention import AttentionBackend
-from sglang.srt.layers.attention.flashinfer_utils import update_flashinfer_indices
+from sglang.srt.layers.attention.flashinfer_utils import (
+    WrapperDispatch,
+    update_flashinfer_indices,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.utils import is_hip
 
@@ -53,10 +56,19 @@ class FlashInferAttnBackend(AttentionBackend):
             device="cuda",
         )
 
+        assert not (
+            model_runner.sliding_window_size is not None
+            and model_runner.has_cross_attention
+        ), "Sliding window and cross attention are not supported together"
+
+        self.num_wrappers = 1
+        self.dispatch_reason = None
         if model_runner.sliding_window_size is not None:
             self.num_wrappers = 2
-        else:
-            self.num_wrappers = 1
+            self.dispatch_reason = WrapperDispatch.SLIDING_WINDOW
+        elif model_runner.has_cross_attention:
+            self.num_wrappers = 2
+            self.dispatch_reason = WrapperDispatch.CROSS_ATTENTION
 
         # NOTE: we do not use ragged attention when there are multiple wrappers
         self.prefill_wrapper_ragged = (
@@ -88,8 +100,12 @@ class FlashInferAttnBackend(AttentionBackend):
         if self.num_wrappers == 1:
             return 0
 
-        # TODO: make sure the idx is related to sliding window size
-        return layer.sliding_window_size == -1
+        if self.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+            return layer.sliding_window_size == -1
+        if self.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
+            return layer.is_cross_attention
+
+        raise ValueError(f"Unknown dispatch reason: {self.dispatch_reason}")
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         if forward_batch.forward_mode.is_decode():