fix: update batch attention logic to handle padding and size conditions (#436)

Author: rebel-jaehunryu

vllm_rbln/v1/attention/backends/flash_attention.py

@@ -1203,7 +1203,7 @@ def build(
             query_start_loc=query_start_loc,
             max_seq_len=query_max_seq_len,
             seq_lens=seq_lens_tensor.to(self.device)
-            if not self.is_batch_attention_opt or is_prefills[0]
+            if not self.is_batch_attention_opt or is_prefills[0] or batch_pad <= 1
             else seq_idx.to(self.device),
             block_tables=block_tables_tensor.to(self.device),
             slot_mapping=slot_mapping,
@@ -1437,15 +1437,15 @@ def forward(
                     value,
                     kv_cache,
                     attn_metadata.cache_seq_lens.to(torch.int32)
-                    if self.is_batch_attention_opt
+                    if self.is_batch_attention_opt and b_size > 1
                     else attn_metadata.cache_seq_lens,
                     attn_metadata.cache_offsets,
                     self.scale,
                     attn_metadata.local_block_tables,
                     self.scale,  # dummy
                 ]
                 if not envs.VLLM_RBLN_USE_CUSTOM_KERNEL:
-                    if self.is_batch_attention_opt:
+                    if self.is_batch_attention_opt and b_size > 1:
                         decode_args.append(attn_metadata.swa_attn_masks)
                     else:
                         decode_args.append(None)