Add new register op get_attention_backend for METAX (#31)

# Description

Add new register op get_attention_backend for METAX

Fixes # (issue)

## Type of change

- [ ] New feature (non-breaking change which adds functionality)

## Changes

Please list the changes introduced in this PR:

- Add register for get_attention_backend in register_ops.py
- Add implement of get_attention_backend in metax.py

# Checklist:

- [x] I have read and followed the [contributing
guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst)
- [x] The functionality is complete
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes

Author: dinghaodhd

transformer_engine/plugin/core/backends/vendor/metax/metax.py

@@ -158,6 +158,23 @@ def get_flash_attention_class(self):
         from .flash_attention import FlashAttentionMETAX
         return FlashAttentionMETAX
 
+    def get_attention_backend(self, attention_params=None):
+        # Import the metax get_attention_backend function
+        try:
+            from transformer_engine_metax.pytorch.attention.dot_product_attention import utils
+            return utils.get_attention_backend(attention_params)
+
+        except ImportError as e:
+            raise RuntimeError(
+                f"Failed to import metax FlashAttention: {e}. "
+                "Please ensure flash-attn is installed and transformer_engine_metax is available."
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get_attention_backend: {e}. "
+                f"Attention_params: {self.attention_params}"
+            )
+
     def quantize(
         self,
         tensor: torch.Tensor,

transformer_engine/plugin/core/backends/vendor/metax/register_ops.py

@@ -197,6 +197,8 @@ def register_builtins(registry) -> None:
 
         # FlashAttention class getter
         OpImpl(op_name="get_flash_attention_class", impl_id="vendor.metax", kind=BackendImplKind.VENDOR, fn=_bind_is_available(backend.get_flash_attention_class, is_avail), vendor="METAX", priority=100),
+                # Attention backend selection
+        OpImpl(op_name="get_attention_backend", impl_id="vendor.metax", kind=BackendImplKind.VENDOR, fn=_bind_is_available(backend.get_attention_backend, is_avail), vendor="METAX", priority=100),
     ]
 
     registry.register_many(impls)