Add KunLunXin platform support (#63)

# Description

Add KunLunXin platform support for Megatron-LM-FL so XPU environments
can be detected and routed through XME initialization before falling
back to CUDA.

## Type of change

- [x] New feature (non-breaking change which adds functionality)
- [ ] Infra/Build change (changes to CI/CD workflows or build scripts)
- [ ] Code refactoring
- [ ] Documentation change
- [ ] Bug fix
- [ ] Breaking change

## Changes

- Add `PlatformKunLunXin`, inheriting from `PlatformCUDA` for XMLIR
CUDA-compatible execution.
- Detect KunLunXin/XPU environments via `XPU=1` or `xpu-smi`.
- Register the KunLunXin platform before CUDA so XPU machines select it
first.
- Defer XME initialization through `ensure_xme_init()` to avoid circular
imports during platform registration.

## Checklist

- [x] I have read and followed the contributing guidelines
- [x] The functionality is complete
- [ ] I have commented my code, particularly in coverage report
uploading steps
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added/updated tests that prove my feature works
- [ ] New and existing unit tests pass locally

Author: Jonahmoon

megatron/plugin/platform/platform_kunlunxin.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2026, Baidu Inc. All rights reserved.
+"""KunLunXin XPU Platform for Megatron-LM-FL.
+
+Inherits from PlatformCUDA since XMLIR makes XPU appear as CUDA.
+XME init is deferred to first use (via _ensure_xme_init) to avoid circular import
+during platform registration (parallel_state calls get_platform() at module level).
+"""
+
+import os
+import shutil
+
+from .platform_cuda import PlatformCUDA
+
+_XME_INITIALIZED = False
+
+
+class PlatformKunLunXin(PlatformCUDA):
+
+    def __init__(self):
+        super().__init__()
+        self._name = "kunlunxin"
+
+    def is_available(self):
+        """Detect KunLunXin XPU. Does NOT trigger XME init to avoid circular import."""
+        xpu_flag = os.getenv("XPU", "")
+        if xpu_flag in ("1", "True", "true", "TRUE"):
+            return True
+        if shutil.which("xpu-smi") is not None and shutil.which("nvidia-smi") is None:
+            return True
+        return False
+
+    @staticmethod
+    def ensure_xme_init():
+        """Call once after get_platform() to init XME patches. Idempotent."""
+        global _XME_INITIALIZED
+        if _XME_INITIALIZED:
+            return
+        _XME_INITIALIZED = True
+        try:
+            from xmegatron_ext import megatron_plugin_init
+            megatron_plugin_init(use_version="0.17.1", check_version=True)
+        except Exception:
+            pass
+
+    def device_name(self, device_index=None):
+        if device_index is None:
+            return "kunlunxin"
+        return f"kunlunxin:{device_index}"
+
+    def current_device_name(self):
+        return f"kunlunxin:{super().current_device()}"

megatron/plugin/platform/platform_manager.py

@@ -30,6 +30,11 @@ def get_platform():
     elif "enflame" in PLATFORMS.keys() and PLATFORMS["enflame"].is_available():
         cur_platform = PLATFORMS["enflame"]
         print(f"Megatron-LM-FL Platform: enflame Selected")
+    elif "kunlunxin" in PLATFORMS.keys() and PLATFORMS["kunlunxin"].is_available():
+        cur_platform = PLATFORMS["kunlunxin"]
+        print(f"Megatron-LM-FL Platform: kunlunxin Selected")
+        # Deferred XME init (after platform selection to avoid circular import)
+        cur_platform.ensure_xme_init()
     elif "cpu" in PLATFORMS.keys() and PLATFORMS["cpu"].is_available():
         cur_platform = PLATFORMS["cpu"]
         print(f"Megatron-LM-FL Platform: cpu Selected")

megatron/plugin/platform/platform_register.py

@@ -51,4 +51,10 @@ def register_platforms() -> None:
     if platform_enflame.is_available():
         PLATFORMS["enflame"] = platform_enflame # use lower keys: enflame
         print(f"Megatron-LM-FL Platform: enflame Registered")
-    
+    
+    # Register KunLunXin Platform
+    from .platform_kunlunxin import PlatformKunLunXin
+    platform_kunlunxin = PlatformKunLunXin()
+    if platform_kunlunxin.is_available():
+        PLATFORMS["kunlunxin"] = platform_kunlunxin
+        print(f"Megatron-LM-FL Platform: kunlunxin Registered")