Skip to content

Commit d7147b5

Browse files
zhangyue207zhangyue207
committed
fa hygon
1 parent 9319161 commit d7147b5

22 files changed

Lines changed: 888 additions & 48 deletions

File tree

‎include/infinicore/adaptor/aten_adaptor.hpp‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,12 @@
55

66
#include <ATen/ATen.h>
77

8-
#ifdef ENABLE_NVIDIA_API
8+
#if defined(ENABLE_NVIDIA_API)
99
#include <ATen/cuda/CUDAContext.h>
1010
#include <c10/cuda/CUDAGuard.h>
11+
#elif defined(ENABLE_HYGON_API)
12+
#include <ATen/hip/HIPContext.h>
13+
#include <c10/hip/HIPGuard.h>
1114
#endif
1215

1316
namespace infinicore::adaptor {
@@ -29,7 +32,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
2932
}
3033

3134
inline at::Device to_at_device(const Device &device) {
32-
if (device.getType() == Device::Type::NVIDIA) {
35+
if (device.getType() == Device::Type::NVIDIA
36+
|| device.getType() == Device::Type::HYGON) {
3337
return at::Device(at::kCUDA, device.getIndex());
3438
} else if (device.getType() == Device::Type::CPU) {
3539
return at::Device(at::kCPU);
@@ -40,8 +44,14 @@ inline at::Device to_at_device(const Device &device) {
4044

4145
at::Tensor to_aten_tensor(const infinicore::Tensor &t);
4246

43-
#ifdef ENABLE_NVIDIA_API
44-
c10::cuda::CUDAStream get_cuda_stream();
47+
#if defined(ENABLE_HYGON_API)
48+
using TorchStream = c10::hip::HIPStream;
49+
using TorchStreamGuard = c10::hip::HIPStreamGuard;
50+
TorchStream get_cuda_stream();
51+
#elif defined(ENABLE_NVIDIA_API)
52+
using TorchStream = c10::cuda::CUDAStream;
53+
using TorchStreamGuard = c10::cuda::CUDAStreamGuard;
54+
TorchStream get_cuda_stream();
4555
#endif
4656
} // namespace infinicore::adaptor
4757

‎include/infinicore/adaptor/flash_attention_adaptor.hpp‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,5 +110,46 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size
110110
bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
111111
int num_splits);
112112

113+
#ifdef ENABLE_HYGON_API
114+
// Hygon-specific wrappers resolved via dlsym at runtime.
115+
116+
std::vector<at::Tensor>
117+
vllm_mha_varlen_fwd(at::Tensor &q,
118+
const at::Tensor &k,
119+
const at::Tensor &v,
120+
std::optional<at::Tensor> &out_,
121+
const at::Tensor &cu_seqlens_q,
122+
const at::Tensor &cu_seqlens_k,
123+
std::optional<at::Tensor> &seqused_k,
124+
std::optional<const at::Tensor> &leftpad_k_,
125+
std::optional<at::Tensor> &block_table_,
126+
std::optional<at::Tensor> &alibi_slopes_,
127+
int max_seqlen_q,
128+
const int max_seqlen_k,
129+
const float p_dropout,
130+
const float softmax_scale,
131+
const bool zero_tensors,
132+
bool is_causal,
133+
int window_size_left,
134+
int window_size_right,
135+
const float softcap,
136+
const bool return_softmax,
137+
std::optional<at::Generator> gen_);
138+
139+
void paged_attention(at::Tensor &out,
140+
at::Tensor &q,
141+
at::Tensor &k_cache,
142+
at::Tensor &v_cache,
143+
double scale,
144+
at::Tensor &block_table,
145+
at::Tensor &context_lens,
146+
std::optional<at::Tensor> alibi_slopes,
147+
const std::string &kv_cache_dtype,
148+
std::optional<at::Tensor> q_scale,
149+
std::optional<at::Tensor> k_scale,
150+
std::optional<at::Tensor> v_scale,
151+
int max_context_len);
152+
#endif // ENABLE_HYGON_API
153+
113154
} // namespace flash
114155
#endif // ENABLE_FLASH_ATTN

‎python/infinicore/__init__.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,17 @@
7676
zeros,
7777
)
7878

79+
# Re-attempt flash_attn preload after _infinicore.so (and its DTK/HIP
80+
# dependencies) are loaded. The initial preload() above may have failed
81+
# because flash_attn_2_cuda.so's dependencies were not yet available.
82+
# Now that _infinicore.so is loaded we can call dlopen(path, RTLD_GLOBAL)
83+
# to upgrade the already-loaded library from RTLD_LOCAL to RTLD_GLOBAL,
84+
# making its symbols visible to dlsym(RTLD_DEFAULT).
85+
with contextlib.suppress(Exception):
86+
from ._preload import preload_flash_attn
87+
88+
preload_flash_attn()
89+
7990
__all__ = [
8091
# Modules.
8192
"context",

‎python/infinicore/_preload.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
import ctypes
2+
import glob
3+
import importlib.util
24
import os
5+
import sys
36
from typing import Iterable, List
47

58

@@ -63,12 +66,94 @@ def preload_hpcc() -> None:
6366
_try_load(prefixes, lib)
6467

6568

69+
def preload_torch_hip() -> None:
70+
"""
71+
Best-effort preload of torch HIP runtime libs with RTLD_GLOBAL.
72+
73+
This helps external extensions resolve c10::hip symbols when they are
74+
not recorded as direct DT_NEEDED dependencies.
75+
"""
76+
spec = importlib.util.find_spec("torch")
77+
if spec is None or not spec.origin:
78+
return
79+
torch_dir = os.path.dirname(spec.origin)
80+
torch_libdir = os.path.join(torch_dir, "lib")
81+
if not os.path.isdir(torch_libdir):
82+
return
83+
84+
libs = [
85+
"libtorch_global_deps.so",
86+
"libc10.so",
87+
"libc10_hip.so",
88+
"libtorch_cpu.so",
89+
"libtorch.so",
90+
"libtorch_hip.so",
91+
]
92+
for lib in libs:
93+
full = os.path.join(torch_libdir, lib)
94+
if os.path.exists(full):
95+
try:
96+
ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL)
97+
except OSError:
98+
# Best-effort preload, continue on errors.
99+
pass
100+
101+
102+
def preload_flash_attn() -> None:
103+
"""
104+
Best-effort preload of flash_attn_2_cuda extension with RTLD_GLOBAL.
105+
106+
InfiniCore hygon wrapper resolves C symbols like `mha_varlen_fwd` from the
107+
flash-attn extension at runtime via dlsym(RTLD_DEFAULT, ...). The symbols
108+
only need to be available when the operator is actually called, not at
109+
library load time. So this preload is a convenience — if it fails, the
110+
symbols will be resolved later when torch + flash_attn are imported by
111+
the application (e.g. InfiniLM).
112+
"""
113+
candidates: List[str] = []
114+
from_env = os.getenv("FLASH_ATTN_PREBUILT")
115+
if from_env:
116+
if os.path.isfile(from_env):
117+
candidates.append(from_env)
118+
elif os.path.isdir(from_env):
119+
candidates.extend(glob.glob(os.path.join(from_env, "flash_attn_2_cuda*.so")))
120+
121+
# Try resolving via Python import metadata.
122+
spec = importlib.util.find_spec("flash_attn_2_cuda")
123+
if spec and spec.origin and os.path.exists(spec.origin):
124+
candidates.append(spec.origin)
125+
126+
# Fallback: scan python paths for extension module.
127+
for p in sys.path:
128+
if not p:
129+
continue
130+
candidates.extend(glob.glob(os.path.join(p, "flash_attn_2_cuda*.so")))
131+
132+
# Common installation locations.
133+
candidates.extend(glob.glob("/usr/local/lib/python*/dist-packages/flash_attn_2_cuda*.so"))
134+
candidates.extend(glob.glob("/root/.infini/lib/flash_attn_2_cuda*.so"))
135+
136+
seen = set()
137+
for so_path in candidates:
138+
if not so_path or so_path in seen:
139+
continue
140+
seen.add(so_path)
141+
if not os.path.exists(so_path):
142+
continue
143+
try:
144+
ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)
145+
return
146+
except OSError:
147+
continue
148+
149+
66150
def _should_preload_device(device_type: str) -> bool:
67151
"""
68152
Check if preload is needed for a specific device type.
69153
"""
70154
device_env_map = {
71155
"METAX": ["HPCC_PATH", "INFINICORE_PRELOAD_HPCC"], # HPCC/METAX
156+
"HYGON": ["DTK_ROOT", "INFINICORE_PRELOAD_TORCH_HIP"],
72157
# Add other device types here as needed:
73158
# "ASCEND": ["ASCEND_PATH"],
74159
# "CAMBRICON": ["NEUWARE_HOME"],
@@ -90,6 +175,8 @@ def preload_device(device_type: str) -> None:
90175
"""
91176
if device_type == "METAX":
92177
preload_hpcc()
178+
elif device_type == "HYGON":
179+
preload_torch_hip()
93180
# Add other device preload functions here as needed:
94181
# elif device_type == "ASCEND":
95182
# preload_ascend()
@@ -103,9 +190,20 @@ def preload() -> None:
103190
This function detects available device types and preloads their runtime libraries
104191
if the environment indicates they are needed.
105192
"""
193+
# Always try torch HIP preload first (best-effort, no-op if torch/HIP is absent).
194+
try:
195+
preload_torch_hip()
196+
except Exception:
197+
pass
198+
try:
199+
preload_flash_attn()
200+
except Exception:
201+
pass
202+
106203
# Device types that may require preload
107204
device_types = [
108205
"METAX", # HPCC/METAX
206+
"HYGON",
109207
# Add other device types here as they are implemented:
110208
# "ASCEND",
111209
# "CAMBRICON",

‎src/infinicore/adaptor/aten_adaptor.cc‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,13 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t) {
3232
options);
3333
}
3434

35-
#ifdef ENABLE_NVIDIA_API
36-
c10::cuda::CUDAStream get_cuda_stream() {
35+
#if defined(ENABLE_HYGON_API)
36+
TorchStream get_cuda_stream() {
37+
return c10::hip::getStreamFromExternal(
38+
hipStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex());
39+
}
40+
#elif defined(ENABLE_NVIDIA_API)
41+
TorchStream get_cuda_stream() {
3742
return c10::cuda::getStreamFromExternal(
3843
cudaStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex());
3944
}

0 commit comments

Comments
 (0)