feat(megatron-lm): reduce extra qkv transpose in attn (#645)

Keep input layout to SBHD layout to reduce extra q,k,v transpose in
attention.

Co-authored-by: RuibinCheung <ruibzhan@amd.com>

Author: Xiaoming-AMD

primus/backends/megatron/core/extensions/primus_turbo.py

@@ -458,16 +458,20 @@ def forward(
         packed_seq_params: PackedSeqParams = None,
     ):
         """Forward."""
+        SUPPORTED_QKV_FORMATS = "sbhd"
+
         packed_seq_kwargs = (
             {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
             if packed_seq_params is not None
             else {}
         )
 
         qkv_format = packed_seq_kwargs.get("qkv_format", self.qkv_format)
-        assert qkv_format in ("sbhd", "bhsd"), "qkv_format only support bshd, but got {qkv_format}"
-        if qkv_format == "sbhd":
-            query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
+        assert (
+            qkv_format in SUPPORTED_QKV_FORMATS
+        ), f"qkv_format only support {SUPPORTED_QKV_FORMATS}, but got {qkv_format}"
+        # NOTE(ruibin): The layout of q, k and v is (S, B, H, D). But attn accept the shape of qkv is (B, S, H, D).
+        query, key, value = [x.permute(1, 0, 2, 3) for x in (query, key, value)]
         mask_type = attn_mask_type.name
         if mask_type == AttnMaskType.causal.name:
             causal = True
@@ -523,9 +527,10 @@ def forward(
             **self.attn_kwargs,
         )
 
-        o = o.reshape(o.shape[0], o.shape[1], -1).transpose(0, 1)
-        if not o.is_contiguous():
-            o = o.contiguous()
+        # NOTE(ruibin): The output of attn is BSHD. Use permute to convert the layout to SBHD.
+        o = o.permute(1, 0, 2, 3).contiguous()
+        o = o.view(o.shape[0], o.shape[1], -1)
+
         return o