permute fusion and padded mla attention

[wenxie-amd]

1. Support moe_permute_fusion

   • Updated to align with the latest Megatron implementation.
   • Added support for the newly required permutation logic in moe_permute_fusion.

2. Support for fused_padded_mla_attention

   • Enables fused attention even when the QK head dim is 192 and V head dim is 128, as seen in DeepSeek-style models.
   • Pads the V head dim to match QK, allowing the use of flash-attn or TE fused attention with uniform head dimension (192).

3. Fix TE flash-attn version compatibility

   • Updated _flash_attn_max_version from PkgVersion("2.7.3") to PkgVersion("3.0.0.post1"), ensuring compatibility with newer versions of flash-attn.

4. Support HSA_NO_SCRATCH_RECLAIM configuration

   • HSA_NO_SCRATCH_RECLAIM can now be configured via environment variables and is properly passed into Slurm jobs for AMD ROCm tuning.

[lhzhang333]

LGTM