[Whisper] Batch encoder forward for concurrent prefill requests (#22361)

python/sglang/srt/models/whisper.py

@@ -446,30 +446,31 @@ class WhisperForConditionalGeneration(torch.nn.Module):
                 forward_batch.encoder_cached if forward_batch.encoder_cached else []
             )
 
-            encoder_list = []
-            for i, (mm_input, cached) in enumerate(
-                zip(mm_inputs_list, encoder_cached_list)
-            ):
+            # Collect features from all uncached requests for batched encoding
+            features_to_encode = []
+            for mm_input, cached in zip(mm_inputs_list, encoder_cached_list):
                 if cached or mm_input is None or not mm_input.mm_items:
                     continue
-
                 features = mm_input.mm_items[0].feature
                 if features.ndim == 2:
                     features = features.unsqueeze(0)
+                features_to_encode.append(features.to(dtype))
 
-                encoder_len = features.shape[-1] // 2
-                encoder_position_ids = torch.arange(encoder_len).to(
-                    features.device, non_blocking=True
+            if features_to_encode:
+                # Batch all features and run encoder once instead of sequentially
+                features_batch = torch.cat(features_to_encode, dim=0)
+                encoder_len = features_batch.shape[-1] // 2
+                encoder_position_ids = torch.arange(
+                    encoder_len, device=features_batch.device
                 )
 
-                req_encoder_output = self.encoder(
-                    features.to(dtype), encoder_position_ids, forward_batch
+                batched_output = self.encoder(
+                    features_batch, encoder_position_ids, forward_batch
+                )
+                # Flatten [N, seq_len, dim] → [N*seq_len, dim] for cross-attention
+                encoder_hidden_states = batched_output.reshape(
+                    -1, batched_output.shape[-1]
                 )
-                req_encoder_output = req_encoder_output.squeeze(0)
-                encoder_list.append(req_encoder_output)
-
-            if encoder_list:
-                encoder_hidden_states = torch.cat(encoder_list, dim=0)
 
         decoder_outputs = self.decoder(
             input_ids, encoder_hidden_states, forward_batch, positions