EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add feature: priority

Browse Source 
Signed-off-by: Jing Wang <jingwang96@qq.com>
This commit is contained in:
Jing Wang
2026-05-12 11:51:57 +00:00
parent d627a45881
commit b6549b6e38
11 changed files with 382 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
csrc/camem_allocator.cpp
View File

@@ -595,10 +595,29 @@ static PyObject* python_try_lock_gpu_offload(PyObject* self, PyObject* args) {
}
static PyObject* python_unlock_gpu_offload(PyObject* self, PyObject* args) {
shm_worker->unlock_gpu();
int keep_wait = 0;
if (!PyArg_ParseTuple(args, "|p", &keep_wait)) {
return NULL;
}
shm_worker->unlock_gpu(keep_wait != 0);
Py_RETURN_NONE;
}
static PyObject* python_start_wait_offload(PyObject* self, PyObject* args) {
shm_worker->start_wait();
Py_RETURN_NONE;
}
static PyObject* python_cancel_wait_offload(PyObject* self, PyObject* args) {
shm_worker->cancel_wait();
Py_RETURN_NONE;
}
static PyObject* python_has_higher_priority_waiter_offload(PyObject* self, PyObject* args) {
bool has_higher = shm_worker->has_higher_priority_waiter();
return PyBool_FromLong(has_higher);
}
static PyMethodDef module_methods[] = {
{"init_module", (PyCFunction)py_init_module, METH_VARARGS,
"Initialize module with python_malloc and python_free callables."},
@@ -619,7 +638,13 @@ static PyMethodDef module_methods[] = {
{"python_try_lock_gpu_offload", (PyCFunction)python_try_lock_gpu_offload,
METH_NOARGS, "Lock GPU."},
{"python_unlock_gpu_offload", (PyCFunction)python_unlock_gpu_offload,
METH_NOARGS, "Unlock GPU."},
METH_VARARGS, "Unlock GPU."},
{"python_start_wait_offload", (PyCFunction)python_start_wait_offload,
METH_NOARGS, "Start waiting for GPU lock."},
{"python_cancel_wait_offload", (PyCFunction)python_cancel_wait_offload,
METH_NOARGS, "Cancel waiting for GPU lock."},
{"python_has_higher_priority_waiter_offload", (PyCFunction)python_has_higher_priority_waiter_offload,
METH_NOARGS, "Check if there is a higher priority waiter."},
{NULL, NULL, 0, NULL} // sentinel
};

39
csrc/vnpu_offload/shm_helper.h
View File

@@ -16,8 +16,8 @@
#include "spdlog/spdlog.h"
#define MAX_WORKERS 60
#define MAX_DEVICES 16
#define MAX_WORKERS 64
#define MAX_DEVICES 32
static inline std::string get_shm_name() {
const char *env_shm_name = getenv("VLLM_VNPU_SHM_NAME");
@@ -34,7 +34,7 @@ static inline std::string get_shm_name() {
}
static constexpr uint32_t heartbeat_us = 1000; // microseconds
static constexpr uint32_t heartbeat_check_everyN = 50;
static constexpr uint32_t heartbeat_check_everyN = 100;
static constexpr uint32_t heartbeat_timeout_us =
heartbeat_check_everyN * heartbeat_us;
@@ -52,6 +52,8 @@ static inline uint64_t heartbeat_ts_us() {
.count());
}
// GPU flag layout (64 bits):
// [lock (1 bit) | reserved (31 bits) | tgid (32 bits)]
static inline uint32_t unpack_lock_field(uint64_t gpu_flag) {
return static_cast<uint32_t>(gpu_flag >> 32);
}
@@ -68,16 +70,43 @@ static inline uint64_t pack_unlocked_tgid(int32_t tgid) {
return static_cast<uint64_t>(tgid);
}
// waiting_worker_flag layout (64 bits):
// [ device_id (5 bits) | priority (3 bits) | timestamp (24 bits) | tgid (32 bits)]
static inline uint32_t unpack_waiting_device_id(uint64_t flag) {
return static_cast<uint32_t>(flag >> 59);
}
static inline uint16_t unpack_waiting_priority(uint64_t flag) {
return static_cast<uint16_t>((flag >> 56) & 0x7);
}
static inline uint32_t unpack_waiting_timestamp_ms(uint64_t flag) {
return static_cast<uint32_t>((flag >> 32) & 0xFFFFFF);
}
static inline int32_t unpack_waiting_tgid(uint64_t flag) {
return static_cast<int32_t>(flag & 0xFFFFFFFF);
}
static inline uint64_t pack_waiting_flag(uint32_t device_id, uint16_t priority,
uint32_t timestamp, int32_t tgid) {
return (static_cast<uint64_t>(device_id & 0x1F) << 59) |
(static_cast<uint64_t>(priority & 0x7) << 56) |
(static_cast<uint64_t>(timestamp & 0xFFFFFF) << 32) |
(static_cast<uint64_t>(tgid) & 0xFFFFFFFF);
}
// mmap usually page-aligned
struct alignas(64) ShmHelper {
struct VramInfo {
uint64_t total_vmem_size;
uint64_t shareable_handle;
};
VramInfo vram_info[MAX_DEVICES]; // support max 16 NPUs
VramInfo vram_info[MAX_DEVICES]; // support max 32 devices
// GPU lock flag
std::atomic<uint64_t> gpu_flag[MAX_DEVICES];
// uint8_t _padding1[64 - sizeof(std::atomic<uint64_t>)];
std::atomic<uint64_t> waiting_worker_flags[MAX_WORKERS];
// request
enum RequestType: uint32_t {

4
csrc/vnpu_offload/shm_manager.cpp
View File

@@ -55,7 +55,7 @@ void ShmManager::run_busy_loop() {
while (!stop_loop_flag.load(std::memory_order_acquire)) {
process_requests();
if (loop_cnt % heartbeat_check_everyN== 0) {
if (loop_cnt % heartbeat_check_everyN == 0) {
check_heart_beats();
}
loop_cnt = (loop_cnt + 1) % heartbeat_check_everyN;
@@ -152,6 +152,8 @@ void ShmManager::check_heart_beats() {
shm_helper->heart_beats[i].tgid = 0;
shm_helper->heart_beats[i].timestamp.store(0,
std::memory_order_release);
// clear waiting flag
shm_helper->waiting_worker_flags[i].store(0, std::memory_order_release);
// check dead lock
for (int gpu_id : valid_gpu_ids) {
uint64_t gpu_flag =

114
csrc/vnpu_offload/shm_worker.cpp
View File

@@ -1,7 +1,29 @@
#include "shm_worker.h"
static inline uint16_t get_shm_priority() {
const char *env_priority = getenv("VLLM_VNPU_PRIORITY");
if (env_priority) {
try {
int p = std::stoi(env_priority);
if (p >= 0 && p <= 7) {
return static_cast<uint16_t>(p);
} else {
spdlog::warn("VLLM_VNPU_PRIORITY should be between 0 and 7, got {}. Using default 0.", p);
}
} catch (...) {
spdlog::warn("Invalid VLLM_VNPU_PRIORITY format. Using default 0.");
}
}
return 0;
}
ShmWorker::ShmWorker() {
this->priority = get_shm_priority();
this->waiting_timestamp = 0;
this->is_waiting = false;
this->is_holding_lock = false;
spdlog::info("vNPU worker initialized with priority {}", priority);
std::string shm_name = get_shm_name();
int shm_fd = shm_open(shm_name.c_str(), O_RDWR, 0666);
if (shm_fd == -1) {
@@ -40,16 +62,18 @@ bool ShmWorker::register_worker(int32_t tgid, int gpu_id,
if (slot == -1) {
return false;
}
this->shm_slot = slot;
*out_shareable_handle = shm_helper->vram_info[gpu_id].shareable_handle;
*out_vmem_size = shm_helper->vram_info[gpu_id].total_vmem_size;
stop_heart_beat.store(false, std::memory_order_release);
heart_beat_thread = std::thread(&ShmWorker::heart_beat_loop, this, slot);
heart_beat_thread = std::thread(&ShmWorker::heart_beat_loop, this);
return true;
}
void ShmWorker::heart_beat_loop(int slot) {
void ShmWorker::heart_beat_loop() {
int slot = this->shm_slot;
while (!stop_heart_beat.load(std::memory_order_acquire)) {
// update heart beat
int32_t shm_tgid =
@@ -64,6 +88,7 @@ void ShmWorker::heart_beat_loop(int slot) {
spdlog::error("TGID {} failed to re-register as worker", tgid);
throw std::runtime_error("Failed to re-register as worker");
}
this->shm_slot = slot;
}
uint64_t now = heartbeat_ts_us();
shm_helper->heart_beats[slot].timestamp.store(now,
@@ -72,32 +97,95 @@ void ShmWorker::heart_beat_loop(int slot) {
}
}
void ShmWorker::start_wait() {
if (is_waiting) return; // Keep the older timestamp if already waiting
// Use lower 24 bits of millisecond timestamp
waiting_timestamp = static_cast<uint32_t>((heartbeat_ts_us() / 1000) & 0xFFFFFF);
uint64_t flag = pack_waiting_flag(this->gpu_id, this->priority, waiting_timestamp, this->tgid);
shm_helper->waiting_worker_flags[this->shm_slot].store(flag, std::memory_order_release);
is_waiting = true;
}
void ShmWorker::cancel_wait() {
if (!is_waiting) return;
shm_helper->waiting_worker_flags[this->shm_slot].store(0, std::memory_order_release);
is_waiting = false;
}
bool ShmWorker::has_higher_priority_waiter() {
for (int i = 0; i < MAX_WORKERS; ++i) {
if (i == this->shm_slot) continue;
uint64_t flag = shm_helper->waiting_worker_flags[i].load(std::memory_order_acquire);
if (flag == 0) continue;
if (unpack_waiting_device_id(flag) != this->gpu_id) continue;
uint16_t other_prio = unpack_waiting_priority(flag);
if (other_prio > this->priority) {
return true; // Found a waiter with higher priority
} else if (other_prio == this->priority) {
if (this->is_holding_lock) {
// doesn't need to yield to same priority waiters
continue;
}
if (!this->is_waiting) {
// an earlier waiter with the same priority
return true;
}
uint32_t other_ts = unpack_waiting_timestamp_ms(flag);
// Same priority, compare timestamps (handle 24-bit wrap-around)
// Using 24-bit unsigned subtraction. If the difference is in the lower half,
// my timestamp is greater (i.e., I started waiting later).
uint32_t diff = (this->waiting_timestamp - other_ts) & 0xFFFFFF;
if (diff > 0 && diff < 0x800000) {
return true; // The other worker started waiting earlier
} else if (diff == 0 && unpack_waiting_tgid(flag) < this->tgid) {
// using tgid if timestamps happen to be exactly the same
return true;
}
}
}
return false;
}
bool ShmWorker::try_lock_gpu(bool &out_self_hold) {
static int retry_cnt = 0;
uint64_t old_flag =
shm_helper->gpu_flag[gpu_id].load(std::memory_order_acquire);
if (unpack_lock_field(old_flag) == 0) { // free
// Check priority: yield if there are higher priority waiters, or same priority waiters who have waited longer.
if (has_higher_priority_waiter()) {
out_self_hold = false;
return false;
}
uint64_t new_flag = pack_locked_tgid(tgid);
if (shm_helper->gpu_flag[gpu_id].compare_exchange_weak(
old_flag, new_flag, std::memory_order_acq_rel,
std::memory_order_acquire)) {
spdlog::info("TGID {} acquired GPU {} lock", tgid, gpu_id);
// spdlog::info("TGID {} acquired GPU {} lock", tgid, gpu_id);
int32_t prev_tgid = unpack_tgid_field(old_flag);
out_self_hold = prev_tgid == tgid;
retry_cnt = 0;
this->is_holding_lock = true;
return true;
}
} else { // locked
if (unpack_tgid_field(old_flag) == tgid) {
spdlog::info("TGID {} already holds the GPU {} lock", tgid, gpu_id);
// spdlog::info("TGID {} already holds the GPU {} lock", tgid, gpu_id);
out_self_hold = true;
retry_cnt = 0;
this->is_holding_lock = true;
return true;
}
}
// failed
if (++retry_cnt % 2000 == 0) {
if (++retry_cnt % 10000 == 0) {
spdlog::info(
"TGID {} trying to acquire GPU {} lock, current lock holder TGID {}",
tgid, gpu_id, unpack_tgid_field(old_flag));
@@ -116,19 +204,23 @@ bool ShmWorker::lock_gpu(bool &out_self_hold) {
}
}
void ShmWorker::unlock_gpu() {
void ShmWorker::unlock_gpu(bool keep_wait) {
if (!keep_wait) {
cancel_wait();
}
uint64_t old_flag =
shm_helper->gpu_flag[gpu_id].load(std::memory_order_acquire);
if (unpack_tgid_field(old_flag) != tgid) {
// spdlog::warn("previous gpu flag {} does not match expected locked flag for "
// "TGID {}. This may be a bug, unless during startup.",
// old_flag, tgid);
spdlog::info("TGID {} does not hold GPU {} lock", tgid, gpu_id);
if (!keep_wait) {
spdlog::info("unlock: TGID {} does not hold GPU {} lock", tgid, gpu_id);
}
} else {
uint64_t new_flag = pack_unlocked_tgid(tgid);
shm_helper->gpu_flag[gpu_id].store(new_flag, std::memory_order_release);
spdlog::info("TGID {} released GPU {} lock", tgid, gpu_id);
// spdlog::info("TGID {} released GPU {} lock", tgid, gpu_id);
}
this->is_holding_lock = false;
}
uint64_t ShmWorker::make_request(uint32_t type, uint64_t parameter) {

16
csrc/vnpu_offload/shm_worker.h
View File

@@ -17,11 +17,21 @@ class ShmWorker {
bool try_lock_gpu(bool &out_self_hold);
bool lock_gpu(bool &out_self_hold);
void unlock_gpu();
void unlock_gpu(bool keep_wait = false);
bool has_higher_priority_waiter();
void start_wait();
void cancel_wait();
private:
int32_t tgid;
int gpu_id;
int shm_slot;
uint16_t priority;
uint32_t waiting_timestamp;
bool is_waiting;
bool is_holding_lock;
ShmHelper *shm_helper;
std::thread heart_beat_thread;
std::atomic<bool> stop_heart_beat;
@@ -31,5 +41,5 @@ class ShmWorker {
int register_worker_shm();
// heart beat
void heart_beat_loop(int slot);
};
void heart_beat_loop();
};

43
csrc/vnpu_offload/vnpu_daemon.cpp
View File

@@ -1,25 +1,29 @@
#include <iostream>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <vector>
#include <atomic>
#include <fcntl.h>
#include <mutex>
#include <signal.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <vector>
#include "acl/acl.h"
#include "shm_manager.h"
#include "npu_helper.h"
#include "shm_manager.h"
#include "spdlog/spdlog.h"
static ShmManager *shm_manager = nullptr;
static inline double TO_GB(size_t bytes) {
return static_cast<double>(bytes) / (1024.0 * 1024.0 * 1024.0);
}
void handle_signal(int sig) {
if (shm_manager) {
shm_manager->stop_busy_loop();
@@ -49,7 +53,7 @@ size_t get_reserved_vram_size() {
reserved_vram_size = size_gb * 1024 * 1024 * 1024;
} catch (const std::exception &e) {
spdlog::warn("Failed to parse VNPU_RESERVED_VRAM_SIZE_GB: {}, using "
"default 8GB",
"default 8 GB",
e.what());
}
}
@@ -68,12 +72,13 @@ void ensure_context(unsigned long long device) {
}
void init_acl() {
int32_t deviceId=0;
int32_t deviceId = 0;
aclError ret = aclrtSetDevice(deviceId);
if (ret != ACL_ERROR_NONE) {
throw std::runtime_error("aclrtSetDevice failed with acl error code: " +
std::to_string(ret) + " " + __FILE__ + ":" + std::to_string(__LINE__));
throw std::runtime_error(
"aclrtSetDevice failed with acl error code: " + std::to_string(ret) +
" " + __FILE__ + ":" + std::to_string(__LINE__));
}
}
@@ -109,8 +114,9 @@ void alloc_physical(uint32_t device_id, aclrtDrvMemHandle &out_mem_handle,
spdlog::error("aclrtGetMemInfo failed, error_code: {}", error_code);
throw std::runtime_error("aclrtGetMemInfo failed");
} else {
spdlog::info("aclrtGetMemInfo succeeded, free_mem: {}, total: {}", free_mem,
total);
spdlog::info(
"aclrtGetMemInfo succeeded, free_mem: {:.2f} GB, total: {:.2f} GB",
TO_GB(free_mem), TO_GB(total));
}
aclrtPhysicalMemProp prop = {};
@@ -129,13 +135,14 @@ void alloc_physical(uint32_t device_id, aclrtDrvMemHandle &out_mem_handle,
error_code);
throw std::runtime_error("aclrtMemGetAllocationGranularity failed");
} else {
spdlog::info("aclrtMemGetAllocationGranularity succeeded, granularity: {}",
spdlog::info("aclrtMemGetAllocationGranularity succeeded, granularity: {} bytes",
granularity);
}
size_t reserved_mem_size = get_reserved_vram_size();
if (free_mem < reserved_mem_size) {
spdlog::error("Not enough free memory to reserve: {}, free_mem: {}",
reserved_mem_size, free_mem);
spdlog::error(
"Not enough free memory to reserve: {:.2f} GB, free_mem: {:.2f} GB",
TO_GB(reserved_mem_size), TO_GB(free_mem));
throw std::runtime_error("Not enough free memory to reserve");
}
out_g_size = free_mem - reserved_mem_size;
@@ -147,8 +154,8 @@ void alloc_physical(uint32_t device_id, aclrtDrvMemHandle &out_mem_handle,
spdlog::error("aclrtMallocPhysical failed, error_code: {}", error_code);
throw std::runtime_error("aclrtMallocPhysical failed");
} else {
spdlog::info("device {} aclrtMallocPhysical succeeded, size: {}", device_id,
out_g_size);
spdlog::info("device {} aclrtMallocPhysical succeeded, size: {:.2f} GB",
device_id, TO_GB(out_g_size));
}
}