Skip to content

Commit 2606921

Browse files
GodKickMyAssGodKickMyAss
committed
feat: add ERNIE 4.5 VL core support
1 parent 8d82139 commit 2606921

13 files changed

Lines changed: 1306 additions & 12 deletions

File tree

‎include/infinicore/ops.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
#include "ops/cross_entropy.hpp"
2727
#include "ops/deepseek_moe.hpp"
2828
#include "ops/embedding.hpp"
29+
#include "ops/ernie45_rope.hpp"
2930
#include "ops/flash_attention.hpp"
3031
#include "ops/fmin.hpp"
3132
#include "ops/fmod.hpp"

‎include/infinicore/ops/ernie45_rope.hpp‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#pragma once
2+
3+
#include "../device.hpp"
4+
#include "../graph/graph.hpp"
5+
#include "../tensor.hpp"
6+
#include "common/op.hpp"
7+
8+
namespace infinicore::op {
9+
10+
INFINICORE_GRAPH_OP_CLASS(Ernie45MRoPE, Tensor, Tensor, const Tensor &, double, size_t, size_t, size_t);
11+
INFINICORE_GRAPH_OP_CLASS(Ernie45VisionRoPE, Tensor, Tensor, const Tensor &, double);
12+
13+
Tensor ernie45_mrope_(Tensor q,
14+
Tensor k,
15+
const Tensor &positions,
16+
double rope_theta,
17+
size_t section_h,
18+
size_t section_w,
19+
size_t section_t);
20+
21+
Tensor ernie45_vision_rope_(Tensor q,
22+
Tensor k,
23+
const Tensor &positions,
24+
double rope_theta);
25+
26+
} // namespace infinicore::op

‎include/infiniop/ops/ernie45_rope.h‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#ifndef __INFINIOP_ERNIE45_ROPE_API_H__
2+
#define __INFINIOP_ERNIE45_ROPE_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
6+
typedef struct InfiniopDescriptor *infiniopErnie45MropeDescriptor_t;
7+
typedef struct InfiniopDescriptor *infiniopErnie45VisionRopeDescriptor_t;
8+
9+
__INFINI_C __export infiniStatus_t infiniopCreateErnie45MropeDescriptor(
10+
infiniopHandle_t handle,
11+
infiniopErnie45MropeDescriptor_t *desc_ptr,
12+
infiniopTensorDescriptor_t q,
13+
infiniopTensorDescriptor_t k,
14+
infiniopTensorDescriptor_t positions,
15+
double rope_theta,
16+
size_t section_h,
17+
size_t section_w,
18+
size_t section_t);
19+
20+
__INFINI_C __export infiniStatus_t infiniopGetErnie45MropeWorkspaceSize(infiniopErnie45MropeDescriptor_t desc,
21+
size_t *size);
22+
23+
__INFINI_C __export infiniStatus_t infiniopErnie45Mrope(
24+
infiniopErnie45MropeDescriptor_t desc,
25+
void *workspace,
26+
size_t workspace_size,
27+
void *q,
28+
void *k,
29+
const void *positions,
30+
void *stream);
31+
32+
__INFINI_C __export infiniStatus_t infiniopDestroyErnie45MropeDescriptor(infiniopErnie45MropeDescriptor_t desc);
33+
34+
__INFINI_C __export infiniStatus_t infiniopCreateErnie45VisionRopeDescriptor(
35+
infiniopHandle_t handle,
36+
infiniopErnie45VisionRopeDescriptor_t *desc_ptr,
37+
infiniopTensorDescriptor_t q,
38+
infiniopTensorDescriptor_t k,
39+
infiniopTensorDescriptor_t positions,
40+
double rope_theta);
41+
42+
__INFINI_C __export infiniStatus_t infiniopGetErnie45VisionRopeWorkspaceSize(infiniopErnie45VisionRopeDescriptor_t desc,
43+
size_t *size);
44+
45+
__INFINI_C __export infiniStatus_t infiniopErnie45VisionRope(
46+
infiniopErnie45VisionRopeDescriptor_t desc,
47+
void *workspace,
48+
size_t workspace_size,
49+
void *q,
50+
void *k,
51+
const void *positions,
52+
void *stream);
53+
54+
__INFINI_C __export infiniStatus_t infiniopDestroyErnie45VisionRopeDescriptor(infiniopErnie45VisionRopeDescriptor_t desc);
55+
56+
#endif

‎src/infinicore/device_event.cc‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,32 @@ DeviceEvent::DeviceEvent(DeviceEvent &&other) noexcept
3838
other.is_recorded_ = false;
3939
}
4040

41+
namespace {
42+
43+
void destroy_event_on_device(infinirtEvent_t event, const Device &device) noexcept {
44+
try {
45+
Device current_device = context::getDevice();
46+
bool switched_device = current_device != device;
47+
if (switched_device) {
48+
context::setDevice(device);
49+
}
50+
context::destroyEvent(event);
51+
if (switched_device) {
52+
context::setDevice(current_device);
53+
}
54+
} catch (...) {
55+
// Event cleanup can happen during interpreter shutdown; do not throw
56+
// from destructors or noexcept move assignment cleanup.
57+
}
58+
}
59+
60+
} // namespace
61+
4162
DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept {
4263
if (this != &other) {
4364
// Clean up current resources
4465
if (event_ != nullptr) {
45-
context::destroyEvent(event_);
66+
destroy_event_on_device(event_, device_);
4667
}
4768

4869
// Transfer ownership
@@ -59,7 +80,7 @@ DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept {
5980

6081
DeviceEvent::~DeviceEvent() {
6182
if (event_ != nullptr) {
62-
context::destroyEvent(event_);
83+
destroy_event_on_device(event_, device_);
6384
}
6485
}
6586

‎src/infinicore/memory.cc‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include "infinicore/memory.hpp"
2+
#include "infinicore/context/context.hpp"
23

34
namespace infinicore {
45

@@ -14,7 +15,25 @@ Memory::~Memory() {
1415
return;
1516
}
1617
try {
17-
deleter_(data_);
18+
Device current_device = context::getDevice();
19+
bool switched_device = current_device != device_;
20+
if (switched_device) {
21+
context::setDevice(device_);
22+
}
23+
try {
24+
deleter_(data_);
25+
} catch (...) {
26+
if (switched_device) {
27+
try {
28+
context::setDevice(current_device);
29+
} catch (...) {
30+
}
31+
}
32+
throw;
33+
}
34+
if (switched_device) {
35+
context::setDevice(current_device);
36+
}
1837
} catch (...) {
1938
// Memory can be released during interpreter/static shutdown after
2039
// allocator metadata has already been torn down. Destructors must not

‎src/infinicore/ops/ernie45_rope/ernie45_rope.cc‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#include "infinicore/ops/ernie45_rope.hpp"
2+
#include "../../utils.hpp"
3+
4+
namespace infinicore::op {
5+
6+
INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Ernie45MRoPE);
7+
INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Ernie45VisionRoPE);
8+
9+
Ernie45MRoPE::Ernie45MRoPE(Tensor q,
10+
Tensor k,
11+
const Tensor &positions,
12+
double rope_theta,
13+
size_t section_h,
14+
size_t section_w,
15+
size_t section_t) {
16+
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, k, positions);
17+
INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, k, positions, rope_theta, section_h, section_w, section_t);
18+
}
19+
20+
void Ernie45MRoPE::execute(Tensor q,
21+
Tensor k,
22+
const Tensor &positions,
23+
double rope_theta,
24+
size_t section_h,
25+
size_t section_w,
26+
size_t section_t) {
27+
INFINICORE_GRAPH_OP_RECORD_OR_RUN(Ernie45MRoPE, q, k, positions, rope_theta, section_h, section_w, section_t);
28+
}
29+
30+
Tensor ernie45_mrope_(Tensor q,
31+
Tensor k,
32+
const Tensor &positions,
33+
double rope_theta,
34+
size_t section_h,
35+
size_t section_w,
36+
size_t section_t) {
37+
Ernie45MRoPE::execute(q, k, positions, rope_theta, section_h, section_w, section_t);
38+
return q;
39+
}
40+
41+
Ernie45VisionRoPE::Ernie45VisionRoPE(Tensor q,
42+
Tensor k,
43+
const Tensor &positions,
44+
double rope_theta) {
45+
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, k, positions);
46+
INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, k, positions, rope_theta);
47+
}
48+
49+
void Ernie45VisionRoPE::execute(Tensor q,
50+
Tensor k,
51+
const Tensor &positions,
52+
double rope_theta) {
53+
INFINICORE_GRAPH_OP_RECORD_OR_RUN(Ernie45VisionRoPE, q, k, positions, rope_theta);
54+
}
55+
56+
Tensor ernie45_vision_rope_(Tensor q,
57+
Tensor k,
58+
const Tensor &positions,
59+
double rope_theta) {
60+
Ernie45VisionRoPE::execute(q, k, positions, rope_theta);
61+
return q;
62+
}
63+
64+
} // namespace infinicore::op

‎src/infinicore/ops/ernie45_rope/ernie45_rope_infiniop.cc‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
#include "infinicore/ops/ernie45_rope.hpp"
2+
3+
#include "../infiniop_impl.hpp"
4+
#include "infiniop/ops/ernie45_rope.h"
5+
6+
namespace infinicore::op::ernie45_rope_impl::infiniop {
7+
8+
namespace mrope {
9+
INFINIOP_CACHABLE_DESCRIPTOR(MropeDescriptor, Ernie45Mrope, 100);
10+
} // namespace mrope
11+
12+
namespace vision {
13+
INFINIOP_CACHABLE_DESCRIPTOR(VisionDescriptor, Ernie45VisionRope, 100);
14+
} // namespace vision
15+
16+
namespace mrope_op {
17+
18+
struct MropePlannedMeta {
19+
std::shared_ptr<mrope::MropeDescriptor> descriptor;
20+
graph::GraphTensor workspace;
21+
graph::GraphTensor q;
22+
graph::GraphTensor k;
23+
graph::GraphTensor positions;
24+
};
25+
26+
void *plan_mrope(Tensor q,
27+
Tensor k,
28+
const Tensor &positions,
29+
double rope_theta,
30+
size_t section_h,
31+
size_t section_w,
32+
size_t section_t) {
33+
size_t key = hash_combine(q, k, positions, rope_theta, section_h, section_w, section_t);
34+
std::shared_ptr<mrope::MropeDescriptor> descriptor;
35+
{
36+
auto device__ = context::getDevice();
37+
auto &cache__ = mrope::caches.getCache(device__);
38+
descriptor = cache__.get(key).value_or(nullptr);
39+
if (!descriptor) {
40+
descriptor = std::make_shared<mrope::MropeDescriptor>(nullptr);
41+
INFINICORE_CHECK_ERROR(infiniopCreateErnie45MropeDescriptor(
42+
context::getInfiniopHandle(device__),
43+
&descriptor->desc,
44+
q->desc(), k->desc(), positions->desc(), rope_theta, section_h, section_w, section_t));
45+
cache__.put(key, descriptor);
46+
}
47+
}
48+
INFINIOP_WORKSPACE_TENSOR(workspace, Ernie45Mrope, descriptor);
49+
return new MropePlannedMeta{
50+
descriptor,
51+
graph::GraphTensor(workspace),
52+
graph::GraphTensor(q),
53+
graph::GraphTensor(k),
54+
graph::GraphTensor(positions)};
55+
}
56+
57+
void run_mrope(void *planned_meta) {
58+
auto *p = reinterpret_cast<MropePlannedMeta *>(planned_meta);
59+
INFINICORE_CHECK_ERROR(
60+
infiniopErnie45Mrope(
61+
p->descriptor->desc,
62+
p->workspace->data(),
63+
p->workspace->numel(),
64+
p->q->data(),
65+
p->k->data(),
66+
p->positions->data(),
67+
context::getStream()));
68+
}
69+
70+
void cleanup_mrope(void **planned_meta_ptr) {
71+
delete *reinterpret_cast<MropePlannedMeta **>(planned_meta_ptr);
72+
*planned_meta_ptr = nullptr;
73+
}
74+
75+
INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Ernie45MRoPE, &plan_mrope, &run_mrope, &cleanup_mrope);
76+
77+
} // namespace mrope_op
78+
79+
namespace vision_op {
80+
81+
struct VisionPlannedMeta {
82+
std::shared_ptr<vision::VisionDescriptor> descriptor;
83+
graph::GraphTensor workspace;
84+
graph::GraphTensor q;
85+
graph::GraphTensor k;
86+
graph::GraphTensor positions;
87+
};
88+
89+
void *plan_vision(Tensor q,
90+
Tensor k,
91+
const Tensor &positions,
92+
double rope_theta) {
93+
size_t key = hash_combine(q, k, positions, rope_theta);
94+
std::shared_ptr<vision::VisionDescriptor> descriptor;
95+
{
96+
auto device__ = context::getDevice();
97+
auto &cache__ = vision::caches.getCache(device__);
98+
descriptor = cache__.get(key).value_or(nullptr);
99+
if (!descriptor) {
100+
descriptor = std::make_shared<vision::VisionDescriptor>(nullptr);
101+
INFINICORE_CHECK_ERROR(infiniopCreateErnie45VisionRopeDescriptor(
102+
context::getInfiniopHandle(device__),
103+
&descriptor->desc,
104+
q->desc(), k->desc(), positions->desc(), rope_theta));
105+
cache__.put(key, descriptor);
106+
}
107+
}
108+
INFINIOP_WORKSPACE_TENSOR(workspace, Ernie45VisionRope, descriptor);
109+
return new VisionPlannedMeta{
110+
descriptor,
111+
graph::GraphTensor(workspace),
112+
graph::GraphTensor(q),
113+
graph::GraphTensor(k),
114+
graph::GraphTensor(positions)};
115+
}
116+
117+
void run_vision(void *planned_meta) {
118+
auto *p = reinterpret_cast<VisionPlannedMeta *>(planned_meta);
119+
INFINICORE_CHECK_ERROR(
120+
infiniopErnie45VisionRope(
121+
p->descriptor->desc,
122+
p->workspace->data(),
123+
p->workspace->numel(),
124+
p->q->data(),
125+
p->k->data(),
126+
p->positions->data(),
127+
context::getStream()));
128+
}
129+
130+
void cleanup_vision(void **planned_meta_ptr) {
131+
delete *reinterpret_cast<VisionPlannedMeta **>(planned_meta_ptr);
132+
*planned_meta_ptr = nullptr;
133+
}
134+
135+
INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Ernie45VisionRoPE, &plan_vision, &run_vision, &cleanup_vision);
136+
137+
} // namespace vision_op
138+
139+
} // namespace infinicore::op::ernie45_rope_impl::infiniop

0 commit comments

Comments
 (0)