feat: support flydsl all2all

aiter/dist/device_communicators/all2all.py

@@ -11,7 +11,10 @@ def _has_module(module_name: str) -> bool:
     The result is cached so that subsequent queries for the same module incur
     no additional overhead.
     """
-    return importlib.util.find_spec(module_name) is not None
+    try:
+        return importlib.util.find_spec(module_name) is not None
+    except ModuleNotFoundError:
+        return False
 
 
 def has_mori() -> bool:
@@ -20,18 +23,23 @@ def has_mori() -> bool:
 
 
 class MoriAll2AllManager(All2AllManagerBase):
+    @staticmethod
+    def _init_mori_shmem(cpu_group) -> None:
+        """Register *cpu_group* with mori's shmem heap and run the barrier."""
+        import mori
+
+        torch._C._distributed_c10d._register_process_group("mori", cpu_group)
+        mori.shmem.shmem_torch_process_group_init("mori")
+
     def __init__(self, cpu_group):
         assert has_mori(), (
             "MoRI kernels not found. Please follow https://github.com/ROCm/mori/blob/main/README.md"
             " to install MoRI kernels."
         )  # noqa
-        import mori
 
         super().__init__(cpu_group)
         self.handle_cache = Cache()
-
-        torch._C._distributed_c10d._register_process_group("mori", cpu_group)
-        mori.shmem.shmem_torch_process_group_init("mori")
+        self._init_mori_shmem(cpu_group)
 
     def _make_all2all_kwargs(
         self,
@@ -96,3 +104,117 @@ def get_handle(self, kwargs):
             mori_kwargs, self._make_handle
         )
         return handle
+
+
+class FlyDSLAll2AllManager(All2AllManagerBase):
+    """
+    EP all2all backend backed by FlyDSL intranode dispatch/combine kernels.
+
+    FlyDSL still uses mori's shmem heap for P2P buffer allocation, so mori
+    must be installed alongside flydsl. The dispatch/combine *kernels* however
+    are entirely FlyDSL-generated, replacing mori's comm primitives.
+
+    TBO multi-instance ops are created via ``create_handle`` (non-cached) so
+    the two ubatch ops are guaranteed to be distinct, independent objects.
+    """
+
+    @staticmethod
+    def _init_mori_shmem(cpu_group) -> None:
+        """Register *cpu_group* with mori's shmem heap and run the barrier."""
+        import mori
+
+        torch._C._distributed_c10d._register_process_group("mori", cpu_group)
+        mori.shmem.shmem_torch_process_group_init("mori")
+
+    @staticmethod
+    def _make_quant_type(input_dtype: torch.dtype, quant_dtype: torch.dtype) -> str:
+        fp8_dtypes = {
+            torch.float8_e4m3fn,
+            torch.float8_e4m3fnuz,
+        }
+        if input_dtype == torch.bfloat16 and quant_dtype in fp8_dtypes:
+            return "fp8_direct_cast"
+        return "none"
+
+    def __init__(self, cpu_group):
+        try:
+            from aiter.ops.flydsl.kernels.flydsl_dispatch_combine_intranode_op import (
+                FlyDSLDispatchCombineConfig,
+                FlyDSLDispatchCombineIntraNodeOp,
+            )
+        except ImportError as e:
+            raise ImportError(
+                "FlyDSL dispatch/combine module not found at "
+                "'aiter.ops.flydsl.kernels.flydsl_dispatch_combine_intranode_op'."
+            ) from e
+
+        assert has_mori(), (
+            "mori is required alongside FlyDSL for shmem buffer allocation. "
+            "Please install mori."
+        )
+        super().__init__(cpu_group)
+        self._flydsl_dispatch_config_cls = FlyDSLDispatchCombineConfig
+        self._flydsl_dispatch_op_cls = FlyDSLDispatchCombineIntraNodeOp
+        if self.internode:
+            raise NotImplementedError(
+                "FlyDSLAll2AllManager currently supports only intranode EP "
+                "dispatch/combine. For inter-node EP, please use the mori "
+                "backend."
+            )
+
+        # FlyDSL uses mori.shmem for P2P buffer allocation internally.
+        # Keep shmem init behavior aligned with MoriAll2AllManager.
+        self._init_mori_shmem(cpu_group)
+        self.handle_cache = Cache()
+
+    def _make_all2all_kwargs(
+        self,
+        rank: int,
+        num_ep_ranks: int,
+        input_dtype: torch.dtype,
+        quant_dtype: torch.dtype,
+        token_hidden_size: int,
+        scale_dim: int,
+        scale_type_size: int,
+        max_num_tokens_per_dp_rank: int,
+        num_local_experts: int,
+        num_experts_per_token: int,
+    ):
+
+        return dict(
+            rank=rank,
+            world_size=num_ep_ranks,
+            # Buffer sized for input dtype (bf16) so one op handles both
+            # bf16 and quantized input without reallocation; the dispatch
+            # kernel specialisation is selected at call time by input.dtype.
+            data_type=input_dtype,
+            hidden_dim=token_hidden_size,
+            scale_dim=scale_dim,
+            scale_type_size=scale_type_size,
+            max_token_type_size=input_dtype.itemsize,
+            max_num_inp_token_per_rank=max_num_tokens_per_dp_rank,
+            num_experts_per_rank=num_local_experts,
+            num_experts_per_token=num_experts_per_token,
+            quant_type=self._make_quant_type(input_dtype, quant_dtype),
+        )
+
+    def _make_handle(self, **kwargs):
+        cfg = self._flydsl_dispatch_config_cls(**kwargs)
+        return self._flydsl_dispatch_op_cls(cfg)
+
+    def get_handle(self, kwargs):
+        flydsl_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("FlyDSL all2all args %s", flydsl_kwargs)
+        return self.handle_cache.get_or_create(flydsl_kwargs, self._make_handle)
+
+    def create_handle(self, kwargs):
+        """Create a fresh, uncached FlyDSL op instance.
+
+        Unlike ``get_handle`` (which caches one op per config), every call
+        returns a new independent op. Callers that need multiple distinct ops
+        for the same config (e.g. ATOM for TBO ubatches) should call this
+        and manage the instances themselves.
+        """
+        flydsl_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("FlyDSL all2all (uncached) args %s", flydsl_kwargs)
+        return self._make_handle(**flydsl_kwargs)

aiter/dist/device_communicators/communicator_cuda.py

@@ -129,6 +129,10 @@ def all2all_manager(self):
                 from .all2all import MoriAll2AllManager
 
                 self._all2all_manager = MoriAll2AllManager(self.cpu_group)
+            elif self.all2all_backend == "flydsl":
+                from .all2all import FlyDSLAll2AllManager
+
+                self._all2all_manager = FlyDSLAll2AllManager(self.cpu_group)
             elif self.all2all_backend == "flashinfer_all2allv":
                 from .all2all import FlashInferAllToAllManager
 

aiter/ops/flydsl/kernels/communication_ops_utils.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""Low-level cross-card (P2P) communication primitives for communication kernels.
+
+These wrap LLVM-dialect global memory ops with explicit memory ordering and
+syncscope -- which the high-level FlyDSL APIs (buffer_ops / Pointer) do not
+expose -- so dispatch/combine can publish and observe data across cards.
+
+Also hosts :class:`GeometryTuningTable`, the per-shape launch-geometry lookup
+shared by the dispatch/combine ops.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Dict, Tuple
+
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import llvm as _llvm_d
+from flydsl.expr import arith
+
+__all__ = [
+    "store_i32_system",
+    "store_i64_global_system",
+    "fence_system_acquire",
+    "load_i64_global",
+    "atomic_add_global_at",
+    "GeometryTuningTable",
+]
+
+
+def _to_ptr_global(v):
+    """Cast an i64 address to ``!llvm.ptr<1>`` (global address space)."""
+    return _llvm_d.IntToPtrOp(
+        _llvm_d.PointerType.get(address_space=1), arith.unwrap(v)
+    ).result
+
+
+def store_i32_system(addr_i64, offset, val):
+    """System-scope release i32 store at ``addr_i64 + offset*4``."""
+    base = arith.unwrap(addr_i64)
+    off = arith.unwrap(offset)
+    val_ = arith.unwrap(val)
+    _i64 = ir.IntegerType.get_signless(64)
+    _i32 = ir.IntegerType.get_signless(32)
+    _nuw = ir.Attribute.parse("#llvm.overflow<none>")
+    off64 = _llvm_d.ZExtOp(_i64, off).res if off.type == _i32 else off
+    byte_off = _llvm_d.MulOp(
+        off64, _llvm_d.ConstantOp(_i64, ir.IntegerAttr.get(_i64, 4)).result, _nuw
+    ).result
+    addr = _llvm_d.AddOp(base, byte_off, _nuw).result
+    gptr = _llvm_d.IntToPtrOp(_llvm_d.PointerType.get(address_space=1), addr).result
+    _llvm_d.StoreOp(
+        val_,
+        gptr,
+        alignment=4,
+        ordering=_llvm_d.AtomicOrdering.release,
+        syncscope="one-as",
+    )
+
+
+def store_i64_global_system(addr_i64, val):
+    """System-scope release i64 store to ``addr_i64``."""
+    gptr = _to_ptr_global(addr_i64)
+    _llvm_d.StoreOp(
+        arith.unwrap(val),
+        gptr,
+        alignment=8,
+        ordering=_llvm_d.AtomicOrdering.release,
+        syncscope="one-as",
+    )
+
+
+def fence_system_acquire():
+    """System-scope acquire fence."""
+    _llvm_d.FenceOp(_llvm_d.AtomicOrdering.acquire, syncscope="one-as")
+
+
+def load_i64_global(addr_i64):
+    """Relaxed global i64 load from ``addr_i64``."""
+    ptr = _to_ptr_global(addr_i64)
+    _i64 = ir.IntegerType.get_signless(64)
+    return _llvm_d.LoadOp(_i64, ptr, alignment=8).result
+
+
+def atomic_add_global_at(addr_i64, val):
+    """Monotonic global ``atomic fetch-and-add``; returns the old value."""
+    ptr = _to_ptr_global(addr_i64)
+    return _llvm_d.AtomicRMWOp(
+        _llvm_d.AtomicBinOp.add,
+        ptr,
+        arith.unwrap(val),
+        _llvm_d.AtomicOrdering.monotonic,
+    ).res
+
+
+@dataclass
+class GeometryTuningTable:
+    """Per-shape token-count -> (block_num, warp_num_per_block) lookup; rounds up
+    to the smallest bucket >= count (largest on overflow, mori parity)."""
+
+    dispatch: Dict[int, Tuple[int, int]] = field(default_factory=dict)
+    combine: Dict[int, Tuple[int, int]] = field(default_factory=dict)
+
+    def __post_init__(self):
+        for phase, tbl in (("dispatch", self.dispatch), ("combine", self.combine)):
+            for n_tok, (bn, wpb) in tbl.items():
+                if bn <= 0 or wpb <= 0:
+                    raise ValueError(
+                        f"GeometryTuningTable.{phase}[{n_tok}] must be positive, "
+                        f"got block_num={bn}, warp_num_per_block={wpb}"
+                    )
+
+    @classmethod
+    def from_tuning_file(
+        cls,
+        path,
+        *,
+        dtype,
+        hidden_dim,
+        zero_copy,
+        topk=None,
+        local_expert_num=None,
+        combine_dtype="bf16",
+    ):
+        """Build a per-op table from a multi-shape tuning JSON, filtered to this
+        op's shape; empty table => cfg defaults."""
+        with open(path, "r", encoding="utf-8") as f:
+            raw = json.load(f)
+
+        def _match(r, want_dtype, need_zc):
+            if (
+                r.get("dtype") != want_dtype
+                or int(r.get("hidden_dim", -1)) != hidden_dim
+            ):
+                return False
+            if topk is not None and "topk" in r and int(r["topk"]) != topk:
+                return False
+            if (
+                local_expert_num is not None
+                and "local_expert_num" in r
+                and int(r["local_expert_num"]) != local_expert_num
+            ):
+                return False
+            if need_zc and bool(r.get("zero_copy", False)) != bool(zero_copy):
+                return False
+            return True
+
+        def _build(rules, want_dtype, need_zc):
+            return {
+                int(r["num_tokens"]): (
+                    int(r["block_num"]),
+                    int(r["warp_num_per_block"]),
+                )
+                for r in rules
+                if _match(r, want_dtype, need_zc)
+            }
+
+        return cls(
+            dispatch=_build(raw.get("dispatch", []), dtype, need_zc=False),
+            combine=_build(raw.get("combine", []), combine_dtype, need_zc=True),
+        )
+
+    def lookup(self, phase, num_tokens):
+        """Smallest bucket >= num_tokens (largest on overflow); None if empty."""
+        tbl = self.dispatch if phase == "dispatch" else self.combine
+        if not tbl:
+            return None
+        if num_tokens in tbl:
+            return tbl[num_tokens]
+        candidates = [k for k in tbl if k >= num_tokens]
+        return tbl[min(candidates)] if candidates else tbl[max(tbl)]