Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -97,8 +97,8 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]
| `--nv-gpu=[y\|n]` | 是否编译英伟达 GPU 接口实现 | n
| `--ascend-npu=[y\|n]` | 是否编译昇腾 NPU 接口实现 | n
| `--cambricon-mlu=[y\|n]` | 是否编译寒武纪 MLU 接口实现 | n
| `--metax-gpu=[y\|n]` | 是否编译沐曦 GPU 接口实现 | n
| `--use-mc=[y\|n]` | 是否沐曦 GPU 接口实现使用maca SDK | n
| `--metax-gpu=[y\|n]` | 是否编译 MetaX/MACA GPU 接口实现 | n
| `--mars-gpu=[y\|n]` | 是否编译 Mars/HPCC GPU 接口实现 | n
| `--moore-gpu=[y\|n]` | 是否编译摩尔线程 GPU 接口实现 | n
| `--iluvatar-gpu=[y\|n]` | 是否编译天数 GPU 接口实现 | n
| `--qy-gpu=[y\|n]` | 是否编译QY GPU 接口实现 | n
Expand Down Expand Up @@ -341,9 +341,9 @@ pip install -e .

```bash
# 测试单算子
python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali]
python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali]
# 测试全部算子
python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --ali]
python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --ali]
```

使用 -h 查看更多参数。
Expand All @@ -352,9 +352,9 @@ python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia

```shell
# 测试单算子
python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali]
python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali]
# 测试全部算子
python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali]
python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali]
```

#### 通信库(InfiniCCL)测试
Expand Down
1 change: 1 addition & 0 deletions include/infinicore.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ typedef enum {
INFINI_DEVICE_HYGON = 8,
INFINI_DEVICE_QY = 9,
INFINI_DEVICE_ALI = 10,
INFINI_DEVICE_MARS = 11,
INFINI_DEVICE_TYPE_COUNT
} infiniDevice_t;

Expand Down
5 changes: 3 additions & 2 deletions include/infinicore/adaptor/aten_adaptor.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
#if defined(ENABLE_HYGON_API)
#include <c10/hip/HIPGuard.h>
#include <c10/hip/HIPStream.h>
#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
Expand Down Expand Up @@ -53,6 +53,7 @@ inline at::Device to_at_device(const Device &device) {
// Treat CUDA-compatible vendor devices as CUDA devices for ATen tensor interoperability.
if (device.getType() == Device::Type::NVIDIA
|| device.getType() == Device::Type::METAX
|| device.getType() == Device::Type::MARS
|| device.getType() == Device::Type::QY
|| device.getType() == Device::Type::ILUVATAR
|| device.getType() == Device::Type::ALI
Expand Down Expand Up @@ -85,7 +86,7 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t);

#if defined(ENABLE_HYGON_API)
c10::hip::HIPStream get_hip_stream();
#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API)
c10::cuda::CUDAStream get_cuda_stream();
#endif

Expand Down
33 changes: 18 additions & 15 deletions include/infinicore/adaptor/flash_attention_adaptor.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,14 @@
#pragma once
#include "aten_adaptor.hpp"

// NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension
// exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds
// where the namespace is empty.
// NVIDIA flash-attn-nvidia.so uses namespace flash. The MetaX and Mars
// extensions export the same entry points at global scope.
//
// Ascend (aclnn C API path) does NOT use the flash:: namespace at all — the aclnn kernels are
// called directly from the dedicated *._ascend.cc implementation files, so this header is only
// included by the NVIDIA/MetaX/QY code paths. We still guard the namespace below so that
// included by the NVIDIA/MetaX/Mars/QY code paths. We still guard the namespace below so that
// existing code compiles unchanged when ENABLE_ASCEND_FLASH_ATTN is defined.
#if !defined(ENABLE_METAX_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
namespace flash {
#endif
std::vector<at::Tensor>
Expand All @@ -22,6 +21,10 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x
std::optional<at::Tensor> &softmax_lse_, // MetaX flash-attn dense fwd ABI includes an optional preallocated LSE tensor
#endif
std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT)
// Extended device ABIs accept an attention mask before the scalar options.
std::optional<at::Tensor> &attn_mask_,
#endif
const float p_dropout,
const float softmax_scale,
bool is_causal,
Expand All @@ -30,10 +33,10 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x
const float softcap,
const bool return_softmax,
std::optional<at::Generator> gen_
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
// MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX)
// Some device extensions append an auxiliary tensor after the generator.
,
std::optional<at::Tensor> &flash_attn_mars_ext_
std::optional<at::Tensor> &s_aux_
#endif
);

Expand All @@ -59,10 +62,10 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_hea
const float softcap,
const bool return_softmax,
std::optional<at::Generator> gen_
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
// MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX)
// Some device extensions append an auxiliary tensor.
,
std::optional<at::Tensor> &flash_attn_mars_ext_
std::optional<at::Tensor> &s_aux_
#endif
);

Expand Down Expand Up @@ -134,14 +137,14 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size
const float softcap,
bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
int num_splits
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
// MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn.
#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX)
// Some device extensions append an auxiliary tensor.
,
std::optional<at::Tensor> &flash_attn_mars_ext_
std::optional<at::Tensor> &s_aux_
#endif
);

#if !defined(ENABLE_METAX_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API) && !defined(ENABLE_ASCEND_FLASH_ATTN)
} // namespace flash
#endif
#endif // ENABLE_FLASH_ATTN
1 change: 1 addition & 0 deletions include/infinicore/device.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ class Device {
HYGON = INFINI_DEVICE_HYGON,
QY = INFINI_DEVICE_QY,
ALI = INFINI_DEVICE_ALI,
MARS = INFINI_DEVICE_MARS,
COUNT = INFINI_DEVICE_TYPE_COUNT,
};

Expand Down
1 change: 1 addition & 0 deletions python/infinicore/device.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ def _from_infinicore_device(infinicore_device: _infinicore.Device):
_infinicore.Device.Type.HYGON: "cuda",
_infinicore.Device.Type.QY: "cuda",
_infinicore.Device.Type.ALI: "cuda",
_infinicore.Device.Type.MARS: "cuda",
}


Expand Down
49 changes: 36 additions & 13 deletions scripts/metax_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,12 @@ def _first_existing_dir(paths: list[str]) -> str:
return ""


def _metax_toolkit_root(use_mc: bool) -> str:
"""Return toolkit root for MetaX builds (MACA when use-mc; otherwise HPCC)."""
if use_mc:
for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"):
v = os.environ.get(key, "").strip()
if v:
return v
return _first_existing_dir(["/opt/maca"])
return _first_existing_dir(["/opt/hpcc"])
def _toolkit_root(env_names: tuple[str, ...], fallback: str) -> str:
for key in env_names:
value = os.environ.get(key, "").strip()
if value:
return value
return _first_existing_dir([fallback])


def _prepend_path_var(name: str, prefixes: list[str]) -> None:
Expand All @@ -37,16 +34,14 @@ def set_env_for_metax_gpu(
"""
Prepend compiler include paths needed when building ATen-enabled C++ against torch headers.

This chooses paths based on xmake backend flags (e.g. --metax-gpu) and toolkit selection
(e.g. MetaX HPCC vs MACA when --use-mc=y).
MetaX always uses the MACA SDK. Mars/HPCC is configured separately.
"""
d = parse_xmake_cli_flag_values(flags)
if not truthy_flag_value(d.get("aten", "n")):
return

if truthy_flag_value(d.get("metax-gpu", "n")):
use_mc = truthy_flag_value(d.get("use-mc", "n"))
root = _metax_toolkit_root(use_mc=use_mc)
root = _toolkit_root(("MACA_PATH", "MACA_HOME", "MACA_ROOT"), "/opt/maca")
if not root:
return
dirs = [
Expand All @@ -64,4 +59,32 @@ def set_env_for_metax_gpu(
]
for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
_prepend_path_var(var, dirs)


def set_env_for_mars_gpu(
flags: str,
*,
parse_xmake_cli_flag_values,
truthy_flag_value,
) -> None:
"""Prepend HPCC compatibility headers for ATen-enabled Mars builds."""
d = parse_xmake_cli_flag_values(flags)
if not truthy_flag_value(d.get("aten", "n")):
return
if not truthy_flag_value(d.get("mars-gpu", "n")):
return

root = _toolkit_root(("HPCC_PATH", "HPCC_HOME"), "/opt/hpcc")
if not root:
return
dirs = [
os.path.join(root, "tools", "cu-bridge", "include"),
os.path.join(root, "include", "hcr"),
os.path.join(root, "include", "common"),
os.path.join(root, "include", "hcsparse"),
os.path.join(root, "include", "hcblas"),
os.path.join(root, "include", "hcsolver"),
os.path.join(root, "include"),
]
for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"):
_prepend_path_var(var, dirs)
35 changes: 26 additions & 9 deletions scripts/set_env.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import os
import platform

from metax_env import set_env_for_metax_gpu
from metax_env import set_env_for_mars_gpu, set_env_for_metax_gpu


def _parse_xmake_cli_flag_values(flags: str):
Expand Down Expand Up @@ -38,12 +38,16 @@ def set_env_by_config(flags: str) -> None:
parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values,
truthy_flag_value=_truthy_flag_value,
)
else:
pass
if _truthy_flag_value(d.get("mars-gpu", "n")):
set_env_for_mars_gpu(
flags,
parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values,
truthy_flag_value=_truthy_flag_value,
)


def set_env():
if os.environ.get("INFINI_ROOT") == None:
if os.environ.get("INFINI_ROOT") is None:
os.environ["INFINI_ROOT"] = os.path.expanduser("~/.infini")

if platform.system() == "Windows":
Expand All @@ -56,10 +60,23 @@ def set_env():
if new_path not in os.environ.get("PATH", ""):
os.environ["PATH"] = f"{new_path}:{os.environ.get('PATH', '')}"

new_lib_path = os.path.expanduser(os.environ.get("INFINI_ROOT") + "/lib")
if new_lib_path not in os.environ.get("LD_LIBRARY_PATH", ""):
os.environ["LD_LIBRARY_PATH"] = (
f"{new_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}"
)
new_lib_paths = []
infinirt_root = os.environ.get("INFINI_RT_ROOT")
if infinirt_root:
for subdir in ("lib", "lib64"):
candidate = os.path.join(infinirt_root, subdir)
if os.path.isdir(candidate):
new_lib_paths.append(candidate)
new_lib_paths.append(os.path.expanduser(os.environ["INFINI_ROOT"] + "/lib"))

current_lib_paths = [
path
for path in os.environ.get("LD_LIBRARY_PATH", "").split(":")
if path
]
for new_lib_path in reversed(new_lib_paths):
if new_lib_path not in current_lib_paths:
current_lib_paths.insert(0, new_lib_path)
os.environ["LD_LIBRARY_PATH"] = ":".join(current_lib_paths)
else:
raise RuntimeError("Unsupported platform.")
4 changes: 4 additions & 0 deletions src/bridge/infini/rt.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ inline constexpr ::infini::rt::Device::Type translate_to(infiniDevice_t device)
return ::infini::rt::Device::Type::kHygon;
case INFINI_DEVICE_ALI:
return ::infini::rt::Device::Type::kThead;
case INFINI_DEVICE_MARS:
return ::infini::rt::Device::Type::kMars;
default:
return ::infini::rt::Device::Type::kCount;
}
Expand All @@ -60,6 +62,8 @@ inline constexpr infiniDevice_t translate_from(::infini::rt::Device::Type device
return INFINI_DEVICE_HYGON;
case ::infini::rt::Device::Type::kThead:
return INFINI_DEVICE_ALI;
case ::infini::rt::Device::Type::kMars:
return INFINI_DEVICE_MARS;
default:
return INFINI_DEVICE_TYPE_COUNT;
}
Expand Down
3 changes: 2 additions & 1 deletion src/infiniccl-test/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ void printUsage() {
<< "[--master-addr <ip>] [--master-port <port>]" << std::endl
<< std::endl;
std::cout << " --<device>" << std::endl;
std::cout << " Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl
std::cout << " Specify the device type --(nvidia|cambricon|ascend|metax|mars|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl
<< std::endl;
std::cout << " --func" << std::endl
<< " allreduce: local all-reduce test across all visible devices." << std::endl
Expand Down Expand Up @@ -64,6 +64,7 @@ ParsedArgs parseArgs(int argc, char *argv[]) {
else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON)
else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND)
else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX)
else PARSE_DEVICE("--mars", INFINI_DEVICE_MARS)
else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE)
else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR)
else PARSE_DEVICE("--qy", INFINI_DEVICE_QY)
Expand Down
Loading
Loading