add vxpu
This commit is contained in:
36
vllm_kunlun/csrc/vxpu_offload/shm_worker.h
Normal file
36
vllm_kunlun/csrc/vxpu_offload/shm_worker.h
Normal file
@@ -0,0 +1,36 @@
|
||||
#pragma once
|
||||
|
||||
#include <vector>
|
||||
#include <atomic>
|
||||
#include <thread>
|
||||
|
||||
#include "xpu_helper.h"
|
||||
#include "shm_helper.h"
|
||||
|
||||
|
||||
class ShmWorker {
|
||||
public:
|
||||
ShmWorker();
|
||||
~ShmWorker();
|
||||
|
||||
bool register_worker(int device_id, XPUIpcMemHandle *out_shareable_handle,
|
||||
uint64_t *out_vmem_size);
|
||||
|
||||
bool try_lock_gpu(bool &out_self_hold);
|
||||
bool lock_gpu(bool &out_self_hold);
|
||||
void unlock_gpu();
|
||||
|
||||
private:
|
||||
int32_t worker_id;
|
||||
int gpu_slot;
|
||||
ShmHelper *shm_helper;
|
||||
std::thread heart_beat_thread;
|
||||
std::atomic<bool> stop_heart_beat;
|
||||
|
||||
// request
|
||||
uint64_t make_request(uint32_t type, uint64_t parameter);
|
||||
int register_worker_shm();
|
||||
|
||||
// heart beat
|
||||
void heart_beat_loop(int slot);
|
||||
};
|
||||
Reference in New Issue
Block a user