[BugFix] Fix cascade attention - RuntimeError: scheduler_metadata must have shape (metadata_size) (vllm-project#17283)

Signed-off-by: Lucas Wilkinson <[email protected]>
Signed-off-by: Mu Huai <[email protected]>

Author: LucasWilkinson

vllm/v1/attention/backends/flash_attn.py

@@ -372,7 +372,7 @@ def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
             suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(
                 self.runner.device)
             prefix_scheduler_metadata = schedule(
-                batch_size=num_reqs,
+                batch_size=1,
                 cu_query_lens=cu_prefix_query_lens,
                 max_query_len=num_actual_tokens,
                 seqlens=prefix_kv_lens,