Skip to content

Commit bf0c825

Browse files
zhangyue207zhangyue207
authored
issue/995 fix paged attn on iluvatar
1 parent 70862bc commit bf0c825

5 files changed

Lines changed: 46 additions & 6 deletions

File tree

‎src/infiniop/ops/paged_attention/cuda/kernel_v2.cuh‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,11 @@ __device__ __forceinline__ float warpReduceMax(float x) {
3030
}
3131

3232
__device__ __forceinline__ unsigned int cvtaToShared(const void *ptr) {
33+
#if defined(ENABLE_ILUVATAR_API)
34+
return static_cast<unsigned int>(reinterpret_cast<uintptr_t>(ptr));
35+
#else
3336
return static_cast<unsigned int>(__cvta_generic_to_shared(ptr));
37+
#endif
3438
}
3539

3640
__device__ __forceinline__ void cpAsyncCaSharedGlobal16(void *dst_shared, const void *src_global) {

‎src/infiniop/ops/paged_attention/operator.cc‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
#include "../../handle.h"
33
#include "infiniop/ops/paged_attention.h"
44

5-
#ifdef ENABLE_NVIDIA_API
5+
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
66
#include "nvidia/paged_attention_nvidia.cuh"
77
#endif
88
#ifdef ENABLE_METAX_API
@@ -36,6 +36,9 @@ __C infiniStatus_t infiniopCreatePagedAttentionDescriptor(
3636
#endif
3737
#ifdef ENABLE_METAX_API
3838
CREATE(INFINI_DEVICE_METAX, metax)
39+
#endif
40+
#ifdef ENABLE_ILUVATAR_API
41+
CREATE(INFINI_DEVICE_ILUVATAR, nvidia)
3942
#endif
4043
default:
4144
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -57,6 +60,9 @@ __C infiniStatus_t infiniopGetPagedAttentionWorkspaceSize(
5760
#endif
5861
#ifdef ENABLE_METAX_API
5962
GET(INFINI_DEVICE_METAX, metax)
63+
#endif
64+
#ifdef ENABLE_ILUVATAR_API
65+
GET(INFINI_DEVICE_ILUVATAR, nvidia)
6066
#endif
6167
default:
6268
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -82,6 +88,9 @@ __C infiniStatus_t infiniopPagedAttention(
8288
#endif
8389
#ifdef ENABLE_METAX_API
8490
CALCULATE(INFINI_DEVICE_METAX, metax)
91+
#endif
92+
#ifdef ENABLE_ILUVATAR_API
93+
CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia)
8594
#endif
8695
default:
8796
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -102,6 +111,9 @@ __C infiniStatus_t infiniopDestroyPagedAttentionDescriptor(
102111
#endif
103112
#ifdef ENABLE_METAX_API
104113
DESTROY(INFINI_DEVICE_METAX, metax)
114+
#endif
115+
#ifdef ENABLE_ILUVATAR_API
116+
DESTROY(INFINI_DEVICE_ILUVATAR, nvidia)
105117
#endif
106118
default:
107119
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

‎src/infiniop/ops/paged_attention_prefill/operator.cc‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
#include "../../handle.h"
33
#include "infiniop/ops/paged_attention_prefill.h"
44

5-
#ifdef ENABLE_NVIDIA_API
5+
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
66
#include "nvidia/paged_attention_prefill_nvidia.cuh"
77
#endif
88
#ifdef ENABLE_METAX_API
@@ -38,6 +38,9 @@ __C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor(
3838
#endif
3939
#ifdef ENABLE_METAX_API
4040
CREATE(INFINI_DEVICE_METAX, metax)
41+
#endif
42+
#ifdef ENABLE_ILUVATAR_API
43+
CREATE(INFINI_DEVICE_ILUVATAR, nvidia)
4144
#endif
4245
default:
4346
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -59,6 +62,9 @@ __C infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize(
5962
#endif
6063
#ifdef ENABLE_METAX_API
6164
GET(INFINI_DEVICE_METAX, metax)
65+
#endif
66+
#ifdef ENABLE_ILUVATAR_API
67+
GET(INFINI_DEVICE_ILUVATAR, nvidia)
6268
#endif
6369
default:
6470
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -87,6 +93,9 @@ __C infiniStatus_t infiniopPagedAttentionPrefill(
8793
#endif
8894
#ifdef ENABLE_METAX_API
8995
CALCULATE(INFINI_DEVICE_METAX, metax)
96+
#endif
97+
#ifdef ENABLE_ILUVATAR_API
98+
CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia)
9099
#endif
91100
default:
92101
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -107,6 +116,9 @@ __C infiniStatus_t infiniopDestroyPagedAttentionPrefillDescriptor(
107116
#endif
108117
#ifdef ENABLE_METAX_API
109118
DESTROY(INFINI_DEVICE_METAX, metax)
119+
#endif
120+
#ifdef ENABLE_ILUVATAR_API
121+
DESTROY(INFINI_DEVICE_ILUVATAR, nvidia)
110122
#endif
111123
default:
112124
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

‎src/infiniop/ops/paged_caching/operator.cc‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
#include "../../handle.h"
33
#include "infiniop/ops/paged_caching.h"
44

5-
#ifdef ENABLE_NVIDIA_API
5+
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API)
66
#include "nvidia/paged_caching_nvidia.cuh"
77
#endif
88
#ifdef ENABLE_METAX_API
@@ -31,6 +31,9 @@ __C infiniStatus_t infiniopCreatePagedCachingDescriptor(
3131
#endif
3232
#ifdef ENABLE_METAX_API
3333
CREATE(INFINI_DEVICE_METAX, metax)
34+
#endif
35+
#ifdef ENABLE_ILUVATAR_API
36+
CREATE(INFINI_DEVICE_ILUVATAR, nvidia)
3437
#endif
3538
default:
3639
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -52,6 +55,9 @@ __C infiniStatus_t infiniopGetPagedCachingWorkspaceSize(
5255
#endif
5356
#ifdef ENABLE_METAX_API
5457
GET(INFINI_DEVICE_METAX, metax)
58+
#endif
59+
#ifdef ENABLE_ILUVATAR_API
60+
GET(INFINI_DEVICE_ILUVATAR, nvidia)
5561
#endif
5662
default:
5763
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -77,6 +83,9 @@ __C infiniStatus_t infiniopPagedCaching(
7783
#endif
7884
#ifdef ENABLE_METAX_API
7985
CALCULATE(INFINI_DEVICE_METAX, metax)
86+
#endif
87+
#ifdef ENABLE_ILUVATAR_API
88+
CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia)
8089
#endif
8190
default:
8291
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -97,6 +106,9 @@ __C infiniStatus_t infiniopDestroyPagedCachingDescriptor(
97106
#endif
98107
#ifdef ENABLE_METAX_API
99108
DESTROY(INFINI_DEVICE_METAX, metax)
109+
#endif
110+
#ifdef ENABLE_ILUVATAR_API
111+
DESTROY(INFINI_DEVICE_ILUVATAR, nvidia)
100112
#endif
101113
default:
102114
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

‎xmake/iluvatar.lua‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,14 @@ target("infiniop-iluvatar")
4242
add_links("cudart", "cublas", "cudnn")
4343

4444
set_warnings("all", "error")
45-
add_cuflags("-Wno-error=unused-private-field")
45+
add_cuflags("-Wno-error=unused-private-field", "-Wno-error=unused-variable", "-Wno-unused-variable")
4646
add_cuflags("-fPIC", "-x", "ivcore", "-std=c++17", {force = true})
4747
if has_config("ivcore-20") then
4848
add_cuflags("--cuda-gpu-arch=ivcore20", {force = true})
4949
end
5050
add_culdflags("-fPIC")
51-
add_cxflags("-fPIC")
52-
add_cxxflags("-fPIC")
51+
add_cxflags("-fPIC", "-Wno-error=unused-variable", "-Wno-unused-variable")
52+
add_cxxflags("-fPIC", "-Wno-error=unused-variable", "-Wno-unused-variable")
5353

5454
-- set_languages("cxx17") 天数似乎不能用这个配置
5555
add_files("../src/infiniop/devices/nvidia/*.cu", "../src/infiniop/ops/*/nvidia/*.cu")

0 commit comments

Comments
 (0)