Add opt-in TorchAO HiFloat8 Linear and AutoEP training

accelerator/abstract_accelerator.py

@@ -215,6 +215,14 @@ def is_triton_supported(self):
     def prefer_triton_grouped_mm(self):
         return False
 
+    def get_hifloat8_provider(self):
+        """Return optional public HiFloat8 APIs for this accelerator."""
+        return None
+
+    def get_grouped_mm(self):
+        """Return a native differentiable grouped GEMM override, if available."""
+        return None
+
     # Graph operations
     @abc.abstractmethod
     def create_graph(self):

accelerator/npu/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Optional Ascend capability discovery."""

accelerator/npu/hifloat8.py

@@ -0,0 +1,29 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Resolve the optional public TorchAO HiFloat8 APIs on NPU."""
+
+import importlib
+from importlib.metadata import PackageNotFoundError, version
+
+from packaging.version import Version
+
+
+def get_hifloat8_provider():
+    hint = ("Install torchao>=0.18,<0.19 and a compatible torchao-npu build "
+            "providing torchao.hifloat8; keep automatic activation enabled.")
+    try:
+        installed_version = Version(version("torchao"))
+    except PackageNotFoundError as error:
+        raise ImportError(f"TorchAO HiFloat8 is unavailable. {hint}") from error
+    if not Version("0.18") <= installed_version < Version("0.19"):
+        raise RuntimeError(f"Unsupported torchao version {installed_version}. {hint}")
+    try:
+        api = importlib.import_module("torchao.hifloat8")
+    except ImportError as error:
+        raise ImportError(f"TorchAO HiFloat8 API is unavailable. {hint}") from error
+    required = ("HiFloat8Config", "assert_hifloat8_training_available")
+    missing = [name for name in required if not hasattr(api, name)]
+    if missing:
+        raise ImportError(f"TorchAO HiFloat8 API is missing {missing}. {hint}")
+    return api

accelerator/npu_accelerator.py

@@ -60,6 +60,16 @@ def is_synchronized_device(self):
     def use_host_timers(self):
         return self.is_synchronized_device()
 
+    def get_hifloat8_provider(self):
+        from .npu.hifloat8 import get_hifloat8_provider
+
+        return get_hifloat8_provider()
+
+    def get_grouped_mm(self):
+        from deepspeed.ops.npu.grouped_mm import grouped_mm
+
+        return grouped_mm
+
     def resolves_data_dependency(self):
         return self.is_synchronized_device()
 

deepspeed/module_inject/auto_ep.py

@@ -126,11 +126,13 @@ def _has_3d_expert_params(module: nn.Module, preset: MoEModelPreset) -> bool:
 
 def _get_num_experts_from_config(model_config, preset: MoEModelPreset) -> int | None:
     """Extract num_experts from model.config using the preset's attribute name."""
+    model_config = getattr(model_config, "text_config", model_config)
     return getattr(model_config, preset.num_experts_attr, None)
 
 
 def _get_top_k_from_config(model_config, preset: MoEModelPreset) -> int | None:
     """Extract top_k from model.config using the preset's attribute name."""
+    model_config = getattr(model_config, "text_config", model_config)
     return getattr(model_config, preset.top_k_attr, None)
 
 
@@ -611,6 +613,8 @@ def ep_parser(self) -> list[MoELayerSpec]:
                     router_logits_capture_mode=forward_contract.router_logits_capture_mode,
                     moe_output_shape=forward_contract.moe_output_shape,
                     e_score_correction_bias_path=e_score_correction_bias_path,
+                    router_topk_sorted=preset.router_topk_sorted,
+                    router_scores_in_gate_dtype=preset.router_scores_in_gate_dtype,
                     expert_activation=expert_activation,
                     expert_activation_alpha=activation_alpha,
                     expert_activation_limit=activation_limit,

deepspeed/module_inject/auto_ep_layer.py

@@ -382,7 +382,7 @@ def permute_by_local_expert(
         n_tokens: original token count before padding (for unpermute)
     """
     from deepspeed.moe.ep_kernels import (generate_permute_indices, permute_rows, permute_rows_supported,
-                                          TOKEN_GROUP_ALIGN_SIZE_M)
+                                          prefer_cpu_permutation_indices, TOKEN_GROUP_ALIGN_SIZE_M)
 
     if local_counts.ndim == 1:
         # [E_local]: already aggregated over sources (ep_degree=1)
@@ -404,9 +404,9 @@ def permute_by_local_expert(
     x_padded_per_expert = n_tokens + num_local_experts * alignment
     padded_max_len = ((x_padded_per_expert + alignment - 1) // alignment) * alignment
 
-    # Use the pure-PyTorch path for host tensors. The CPU accelerator reports
-    # CPU tensors as "on accelerator", but Triton still requires a GPU driver.
-    use_cpu = tokens.device.type == "cpu"
+    # Small NPU groups retain the lower-latency host reference path.
+    use_cpu = (local_counts_flat.device.type == "cpu"
+               or prefer_cpu_permutation_indices(tokens.device.type, local_counts_flat.numel()))
     counts_for_permute = local_counts_flat.cpu() if use_cpu else local_counts_flat
     with torch.no_grad():
         permuted_indices, m_sizes, _offsets = generate_permute_indices(
@@ -417,7 +417,7 @@ def permute_by_local_expert(
             alignment,
             use_cpu=use_cpu,
         )
-    if not use_cpu:
+    if tokens.device.type != "cpu":
         permuted_indices = permuted_indices.to(tokens.device)
         m_sizes = m_sizes.to(tokens.device)
 
@@ -594,6 +594,8 @@ def __init__(
                 route_scale=spec.route_scale,
                 gate_bias=spec.gate_bias,
                 group_score_func=spec.group_score_func,
+                topk_sorted=spec.router_topk_sorted,
+                scores_in_gate_dtype=spec.router_scores_in_gate_dtype,
             )
             # Copy gate weights
             _copy_parameter_data(self.router.gate.weight, source_gate.weight)

deepspeed/module_inject/auto_ep_presets/base.py

@@ -58,6 +58,8 @@ class MoEModelPreset:
     min_transformers_version: str | None = None
     validated_transformers_versions: str = ""
     docs_support_notes: str = ""
+    router_topk_sorted: bool = False
+    router_scores_in_gate_dtype: bool = False
 
 
 @dataclass
@@ -97,6 +99,8 @@ class MoELayerSpec:
     router_logits_capture_mode: Literal["raw", "post_score"] = "post_score"
     moe_output_shape: Literal["batched", "flat"] = "batched"
     e_score_correction_bias_path: str | None = None
+    router_topk_sorted: bool = False
+    router_scores_in_gate_dtype: bool = False
     #: Expert MLP activation resolved for this layer (see MoEModelPreset.expert_activation).
     expert_activation: str = "swiglu"
     expert_activation_alpha: float = 1.702

deepspeed/module_inject/auto_ep_presets/qwen3_5_moe.py

@@ -20,7 +20,7 @@
 PRESET_NAME = "qwen3_5_moe"
 
 PRESET = MoEModelPreset(
-    moe_layer_pattern=r"model\.layers\.\d+\.mlp",
+    moe_layer_pattern=r"model\.(?:language_model\.)?layers\.\d+\.mlp",
     router_pattern="gate",
     experts_pattern="experts",
     expert_storage="fused_3d",
@@ -37,13 +37,11 @@
     shared_experts_pattern="shared_expert",
     shared_experts_gate_pattern="shared_expert_gate",
     preset_adapter="qwen3_5_moe",
-    hf_model_types=("qwen3_5_moe_text", ),
-    unsupported_hf_model_type_notes={
-        "qwen3_5_moe": ("AutoEP supports the Qwen3.5 text backbone preset path; pass the "
-                        "text-backbone model/config with model_type='qwen3_5_moe_text'.")
-    },
+    hf_model_types=("qwen3_5_moe_text", "qwen3_5_moe"),
     min_transformers_version="5.2.0",
-    docs_support_notes="Requires the Qwen3.5 text-backbone qwen3_5_moe_text model type.",
+    docs_support_notes="Text backbone and the language-model path of the Qwen3.5 MoE wrapper.",
+    router_topk_sorted=True,
+    router_scores_in_gate_dtype=True,
 )
 
 
@@ -61,6 +59,7 @@ def adjust_forward_contract(self, contract: ForwardContract) -> ForwardContract:
             return_router_logits=True,
             capture_target="router",
             capture_index=1,
+            router_logits_capture_mode="raw",
         )
 
     def retarget_transformers_output_recorders(

deepspeed/moe/ep_experts.py

@@ -300,6 +300,35 @@ def _run_experts_triton_grouped_mm(
 # ---------------------------------------------------------------------------
 
 
+def _run_experts_with_grouped_mm(w1,
+                                 w2,
+                                 w3,
+                                 x,
+                                 num_tokens_per_expert,
+                                 activation="swiglu",
+                                 alpha=1.702,
+                                 limit=7.0,
+                                 *,
+                                 grouped_mm):
+    """Keep routing, padding and activation independent of the GEMM backend."""
+    offsets = num_tokens_per_expert.cumsum(0).to(torch.int64)
+    rows = int(offsets[-1].item())
+    if rows < 0 or rows > x.shape[0]:
+        raise ValueError("Expert token counts exceed the available rows")
+    inputs = x[:rows]
+
+    def mm(lhs, weight):
+        return grouped_mm(lhs, weight.to(x.dtype), offsets, trans_b=True).to(lhs.dtype)
+
+    gate = mm(inputs, w1)
+    up = mm(inputs, w3)
+    hidden = apply_expert_activation(gate, up, activation, alpha, limit, fused=False)
+    output = mm(hidden, w2)
+    if rows == x.shape[0]:
+        return output
+    return torch.cat((output, output.new_zeros((x.shape[0] - rows, output.shape[-1]))))
+
+
 class GroupedExperts(nn.Module):
     """Grouped expert computation for MoE layers.
 
@@ -356,6 +385,7 @@ def __init__(
         self.w3.is_expert_group = True
         self.use_triton_grouped_mm = False
         self.use_grouped_mm = use_grouped_mm
+        self.grouped_mm = None
 
         # Resolve the Triton path. The device-specific decision is delegated to
         # the accelerator backend (e.g. the CUDA backend prefers Triton on
@@ -364,7 +394,8 @@ def __init__(
         if use_grouped_mm and not disable_triton_grouped_mm:
             self.use_triton_grouped_mm = get_accelerator().prefer_triton_grouped_mm()
 
-        if use_grouped_mm and not hasattr(torch, "_grouped_mm") and not self.use_triton_grouped_mm:
+        if (use_grouped_mm and not hasattr(torch, "_grouped_mm") and not self.use_triton_grouped_mm
+                and get_accelerator().get_grouped_mm() is None):
             raise RuntimeError("GroupedExperts was constructed with use_grouped_mm=True but "
                                "torch._grouped_mm is not available in this PyTorch build. "
                                "Upgrade PyTorch to a build that provides torch._grouped_mm, install "
@@ -392,6 +423,20 @@ def forward(
         """
 
         act = (self.activation, self.activation_alpha, self.activation_limit)
+        grouped_mm = self.grouped_mm
+        accelerator = get_accelerator()
+        if grouped_mm is None and self.use_grouped_mm and accelerator.on_accelerator(x):
+            grouped_mm = accelerator.get_grouped_mm()
+        if grouped_mm is not None:
+            if getattr(self, "gate_up_impl", "auto") == "fused":
+                raise RuntimeError("The accelerator grouped GEMM override requires separate gate/up projections")
+            return _run_experts_with_grouped_mm(self.w1,
+                                                self.w2,
+                                                self.w3,
+                                                x,
+                                                num_tokens_per_expert,
+                                                *act,
+                                                grouped_mm=grouped_mm)
         if self.use_triton_grouped_mm:
             return _run_experts_triton_grouped_mm(self.w1, self.w2, self.w3, x, num_tokens_per_expert, *act)
         elif self.use_grouped_mm:

deepspeed/moe/ep_kernels.py

@@ -37,8 +37,7 @@
 
     _TRITON_AVAILABLE = True
 except ImportError:
-    logger.info("Triton not available; using pure-PyTorch CPU fallback for "
-                "permutation index generation.")
+    logger.info("Triton not available; using the available PyTorch permutation-index backend.")
 
 # ---------------------------------------------------------------------------
 # Alignment constant
@@ -52,6 +51,14 @@
  - mxfp8: 32 (scaling block size)
 """
 
+NPU_DEVICE_PERMUTE_MIN_GROUPS = 64
+
+
+def prefer_cpu_permutation_indices(device_type: str, group_count: int) -> bool:
+    """Use host indexing when device launch overhead dominates small NPU groups."""
+    return device_type == "cpu" or (device_type == "npu" and group_count < NPU_DEVICE_PERMUTE_MIN_GROUPS)
+
+
 # ---------------------------------------------------------------------------
 # Utility: round up
 # ---------------------------------------------------------------------------
@@ -115,10 +122,16 @@ def fill_indices_wrapper(
     block_size: int = 128,
     max_blocks: int = 1024,
 ) -> torch.Tensor:
-    """Launch the Triton kernel to fill permutation indices.
-
-    Falls back to :func:`fill_indices_cpu` when Triton is unavailable.
-    """
+    """Fill permutation indices with the available device implementation."""
+    if tokens_per_expert_group.device.type == "npu":
+        return fill_indices_torch(
+            tokens_per_expert_group,
+            start_index_values,
+            write_offsets,
+            experts_per_rank,
+            num_ranks,
+            max_len,
+        )
     if not _TRITON_AVAILABLE:
         return fill_indices_cpu(
             tokens_per_expert_group,
@@ -182,6 +195,31 @@ def fill_indices_cpu(
     return permuted_indices
 
 
+def fill_indices_torch(
+    tokens_per_expert_group: torch.Tensor,
+    start_index_values: torch.Tensor,
+    write_offsets: torch.Tensor,
+    experts_per_rank: int,
+    num_ranks: int,
+    max_len: int,
+) -> torch.Tensor:
+    """Build the padded permutation without host counts; AutoEP supplies ``max_len >= token count``."""
+    counts = tokens_per_expert_group.to(torch.int64)
+    source_ends = start_index_values + counts
+    source_prefix = (counts.view(num_ranks, experts_per_rank).cumsum(0) -
+                     counts.view(num_ranks, experts_per_rank)).reshape(-1)
+    source_rows = torch.arange(max_len, device=counts.device, dtype=torch.int64)
+    valid = source_rows < source_ends[-1]
+    group_ids = torch.searchsorted(source_ends, source_rows, right=True).clamp_max(counts.numel() - 1)
+    expert_ids = group_ids.remainder(experts_per_rank)
+    destinations = (write_offsets[expert_ids] + source_prefix[group_ids] + source_rows - start_index_values[group_ids])
+    # Extra slot absorbs all rows beyond the real token count; it is discarded.
+    destinations = torch.where(valid, destinations, max_len).to(torch.int64)
+    result = torch.full((max_len + 1, ), -1, dtype=torch.int32, device=counts.device)
+    result.scatter_(0, destinations, source_rows.to(torch.int32))
+    return result[:-1]
+
+
 # ===================================================================
 # generate_permute_indices
 # ===================================================================

deepspeed/moe/ep_router.py

@@ -60,6 +60,8 @@ def __init__(
         route_scale: float,
         gate_bias: bool,
         group_score_func: str = "top2_sum",
+        topk_sorted: bool = False,
+        scores_in_gate_dtype: bool = False,
     ):
         super().__init__()
         self.gate = nn.Linear(dim, num_experts, bias=gate_bias)
@@ -71,6 +73,8 @@ def __init__(
         self.route_norm = route_norm
         self.route_scale = route_scale
         self.group_score_func = group_score_func
+        self.topk_sorted = topk_sorted
+        self.scores_in_gate_dtype = scores_in_gate_dtype
         # Trainable expert score correction bias (e.g. DeepSeek-V3/Moonlight noaux_tc).
         # Separate from the dynamic load-balancing expert_bias passed in forward().
         self.e_score_correction_bias = None
@@ -152,6 +156,7 @@ def forward(
         """
         # Gate projection -> (T, num_experts)
         scores = self.gate(x)
+        gate_dtype = scores.dtype
 
         # Scoring in float32 to avoid loss explosion
         if self.score_func == "sigmoid":
@@ -172,7 +177,7 @@ def forward(
             scores_for_choice = self._get_node_limited_routing_scores(scores_for_choice)
 
         # Select top-k experts per token
-        _, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)
+        _, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=self.topk_sorted)
 
         # Gather original (unbiased) scores for selected experts
         top_scores = scores.gather(dim=1, index=selected_experts_indices)
@@ -183,6 +188,8 @@ def forward(
             top_scores = top_scores / denominator
 
         top_scores = top_scores * self.route_scale
+        if self.scores_in_gate_dtype:
+            top_scores = top_scores.to(gate_dtype)
 
         num_tokens_per_expert = count_tokens_per_expert(selected_experts_indices, self.num_experts)
 

deepspeed/ops/npu/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Native Ascend operator adapters."""