perf: reuse MoonViT FA3 max-seqlen metadata (#30878)

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

@@ -458,8 +458,19 @@ class VisionFlash3Attention(nn.Module):
         else:
             cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
             cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
-            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-            max_seqlen = seq_lens.max().item()
+            # Some vision encoders precompute this scalar once per encoder
+            # forward and share it across all of their attention blocks.  Use
+            # that value when available: deriving it here requires a
+            # GPU-to-host sync, so repeating it per block serializes the ViT
+            # launch stream for variable-size images.
+            max_seqlen = kwargs.get("max_seqlen")
+            if max_seqlen is None:
+                seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+                max_seqlen = int(seq_lens.max().item())
+            elif isinstance(max_seqlen, torch.Tensor):
+                max_seqlen = int(max_seqlen.item())
+            else:
+                max_seqlen = int(max_seqlen)
 
             fa_kwargs = dict(
                 cu_seqlens_q=cu_seqlens,

python/sglang/srt/models/kimi_k25.py

@@ -150,6 +150,7 @@ class MoonViTEncoderLayer(nn.Module):
             hidden_states,
             cu_seqlens=cu_seqlens,
             position_embeddings=rope_freqs_cis,
+            max_seqlen=max_seqlen,
         )
 
         hidden_states = residual + hidden_states
@@ -469,7 +470,10 @@ class MoonViT3dEncoder(nn.Module):
             )
         )
 
-        max_seqlen = lengths.max()
+        # FlashAttention needs a host integer.  Compute it once per MoonViT
+        # forward and pass it to every encoder block instead of synchronizing
+        # once per block inside the attention backend.
+        max_seqlen = int(lengths.max().item())
         cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32)
 
         for block in self.blocks: