diff --git a/python/infinicore/_preload.py b/python/infinicore/_preload.py index 396e68dbd..153356ccb 100644 --- a/python/infinicore/_preload.py +++ b/python/infinicore/_preload.py @@ -191,6 +191,10 @@ def _should_preload_device(device_type: str) -> bool: return True if device_type == "CAMBRICON": return importlib.util.find_spec("torch_mlu") is not None + if device_type == "ILUVATAR": + return importlib.util.find_spec("vllm_iluvatar") is not None or os.path.isdir( + "/usr/local/corex" + ) return False @@ -213,6 +217,9 @@ def preload_device(device_type: str) -> None: # Add other device preload functions here as needed: elif device_type == "CAMBRICON": preload_cambricon() + elif device_type == "ILUVATAR": + preload_torch() + preload_flash_attn() # etc. @@ -231,6 +238,7 @@ def preload() -> None: "HYGON", "ASCEND", "CAMBRICON", + "ILUVATAR", # Add other device types here as they are implemented: # etc. ] diff --git a/src/infiniop/ops/causal_conv1d/operator.cc b/src/infiniop/ops/causal_conv1d/operator.cc index a4a46f658..7c3d02242 100644 --- a/src/infiniop/ops/causal_conv1d/operator.cc +++ b/src/infiniop/ops/causal_conv1d/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/causal_conv1d.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) #include "nvidia/causal_conv1d_nvidia.cuh" #endif #ifdef ENABLE_METAX_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreateCausalConv1dDescriptor( #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif @@ -73,6 +76,9 @@ __INFINI_C infiniStatus_t infiniopGetCausalConv1dWorkspaceSize( #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif @@ -116,6 +122,9 @@ __INFINI_C infiniStatus_t infiniopCausalConv1d( #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif @@ -145,6 +154,9 @@ __INFINI_C infiniStatus_t infiniopDestroyCausalConv1dDescriptor( #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif diff --git a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc index c5db78af8..279fa1fd3 100644 --- a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc @@ -4,7 +4,7 @@ #include "../../handle.h" #include "infiniop/ops/chunk_gated_delta_rule.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) #include "nvidia/chunk_gated_delta_rule_nvidia.cuh" #endif #ifdef ENABLE_METAX_API @@ -50,6 +50,9 @@ __INFINI_C infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif @@ -79,6 +82,9 @@ __INFINI_C infiniStatus_t infiniopGetChunkGatedDeltaRuleWorkspaceSize( #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif @@ -115,6 +121,9 @@ __INFINI_C infiniStatus_t infiniopChunkGatedDeltaRule( #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif @@ -143,6 +152,9 @@ __INFINI_C infiniStatus_t infiniopDestroyChunkGatedDeltaRuleDescriptor( #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc index 5c7ed9246..bc6c65508 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc @@ -4,7 +4,7 @@ #include "../../handle.h" #include "infiniop/ops/recurrent_gated_delta_rule.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) #include "nvidia/recurrent_gated_delta_rule_nvidia.cuh" #endif #ifdef ENABLE_METAX_API @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor( #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif @@ -81,6 +84,9 @@ __INFINI_C infiniStatus_t infiniopGetRecurrentGatedDeltaRuleWorkspaceSize( #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif @@ -121,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopRecurrentGatedDeltaRule( #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif @@ -152,6 +161,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRecurrentGatedDeltaRuleDescriptor( #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia) #endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia) +#endif #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif