Bundle direct2560 grouped matmul for Ascend

CMakeLists.txt

@@ -48,4 +48,8 @@ else()
   message(STATUS "DLINFER_BUILD_DICP=OFF: skipping DICP vendor build (set env DLINFER_BUILD_DICP=ON to enable)")
 endif()
 
+if(DEVICE STREQUAL "ascend")
+  add_subdirectory(dlinfer/vendor/ascend/csrc/grouped_matmul_direct)
+endif()
+
 install(CODE "message(STATUS \"Install completed for device: ${DEVICE}\")")

benchmark/probe_grouped_matmul_direct.py

@@ -0,0 +1,71 @@
+"""Device correctness probe for DLInfer's bundled direct2560 extension."""
+
+import argparse
+
+from dlinfer.vendor.ascend import grouped_matmul_direct
+import torch
+import torch_npu
+
+
+def run(dtype: torch.dtype, iterations: int = 1) -> None:
+    if not grouped_matmul_direct.is_available():
+        raise RuntimeError(grouped_matmul_direct.unavailable_reason())
+
+    experts, hidden_size, output_size, tokens = 2560, 16, 16, 8
+    device = torch.device("npu:0")
+    x = torch.randn(tokens, hidden_size, device=device, dtype=dtype)
+    stored_weight = torch.randn(
+        experts, output_size, hidden_size, device=device, dtype=dtype
+    )
+    weight = stored_weight.transpose(1, 2)
+    group_list = torch.zeros(experts, device=device, dtype=torch.int64)
+    group_list[:tokens] = 1
+
+    references = []
+    row_start = 0
+    for expert_start in range(0, experts, 1024):
+        expert_end = min(expert_start + 1024, experts)
+        chunk_groups = group_list[expert_start:expert_end]
+        row_end = row_start + int(chunk_groups.sum().cpu())
+        if row_end > row_start:
+            references.append(
+                torch.ops.npu.npu_grouped_matmul(
+                    [x[row_start:row_end]],
+                    [weight[expert_start:expert_end]],
+                    group_list=chunk_groups,
+                    split_item=2,
+                    group_type=0,
+                    group_list_type=1,
+                )[0]
+            )
+        row_start = row_end
+
+    reference = torch.cat(references, dim=0)
+    for iteration in range(iterations):
+        print(f"ITERATION {iteration + 1}/{iterations} begin", flush=True)
+        direct = grouped_matmul_direct.grouped_matmul(x, stored_weight, group_list, 1)
+        print(f"ITERATION {iteration + 1}/{iterations} submitted", flush=True)
+        torch_npu.npu.synchronize()
+        print(f"ITERATION {iteration + 1}/{iterations} synchronized", flush=True)
+        torch.testing.assert_close(direct, reference, rtol=0, atol=0)
+    print(
+        f"PASS dtype={dtype} iterations={iterations} "
+        f"max_abs={(direct - reference).abs().max().item()}"
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dtype", choices=("fp16", "bf16", "both"), default="both")
+    parser.add_argument("--iterations", type=int, default=1)
+    args = parser.parse_args()
+    if args.iterations < 1:
+        parser.error("--iterations must be at least 1")
+    if args.dtype in ("fp16", "both"):
+        run(torch.float16, args.iterations)
+    if args.dtype in ("bf16", "both"):
+        run(torch.bfloat16, args.iterations)
+
+
+if __name__ == "__main__":
+    main()

dlinfer/__init__.py

@@ -1,4 +1,29 @@
 # Copyright (c) 2024, DeepLink. All rights reserved.
+import os
+from pathlib import Path
+
+
+def _register_bundled_ascend_opp() -> None:
+    """Expose wheel-local Ascend operators before torch_npu initializes."""
+    if os.environ.get("DLINFER_GMM_EXPERIMENT", "chunked") != "direct2560":
+        return
+    ascend_dir = Path(__file__).parent / "vendor/ascend"
+    candidates = (
+        ascend_dir / "grouped_matmul_direct",
+        ascend_dir / "csrc/grouped_matmul_direct/vendor",
+    )
+    bundled_opp = next((path for path in candidates if path.is_dir()), None)
+    if bundled_opp is None:
+        return
+    bundled_opp_str = str(bundled_opp)
+    configured = os.environ.get("ASCEND_CUSTOM_OPP_PATH", "")
+    paths = [path for path in configured.split(":") if path]
+    if bundled_opp_str not in paths:
+        os.environ["ASCEND_CUSTOM_OPP_PATH"] = ":".join([bundled_opp_str, *paths])
+
+
+_register_bundled_ascend_opp()
+
 import dlinfer.vendor as vendor
 
 vendor.vendor_torch_init()

dlinfer/vendor/ascend/csrc/grouped_matmul_direct/CMakeLists.txt

@@ -0,0 +1,46 @@
+include(ascend)
+
+set(GMM_DIRECT_VENDOR_DIR ${CMAKE_CURRENT_SOURCE_DIR}/vendor)
+set(GMM_DIRECT_OPAPI ${GMM_DIRECT_VENDOR_DIR}/op_api/lib/libcust_opapi.so)
+
+if(NOT EXISTS ${GMM_DIRECT_OPAPI})
+  message(FATAL_ERROR "Missing bundled direct2560 op-api: ${GMM_DIRECT_OPAPI}")
+endif()
+
+add_library(_grouped_matmul_direct MODULE grouped_matmul_direct.cpp)
+find_library(TORCH_PYTHON_LIBRARY torch_python
+  PATHS "${TORCH_INSTALL_PREFIX}/lib"
+  REQUIRED
+  NO_DEFAULT_PATH
+)
+set_target_properties(_grouped_matmul_direct PROPERTIES
+  PREFIX ""
+  BUILD_RPATH "${GMM_DIRECT_VENDOR_DIR}/lib"
+  INSTALL_RPATH "\$ORIGIN/grouped_matmul_direct/op_api/lib"
+)
+
+target_compile_definitions(_grouped_matmul_direct PRIVATE
+  TORCH_EXTENSION_NAME=_grouped_matmul_direct
+  _GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI}
+)
+target_compile_options(_grouped_matmul_direct PRIVATE -O2 -Wall -Wextra)
+target_include_directories(_grouped_matmul_direct PRIVATE
+  ${GMM_DIRECT_VENDOR_DIR}/op_api/include
+  ${TORCH_INCLUDE_DIRS}
+  ${CANN_INCLUDE_DIRS}
+  ${CANN_INCLUDE_DIRS}/aclnn
+)
+target_link_libraries(_grouped_matmul_direct PRIVATE
+  Python::Python
+  torch
+  ${TORCH_PYTHON_LIBRARY}
+  ${GMM_DIRECT_OPAPI}
+  ${CANN_LIBRARIES}
+)
+
+install(TARGETS _grouped_matmul_direct
+  LIBRARY DESTINATION dlinfer/vendor/ascend
+)
+install(DIRECTORY ${GMM_DIRECT_VENDOR_DIR}/
+  DESTINATION dlinfer/vendor/ascend/grouped_matmul_direct
+)

dlinfer/vendor/ascend/csrc/grouped_matmul_direct/grouped_matmul_direct.cpp

@@ -0,0 +1,163 @@
+// Copied from DeepLink-org/DLBlas csrc/ascend/grouped_matmul_direct and
+// adapted for DLInfer pybind packaging.
+#include <acl/acl.h>
+#include <aclnn/aclnn_base.h>
+#include <aclnnop/aclnn_grouped_matmul_v5.h>
+#include <torch/extension.h>
+
+#include <array>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+namespace {
+
+constexpr int64_t kDirectGroups = 2560;
+
+at::Tensor GroupedMatmulDirect(const at::Tensor& x, const at::Tensor& weight, const at::Tensor& group_list, int64_t group_list_type, uint64_t stream_handle);
+
+aclDataType ToAclDataType(const at::ScalarType dtype) {
+    switch (dtype) {
+        case at::ScalarType::Half:
+            return ACL_FLOAT16;
+        case at::ScalarType::BFloat16:
+            return ACL_BF16;
+        case at::ScalarType::Long:
+            return ACL_INT64;
+        default:
+            TORCH_CHECK(false, "unsupported dtype for direct grouped matmul: ", dtype);
+    }
+}
+
+struct AclTensorDeleter {
+    void operator()(aclTensor* tensor) const {
+        if (tensor != nullptr) {
+            aclDestroyTensor(tensor);
+        }
+    }
+};
+
+struct AclTensorListDeleter {
+    void operator()(aclTensorList* tensors) const {
+        if (tensors != nullptr) {
+            aclDestroyTensorList(tensors);
+        }
+    }
+};
+
+using AclTensorPtr = std::unique_ptr<aclTensor, AclTensorDeleter>;
+using AclTensorListPtr = std::unique_ptr<aclTensorList, AclTensorListDeleter>;
+
+AclTensorPtr MakeAclTensor(const at::Tensor& tensor) {
+    std::vector<int64_t> shape(tensor.sizes().begin(), tensor.sizes().end());
+    std::vector<int64_t> strides(tensor.strides().begin(), tensor.strides().end());
+    auto* acl_tensor = aclCreateTensor(shape.data(),
+                                       shape.size(),
+                                       ToAclDataType(tensor.scalar_type()),
+                                       strides.data(),
+                                       tensor.storage_offset(),
+                                       ACL_FORMAT_ND,
+                                       shape.data(),
+                                       shape.size(),
+                                       tensor.data_ptr());
+    TORCH_CHECK(acl_tensor != nullptr, "aclCreateTensor failed");
+    return AclTensorPtr(acl_tensor);
+}
+
+AclTensorPtr MakeTransposedWeightAclTensor(const at::Tensor& tensor) {
+    TORCH_CHECK(tensor.dim() == 3 && tensor.is_contiguous(), "stored weight must be contiguous [E, N, K]");
+    const std::array<int64_t, 3> view_shape{tensor.size(0), tensor.size(2), tensor.size(1)};
+    const std::array<int64_t, 3> view_strides{tensor.stride(0), tensor.stride(2), tensor.stride(1)};
+    const std::array<int64_t, 3> storage_shape{tensor.size(0), tensor.size(1), tensor.size(2)};
+    auto* acl_tensor = aclCreateTensor(view_shape.data(),
+                                       view_shape.size(),
+                                       ToAclDataType(tensor.scalar_type()),
+                                       view_strides.data(),
+                                       tensor.storage_offset(),
+                                       ACL_FORMAT_ND,
+                                       storage_shape.data(),
+                                       storage_shape.size(),
+                                       tensor.data_ptr());
+    TORCH_CHECK(acl_tensor != nullptr, "aclCreateTensor failed for weight");
+    return AclTensorPtr(acl_tensor);
+}
+
+AclTensorListPtr MakeAclTensorList(aclTensor* tensor) {
+    std::array<aclTensor*, 1> tensors{tensor};
+    auto* list = aclCreateTensorList(tensors.data(), tensors.size());
+    TORCH_CHECK(list != nullptr, "aclCreateTensorList failed");
+    return AclTensorListPtr(list);
+}
+
+at::Tensor GroupedMatmulDirect(const at::Tensor& x, const at::Tensor& weight, const at::Tensor& group_list, int64_t group_list_type, uint64_t stream_handle) {
+    TORCH_CHECK(x.device().type() == c10::DeviceType::PrivateUse1, "x must be an NPU tensor");
+    TORCH_CHECK(weight.device() == x.device() && group_list.device() == x.device(), "x, weight and group_list must be on the same NPU device");
+    TORCH_CHECK(x.dim() == 2, "x must have shape [M, K]");
+    TORCH_CHECK(weight.dim() == 3 && weight.size(0) == kDirectGroups, "weight must have stored shape [2560, N, K]");
+    TORCH_CHECK(group_list.dim() == 1 && group_list.numel() == kDirectGroups, "group_list must contain 2560 entries");
+    TORCH_CHECK(group_list.scalar_type() == at::ScalarType::Long, "group_list must be int64");
+    TORCH_CHECK(group_list_type == 1, "direct2560 only supports group_list_type=1");
+    TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || x.scalar_type() == at::ScalarType::BFloat16, "direct2560 only supports float16 and bfloat16");
+    TORCH_CHECK(weight.scalar_type() == x.scalar_type(), "x and weight dtypes must match");
+    TORCH_CHECK(x.size(1) == weight.size(2), "x K must match weight K");
+
+    auto out = at::empty({x.size(0), weight.size(1)}, x.options());
+    auto workspace = at::Tensor();
+
+    auto x_acl = MakeAclTensor(x);
+    auto weight_acl = MakeTransposedWeightAclTensor(weight);
+    auto group_list_acl = MakeAclTensor(group_list);
+    auto out_acl = MakeAclTensor(out);
+    auto x_list = MakeAclTensorList(x_acl.get());
+    auto weight_list = MakeAclTensorList(weight_acl.get());
+    auto out_list = MakeAclTensorList(out_acl.get());
+    // aclTensorList owns the tensors passed to aclCreateTensorList. Transfer
+    // ownership to the lists so the individual guards do not destroy them a
+    // second time when this function returns.
+    x_acl.release();
+    weight_acl.release();
+    out_acl.release();
+
+    uint64_t workspace_size = 0;
+    aclOpExecutor* executor = nullptr;
+    const auto status = aclnnDlinferGroupedMatmulDirectV5GetWorkspaceSize(x_list.get(),
+                                                                          weight_list.get(),
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          group_list_acl.get(),
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          2,
+                                                                          0,
+                                                                          group_list_type,
+                                                                          0,
+                                                                          nullptr,
+                                                                          out_list.get(),
+                                                                          nullptr,
+                                                                          nullptr,
+                                                                          &workspace_size,
+                                                                          &executor);
+    TORCH_CHECK(status == ACL_SUCCESS, "aclnnDlinferGroupedMatmulDirectV5GetWorkspaceSize failed, status=", status);
+    TORCH_CHECK(executor != nullptr, "GroupedMatmul returned a null executor");
+
+    void* workspace_addr = nullptr;
+    if (workspace_size > 0) {
+        workspace = at::empty({static_cast<int64_t>(workspace_size)}, x.options().dtype(at::ScalarType::Byte));
+        workspace_addr = workspace.data_ptr();
+    }
+
+    TORCH_CHECK(stream_handle != 0, "current NPU stream handle must not be null");
+    const auto stream = reinterpret_cast<aclrtStream>(stream_handle);
+    const auto execute_status = aclnnDlinferGroupedMatmulDirectV5(workspace_addr, workspace_size, executor, stream);
+    TORCH_CHECK(execute_status == ACL_SUCCESS, "aclnnDlinferGroupedMatmulDirectV5 failed, status=", execute_status);
+    return out;
+}
+
+}  // namespace
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { module.def("grouped_matmul", &GroupedMatmulDirect, "DLInfer bundled GroupedMatmul direct2560"); }

dlinfer/vendor/ascend/csrc/grouped_matmul_direct/vendor/op_api/include/aclnnop/aclnn_grouped_matmul.h

@@ -0,0 +1,63 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+#ifndef OP_API_INC_GROUPED_MATMUL_H
+#define OP_API_INC_GROUPED_MATMUL_H
+#include "aclnn/aclnn_base.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief aclnnDlinferGroupedMatmulDirect的第一段接口，根据具体的计算流程，计算workspace大小。
+ * @domain aclnn_ops_infer
+ *
+ * @param [in] x: 表示公式中的x，数据类型支持FLOAT16、BFLOAT16、INT8、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] weight:
+ * 表示公式中的weight，数据类型支持FLOAT16、BFLOAT16、INT8、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] biasOptional:
+ * 表示公式中的bias，数据类型支持BLOAT16、FLOAT16、FLOAT32、INT32数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] scaleOptional: 表示量化参数，数据类型支持UINT64数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] offsetOptional: 表示量化参数，数据类型支持FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] antiquantScaleOptional:
+ * 表示伪量化参数，数据类型支持FLOAT16，BFLOAT16数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] antiquantOffsetOptional:
+ * 表示伪量化参数，数据类型支持FLOAT16，BFLOAT16数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [in] groupListOptional: 可选参数，代表输入和输出M轴上的索引情况，数据类型支持INT64，支持的最大长度为128个。
+ * @param [in] splitItem:
+ * 整数型参数，代表输出是否要做tensor切分，0/1代表输出为多tensor；2/3代表输出为单tensor，默认值为0。
+ * @param [out] y: 表示公式中的y，数据类型支持FLOAT16、BFLOAT16、INT8、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。
+ * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
+ * @param [out] executor: 返回op执行器，包含算子计算流程。
+ * @return aclnnStatus: 返回状态码。
+ */
+__attribute__((visibility("default"))) aclnnStatus aclnnDlinferGroupedMatmulDirectGetWorkspaceSize(
+    const aclTensorList* x, const aclTensorList* weight, const aclTensorList* biasOptional,
+    const aclTensorList* scaleOptional, const aclTensorList* offsetOptional,
+    const aclTensorList* antiquantScaleOptional, const aclTensorList* antiquantOffsetOptional,
+    const aclIntArray* groupListOptional, int64_t splitItem, const aclTensorList* y, uint64_t* workspaceSize,
+    aclOpExecutor** executor);
+
+/**
+ * @brief aclnnDlinferGroupedMatmulDirect的第二段接口，用于执行计算。
+ * @param [in] workspace: 在npu device侧申请的workspace内存起址。
+ * @param [in] workspaceSize: 在npu device侧申请的workspace大小，由第一段接口aclnnDlinferGroupedMatmulDirectGetWorkspaceSize获取。
+ * @param [in] executor: op执行器，包含了算子计算流程。
+ * @param [in] stream: acl stream流。
+ * @return aclnnStatus: 返回状态码。
+ */
+__attribute__((visibility("default"))) aclnnStatus aclnnDlinferGroupedMatmulDirect(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
+                                         aclrtStream stream);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif