[KMCompiler]Fused moe: add thead (PPU-ZW810E) backend

src/flaggems_vllm/ops/__init__.py

@@ -97,6 +97,7 @@
     SUPPORTED_FP8_DTYPE,
     per_token_group_quant_fp8,
 )
+from flaggems_vllm.ops.permute_copy import permute_copy
 from flaggems_vllm.ops.persistent_topk import persistent_topk
 from flaggems_vllm.ops.reglu import dreglu, reglu
 from flaggems_vllm.ops.reshape_and_cache import reshape_and_cache
@@ -195,6 +196,7 @@
     "parallel_nsa_compression",
     "pack_seq_triton",
     "per_token_group_quant_fp8",
+    "permute_copy",
     "persistent_topk",
     "reglu",
     "reshape_and_cache",

src/flaggems_vllm/ops/permute_copy.py

@@ -0,0 +1,239 @@
+# Copyright 2026, The FlagOS Contributors.
+#
+# 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.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+from flaggems_vllm.runtime import torch_device_fn
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _permute_copy_kernel_1d(
+    src_ptr,
+    dst_ptr,
+    n_elements,
+    src_stride0,
+    dst_stride0,
+    perm0,
+    rank,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid.to(tl.int64) * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # Single dimension - no permutation possible
+    o0 = offsets
+
+    # Compute source offset
+    # input_idx[d] = output_idx[perm_inv[d]]
+    # For 1D: perm_inv[0] = 0 if perm[0] == 0
+    perm_inv0 = perm0
+    src_offset = tl.where(perm_inv0 == 0, o0, 0) * src_stride0
+
+    # Compute destination offset
+    dst_offset = o0 * dst_stride0
+
+    vals = tl.load(src_ptr + src_offset, mask=mask)
+    tl.store(dst_ptr + dst_offset, vals, mask=mask)
+
+
+@triton.jit
+def _permute_copy_kernel_2d(
+    src_ptr,
+    dst_ptr,
+    n_elements,
+    dst_shape1,
+    src_stride0,
+    src_stride1,
+    dst_stride0,
+    dst_stride1,
+    perm0,
+    perm1,
+    rank,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid.to(tl.int64) * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # Compute output multi-dimensional indices
+    o1 = offsets % dst_shape1
+    o0 = offsets // dst_shape1
+
+    # Compute perm_inv
+    # perm_inv[d] = j such that perm[j] = d
+    perm_inv0 = tl.where(perm1 == 0, 1, 0)
+    perm_inv1 = tl.where(perm0 == 1, 0, 1)
+
+    # input_idx[d] = output_idx[perm_inv[d]]
+    # src_offset = sum(input_idx[d] * src_stride[d])
+    #             = output_idx[perm_inv[0]] * src_stride0 + output_idx[perm_inv[1]] * src_stride1
+    src_offset = (
+        tl.where(perm_inv0 == 0, o0, o1) * src_stride0
+        + tl.where(perm_inv1 == 0, o0, o1) * src_stride1
+    )
+
+    # Compute destination offset
+    dst_offset = o0 * dst_stride0 + o1 * dst_stride1
+
+    vals = tl.load(src_ptr + src_offset, mask=mask)
+    tl.store(dst_ptr + dst_offset, vals, mask=mask)
+
+
+@triton.jit
+def _permute_copy_kernel_3d(
+    src_ptr,
+    dst_ptr,
+    n_elements,
+    dst_shape2,
+    dst_shape1,
+    src_stride0,
+    src_stride1,
+    src_stride2,
+    dst_stride0,
+    dst_stride1,
+    dst_stride2,
+    perm0,
+    perm1,
+    perm2,
+    rank,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid.to(tl.int64) * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # Compute output multi-dimensional indices
+    temp = offsets
+    o2 = temp % dst_shape2
+    temp = temp // dst_shape2
+    o1 = temp % dst_shape1
+    o0 = temp // dst_shape1
+
+    # Compute perm_inv
+    # perm_inv[d] = j such that perm[j] = d
+    perm_inv0 = tl.where(perm1 == 0, 1, tl.where(perm2 == 0, 2, 0))
+    perm_inv1 = tl.where(perm0 == 1, 0, tl.where(perm2 == 1, 2, 1))
+    perm_inv2 = tl.where(perm0 == 2, 0, tl.where(perm1 == 2, 1, 2))
+
+    # input_idx[d] = output_idx[perm_inv[d]]
+    # src_offset = sum(input_idx[d] * src_stride[d])
+    src_offset = (
+        tl.where(perm_inv0 == 0, o0, tl.where(perm_inv0 == 1, o1, o2)) * src_stride0
+        + tl.where(perm_inv1 == 0, o0, tl.where(perm_inv1 == 1, o1, o2)) * src_stride1
+        + tl.where(perm_inv2 == 0, o0, tl.where(perm_inv2 == 1, o1, o2)) * src_stride2
+    )
+
+    # Compute destination offset
+    dst_offset = o0 * dst_stride0 + o1 * dst_stride1 + o2 * dst_stride2
+
+    vals = tl.load(src_ptr + src_offset, mask=mask)
+    tl.store(dst_ptr + dst_offset, vals, mask=mask)
+
+
+def permute_copy(x: torch.Tensor, dims):
+    """
+    Wrapper for aten::permute_copy
+    Creates and returns a copy of x with permuted dimensions.
+    """
+    logger.debug("GEMS PERMUTE_COPY")
+    ndim = x.ndim
+    if ndim == 0:
+        # Scalar tensor - just return a copy
+        return x.clone()
+
+    # Normalize dims (handle negative indices)
+    dims = [d if d >= 0 else d + ndim for d in dims]
+
+    # Compute output shape
+    out_shape = [x.shape[d] for d in dims]
+
+    # Check for empty tensor
+    if x.numel() == 0:
+        return torch.empty(out_shape, dtype=x.dtype, device=x.device)
+
+    out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
+
+    # Ensure both tensors are contiguous for efficient memory access
+    src = x.contiguous() if not x.is_contiguous() else x
+    if not out.is_contiguous():
+        out = out.contiguous()
+
+    n_elements = out.numel()
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    with torch_device_fn.device(x.device):
+        if ndim == 1:
+            _permute_copy_kernel_1d[grid](
+                src,
+                out,
+                n_elements,
+                src.stride(0),
+                out.stride(0),
+                dims[0],
+                ndim,
+                # 1024 balances occupancy and register pressure for element-wise copy
+                BLOCK_SIZE=1024,
+            )
+        elif ndim == 2:
+            _permute_copy_kernel_2d[grid](
+                src,
+                out,
+                n_elements,
+                out.shape[1],
+                src.stride(0),
+                src.stride(1),
+                out.stride(0),
+                out.stride(1),
+                dims[0],
+                dims[1],
+                ndim,
+                # 1024 balances occupancy and register pressure for element-wise copy
+                BLOCK_SIZE=1024,
+            )
+        elif ndim == 3:
+            _permute_copy_kernel_3d[grid](
+                src,
+                out,
+                n_elements,
+                out.shape[2],
+                out.shape[1],
+                src.stride(0),
+                src.stride(1),
+                src.stride(2),
+                out.stride(0),
+                out.stride(1),
+                out.stride(2),
+                dims[0],
+                dims[1],
+                dims[2],
+                ndim,
+                # 1024 balances occupancy and register pressure for element-wise copy
+                BLOCK_SIZE=1024,
+            )
+        else:
+            # For ranks > 3, use permute (view) + contiguous (copy)
+            # to avoid recursive dispatch through torch.permute_copy
+            return x.permute(dims).contiguous()
+    return out

src/flaggems_vllm/runtime/backend/_thead/fused/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2026 FlagOS Contributors
+#
+# 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.
+
+from flaggems_vllm.runtime.backend._thead.fused.fused_moe import (
+    fused_experts_impl,
+    inplace_fused_experts,
+    outplace_fused_experts,
+)
+
+__all__ = [
+    "fused_experts_impl",
+    "inplace_fused_experts",
+    "outplace_fused_experts",
+]