support more general inference case that query length > 1

[yidong72]

Currently the inference requires the query tensor has length 1. However, there are use cases that the query tensor length > 1.
Note, this fix requires

1. TE to use the flash_atten > 2.1 so the correct attention mask is applied according to this document.
2. TE to enable flash attention by removing this line https://github.com/NVIDIA/TransformerEngine/blob/0fbc76af3733ae997394eaf82b78ff9c0498fe91/transformer_engine/pytorch/attention.py#L2732
3. Latest PyTorch version that has the following bug fix so the batch size 1 value tensor has correct stride after transpose ops, which is required for flash_attn. Or there is a work around by using tensor copy.

--- a/megatron/core/transformer/custom_layers/transformer_engine.py
+++ b/megatron/core/transformer/custom_layers/transformer_engine.py
@@ -457,6 +457,10 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
 
         if self.config.apply_rope_fusion and qkv_format == 'bshd':
             query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
+        
+        new_value = torch.zeros_like(key)
+        new_value[:] = value
+        value = new_value
 
         if self.te_forward_mask_type:
             core_attn_out = super().forward(

[github-actions]

Marking as stale. No activity in 60 days.