[feat] add biren supa accelerator support (#10844)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/llamafactory/extras/env.py

@@ -33,6 +33,8 @@ def print_env() -> None:
     import transformers
     from transformers.utils import is_torch_cuda_available, is_torch_npu_available
 
+    from .misc import is_torch_supa_available
+
     info = OrderedDict(
         {
             "`llamafactory` version": VERSION,
@@ -46,7 +48,15 @@ def print_env() -> None:
         }
     )
 
-    if is_torch_cuda_available():
+    if is_torch_supa_available():
+        info["PyTorch version"] += " (SUPA)"
+        info["SUPA type"] = torch.supa.get_device_name()
+        info["SUPA number"] = torch.supa.device_count()
+        try:
+            info["SUPA memory"] = f"{torch.supa.mem_get_info()[1] / (1024**3):.2f}GB"
+        except Exception:
+            pass
+    elif is_torch_cuda_available():
         info["PyTorch version"] += " (GPU)"
         info["GPU type"] = torch.cuda.get_device_name()
         info["GPU number"] = torch.cuda.device_count()

src/llamafactory/extras/misc.py

@@ -15,6 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import gc
 import os
 import socket
@@ -38,9 +39,25 @@ from transformers.utils.versions import require_version
 from . import logging
 
 
-_is_fp16_available = is_torch_npu_available() or is_torch_cuda_available()
+@functools.lru_cache
+def is_torch_supa_available() -> bool:
+    """Check if the supa (SUPA PrivateUse1) accelerator is available."""
+    try:
+        import torch_supa  # noqa: F401  # registers the "supa" device
+
+        return hasattr(torch, "supa") and torch.supa.is_available()
+    except Exception:
+        return False
+
+
+_is_fp16_available = is_torch_npu_available() or is_torch_supa_available() or is_torch_cuda_available()
 try:
-    _is_bf16_available = is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported())
+    if is_torch_supa_available():
+        _is_bf16_available = torch.supa.is_bf16_supported()
+    else:
+        _is_bf16_available = is_torch_bf16_gpu_available() or (
+            is_torch_npu_available() and torch.npu.is_bf16_supported()
+        )
 except Exception:
     _is_bf16_available = False
 
@@ -147,6 +164,8 @@ def get_current_device() -> "torch.device":
         device = "xpu:{}".format(os.getenv("LOCAL_RANK", "0"))
     elif is_torch_npu_available():
         device = "npu:{}".format(os.getenv("LOCAL_RANK", "0"))
+    elif is_torch_supa_available():
+        device = "supa:{}".format(os.getenv("LOCAL_RANK", "0"))
     elif is_torch_mps_available():
         device = "mps:{}".format(os.getenv("LOCAL_RANK", "0"))
     elif is_torch_cuda_available():
@@ -163,6 +182,8 @@ def get_device_name() -> str:
         device = "xpu"
     elif is_torch_npu_available():
         device = "npu"
+    elif is_torch_supa_available():
+        device = "supa"
     elif is_torch_mps_available():
         device = "mps"
     elif is_torch_cuda_available():
@@ -190,6 +211,8 @@ def get_device_count() -> int:
         return torch.xpu.device_count()
     elif is_torch_npu_available():
         return torch.npu.device_count()
+    elif is_torch_supa_available():
+        return torch.supa.device_count()
     elif is_torch_mps_available():
         return torch.mps.device_count()
     elif is_torch_cuda_available():
@@ -211,6 +234,8 @@ def get_current_memory() -> tuple[int, int]:
         return torch.xpu.mem_get_info()
     elif is_torch_npu_available():
         return torch.npu.mem_get_info()
+    elif is_torch_supa_available():
+        return torch.supa.mem_get_info()
     elif is_torch_mps_available():
         return torch.mps.current_allocated_memory(), torch.mps.recommended_max_memory()
     elif is_torch_cuda_available():
@@ -225,6 +250,8 @@ def get_peak_memory() -> tuple[int, int]:
         return torch.xpu.max_memory_allocated(), torch.xpu.max_memory_reserved()
     elif is_torch_npu_available():
         return torch.npu.max_memory_allocated(), torch.npu.max_memory_reserved()
+    elif is_torch_supa_available():
+        return torch.supa.max_memory_allocated(), torch.supa.max_memory_reserved()
     elif is_torch_mps_available():
         return torch.mps.current_allocated_memory(), -1
     elif is_torch_cuda_available():
@@ -251,7 +278,11 @@ def infer_optim_dtype(model_dtype: Optional["torch.dtype"]) -> "torch.dtype":
 def is_accelerator_available() -> bool:
     r"""Check if the accelerator is available."""
     return (
-        is_torch_xpu_available() or is_torch_npu_available() or is_torch_mps_available() or is_torch_cuda_available()
+        is_torch_xpu_available()
+        or is_torch_npu_available()
+        or is_torch_supa_available()
+        or is_torch_mps_available()
+        or is_torch_cuda_available()
     )
 
 
@@ -285,6 +316,8 @@ def torch_gc() -> None:
         torch.xpu.empty_cache()
     elif is_torch_npu_available():
         torch.npu.empty_cache()
+    elif is_torch_supa_available():
+        torch.supa.empty_cache()
     elif is_torch_mps_available():
         torch.mps.empty_cache()
     elif is_torch_cuda_available():

src/llamafactory/hparams/parser.py

@@ -33,7 +33,13 @@ from transformers.utils import is_torch_bf16_gpu_available, is_torch_npu_availab
 
 from ..extras import logging
 from ..extras.constants import CHECKPOINT_NAMES, EngineName
-from ..extras.misc import check_dependencies, check_version, get_current_device, is_env_enabled
+from ..extras.misc import (
+    check_dependencies,
+    check_version,
+    get_current_device,
+    is_env_enabled,
+    is_torch_supa_available,
+)
 from ..extras.packages import is_mcore_adapter_available, is_megatron_bridge_available
 from .data_args import DataArguments
 from .evaluation_args import EvaluationArguments
@@ -223,6 +229,20 @@ def _set_env_vars() -> None:
         # avoid use fork method on NPU devices, see https://github.com/hiyouga/LLaMA-Factory/issues/7447
         os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
+    if is_torch_supa_available():
+        # register the fused SUDNN operators (torch.ops.sudnn.*) used by the supa kernel path.
+        # torch_supa itself is auto-loaded on `import torch`, but torch_supa_ext is not.
+        try:
+            import torch_supa_ext  # noqa: F401
+        except Exception:
+            logger.warning_rank0_once(
+                "Failed to import torch_supa_ext; fused supa operators (e.g. sudnn.rms_norm_func) "
+                "will be unavailable and the eager path may be extremely slow."
+            )
+
+        # supa uses fork-unsafe device handles; align with the NPU workaround.
+        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 
 def _verify_model_args(
     model_args: "ModelArguments",
@@ -260,6 +280,9 @@ def _check_extra_dependencies(
         check_version("accelerate-kt", mandatory=True)
 
     if model_args.use_unsloth:
+        if is_torch_supa_available():
+            raise ValueError("`use_unsloth` is not supported on supa devices.")
+
         check_version("unsloth", mandatory=True)
 
     if model_args.enable_liger_kernel:
@@ -496,7 +519,11 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
         raise ValueError("Please use scripts/pissa_init.py to initialize PiSSA in DeepSpeed ZeRO-3.")
 
     if finetuning_args.pure_bf16:
-        if not (is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported())):
+        if not (
+            is_torch_bf16_gpu_available()
+            or (is_torch_npu_available() and torch.npu.is_bf16_supported())
+            or (is_torch_supa_available() and torch.supa.is_bf16_supported())
+        ):
             raise ValueError("This device does not support `pure_bf16`.")
 
         if is_deepspeed_zero3_enabled():

src/llamafactory/model/model_utils/attention.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING
 
 from ...extras import logging
 from ...extras.constants import AttentionFunction
+from ...extras.misc import is_torch_supa_available
 from ...extras.packages import is_torch_version_greater_than
 
 
@@ -63,7 +64,12 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
             model_args.flash_attn = AttentionFunction.DISABLED
 
     if model_args.flash_attn == AttentionFunction.AUTO:
-        return
+        if is_torch_supa_available():
+            # supa has no FlashAttention build, and AUTO would otherwise be resolved against the
+            # cuda-masqueraded device (which may pick FA2); pin SDPA, the fastest supported backend.
+            requested_attn_implementation = "sdpa"
+        else:
+            return
 
     elif model_args.flash_attn == AttentionFunction.DISABLED:
         requested_attn_implementation = "eager"

src/llamafactory/model/model_utils/liger_kernel.py

@@ -37,6 +37,10 @@ def apply_liger_kernel(
     if not is_trainable or not model_args.enable_liger_kernel:
         return
 
+    if get_device_name() == "supa":
+        # liger kernels can not run on supa currently.
+        return
+
     model_type = getattr(config, "model_type", None)
     if model_type == "gemma":
         from liger_kernel.transformers import apply_liger_kernel_to_gemma as apply_liger_kernel

src/llamafactory/model/patcher.py

@@ -480,7 +480,7 @@ def patch_model(
         prepare_model_for_training(model, model_args)
         autocast_projector_dtype(model, model_args)
         add_z3_leaf_module(model)
-
+        
         if getattr(model.config, "model_type", None) in [
             "qwen3_5",
             "qwen3_5_moe",

src/llamafactory/train/callbacks.py

@@ -33,7 +33,14 @@ from typing_extensions import override
 
 from ..extras import logging
 from ..extras.constants import TRAINER_LOG, V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME
-from ..extras.misc import get_peak_memory, is_env_enabled, is_torch_cuda_available, is_torch_npu_available, use_ray
+from ..extras.misc import (
+    get_peak_memory,
+    is_env_enabled,
+    is_torch_cuda_available,
+    is_torch_npu_available,
+    is_torch_supa_available,
+    use_ray,
+)
 from ..extras.packages import is_safetensors_available
 
 
@@ -504,6 +511,8 @@ class ModuleProfilerCallback(TrainerCallback):
     @staticmethod
     def _get_accelerator():
         """Detect available accelerator and return (event_factory, synchronize_fn)."""
+        if is_torch_supa_available():
+            return torch.supa.Event, torch.supa.synchronize
         if is_torch_cuda_available():
             return torch.cuda.Event, torch.cuda.synchronize
         if is_torch_npu_available():

src/llamafactory/train/tuner.py

@@ -291,6 +291,7 @@ class Worker:
         RAY_NOSET_VISIBLE_DEVICES_LIST = [
             "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
             "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES",
+            "RAY_EXPERIMENTAL_NOSET_SUPA_VISIBLE_DEVICES",
         ]
         is_ray_noset_visible_devices = any(os.environ.get(env_var, None) for env_var in RAY_NOSET_VISIBLE_DEVICES_LIST)
         if is_ray_noset_visible_devices:

src/llamafactory/v1/accelerator/helper.py

@@ -45,6 +45,7 @@ class DeviceType(StrEnum):
     META = "meta"
     MPS = "mps"
     NPU = "npu"
+    SUPA = "supa"
     XPU = "xpu"
 
 

src/llamafactory/v1/plugins/model_plugins/kernels/interface.py

@@ -25,11 +25,13 @@ from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel  # noqa: F401
 from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel  # noqa: F401
 from .ops.mlp.npu_swiglu import NpuSwiGluKernel  # noqa: F401
 from .ops.rms_norm.npu_rms_norm import NpuRMSNormKernel  # noqa: F401
+from .ops.rms_norm.supa_rms_norm import SupaRMSNormKernel  # noqa: F401
 from .ops.rope.npu_rope import NpuRoPEKernel  # noqa: F401
 
 
 _AUTO_KERNELS = {
     DeviceType.NPU: ("npu_fused_moe", "npu_fused_rmsnorm", "npu_fused_rope", "npu_fused_swiglu"),
+    DeviceType.SUPA: ("supa_fused_rmsnorm",),
 }
 
 

src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/supa_rms_norm.py

@@ -0,0 +1,158 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The definition of the SUPA fused RMSNorm kernel."""
+
+import types
+from functools import lru_cache
+
+import torch
+
+from ......accelerator.helper import DeviceType, get_current_accelerator
+from ......utils.logging import get_logger
+from ......utils.types import HFModel
+from ...base import BaseKernel, KernelPlugin
+
+
+logger = get_logger(__name__)
+
+
+# Standard-convention RMSNorm classes: ``y = x * rsqrt(mean(x^2) + eps) * weight``.
+# The fused ``sudnn.rms_norm_func`` matches this convention, so only these classes are patched.
+# Variants with a different weight convention (e.g. Gemma's ``1 + weight``, or gated/residual
+# RMSNorm) are intentionally excluded and fall back to the eager path. Add new class names here
+# after verifying they follow the standard convention.
+_SUPPORTED_RMSNORM_CLASSES = frozenset(
+    {
+        "LlamaRMSNorm",
+        "MistralRMSNorm",
+        "MixtralRMSNorm",
+        "Phi3RMSNorm",
+        "InternLM2RMSNorm",
+        "GlmRMSNorm",
+        "DeepseekV2RMSNorm",
+        "DeepseekV3RMSNorm",
+        "Qwen2RMSNorm",
+        "Qwen2MoeRMSNorm",
+        "Qwen2VLRMSNorm",
+        "Qwen2_5_VLRMSNorm",
+        "Qwen2_5_VLTextRMSNorm",
+        "Qwen3RMSNorm",
+        "Qwen3MoeRMSNorm",
+        "Qwen3VLTextRMSNorm",
+        "Qwen3VLMoeTextRMSNorm",
+    }
+)
+
+
+
+@lru_cache
+def _is_sudnn_rms_norm_available() -> bool:
+    """Return whether the fused SUPA RMSNorm operator is registered."""
+    try:
+        import torch_supa  # noqa: F401
+        import torch_supa_ext  # noqa: F401
+    except Exception:
+        return False
+
+    return hasattr(torch.ops, "sudnn") and hasattr(torch.ops.sudnn, "rms_norm_func")
+
+
+class SupaRMSNormFunction(torch.autograd.Function):
+    """Autograd wrapper around the fused SUPA RMSNorm forward operator."""
+
+    @staticmethod
+    def forward(ctx, hidden_states: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+        weight = weight.to(dtype=hidden_states.dtype)
+        bias = torch.zeros_like(weight)
+        output = torch.ops.sudnn.rms_norm_func(hidden_states, weight, bias, eps, False, False)
+        ctx.save_for_backward(hidden_states, weight)
+        ctx.eps = eps
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        hidden_states, weight = ctx.saved_tensors
+        eps = ctx.eps
+        hidden_size = hidden_states.shape[-1]
+
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        rrms = torch.rsqrt(variance + eps)
+        x_norm = hidden_states * rrms
+
+        grad_input = grad_weight = None
+        if ctx.needs_input_grad[1]:
+            grad_weight = (grad_output * x_norm).reshape(-1, hidden_size).sum(0).to(weight.dtype)
+
+        if ctx.needs_input_grad[0]:
+            grad_norm = grad_output * weight
+            dot = (grad_norm * hidden_states).sum(-1, keepdim=True)
+            grad_input = rrms * grad_norm - (rrms.pow(3) / hidden_size) * hidden_states * dot
+            grad_input = grad_input.to(hidden_states.dtype)
+
+        return grad_input, grad_weight, None
+
+
+def supa_rms_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    """Run the fused kernel for standard RMSNorm modules."""
+    if hidden_states.device.type != DeviceType.SUPA:
+        return self._supa_rmsnorm_original_forward(hidden_states)
+
+    return SupaRMSNormFunction.apply(hidden_states, self.weight, self.variance_epsilon)
+
+
+@KernelPlugin("supa_fused_rmsnorm").register()
+class SupaRMSNormKernel(BaseKernel):
+    """SUPA fused RMSNorm kernel for trainable models."""
+
+    @staticmethod
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current != DeviceType.SUPA:
+            raise RuntimeError(f"SupaRMSNormKernel requires SUPA, current accelerator is {current}.")
+
+    @staticmethod
+    def _apply(**kwargs) -> HFModel:
+        model = kwargs["model"]
+
+        if not _is_sudnn_rms_norm_available():
+            logger.warning_rank0(
+                "Fused SUPA RMSNorm (torch.ops.sudnn.rms_norm_func) is unavailable; "
+                "the eager RMSNorm path will be used. Ensure torch_supa_ext is installed."
+            )
+            return model
+
+        patched_count = 0
+        patched_classes = set()
+        for module in model.modules():
+            cls = module.__class__
+            if getattr(module, "_supa_rmsnorm_patched", False):
+                continue
+            if cls.__name__ not in _SUPPORTED_RMSNORM_CLASSES:
+                continue
+
+            module._supa_rmsnorm_original_forward = module.forward
+            module.forward = types.MethodType(supa_rms_norm_forward, module)
+            module._supa_rmsnorm_patched = True
+            patched_count += 1
+            patched_classes.add(cls)
+
+        if patched_count:
+            logger.info_rank0(
+                "Applied fused SUPA RMSNorm to {} modules: {}.".format(
+                    patched_count, ", ".join(sorted(cls.__name__ for cls in patched_classes))
+                )
+            )
+
+        return model