FA3 speed up: skip len operation and get batch size directly from forward batch (#5969)

Signed-off-by: Lifu Huang <lifu.hlf@gmail.com>

FA3 speed up: skip len operation and get batch size directly from forward batch (#5969)
Signed-off-by: Lifu Huang <lifu.hlf@gmail.com>
1acca3a2 · Lifu Huang · GitHub · 6ea1e6ac · 1acca3a2
Unverified Commit 1acca3a2 authored May 02, 2025 by Lifu Huang Committed by GitHub May 02, 2025
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

python/sglang/srt/layers/attention/flashattention_backend.py python/sglang/srt/layers/attention/flashattention_backend.py +1 -1

No files found.
--- a/python/sglang/srt/layers/attention/flashattention_backend.py
+++ b/python/sglang/srt/layers/attention/flashattention_backend.py
@@ -338,7 +338,7 @@ class FlashAttentionBackend(AttentionBackend):
        """Initialize forward metadata hence all layers in the forward pass can reuse it."""
        metadata = FlashAttentionMetadata()
        seqlens_in_batch = forward_batch.seq_lens
-        batch_size = len(seqlens_in_batch)
+        batch_size = forward_batch.batch_size
        device = seqlens_in_batch.device

        if forward_batch.forward_mode.is_decode_or_idle():