Skip to content

fix(laguna) LagunaBackend ABI mismatch from KvFlashPager #477

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
davide221 wants to merge 1 commit into
base: codex/ddtree-ggml-graph-opt
Choose a base branch
from
+81 −13
Draft

fix(laguna) LagunaBackend ABI mismatch from KvFlashPager #477

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 41 additions & 12 deletions server/src/common/kvflash_pager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,10 +144,13 @@ class KvFlashPager {
has_pending_page_in_ = false;

#ifdef KVFLASH_HAS_ASYNC_DMA
if (cudaStreamCreate(&page_stream_) != cudaSuccess) {
cudaStream_t stream = nullptr;
if (cudaStreamCreate(&stream) != cudaSuccess) {
std::fprintf(stderr, "[kvflash] cudaStreamCreate failed; paging "
"falls back to the (blocking) default stream\n");
page_stream_ = nullptr;
} else {
page_stream_ = stream;
}
#endif
return true;
Expand All @@ -164,7 +167,7 @@ class KvFlashPager {
void reset() {
#ifdef KVFLASH_HAS_ASYNC_DMA
if (page_stream_) {
cudaStreamSynchronize(page_stream_);
cudaStreamSynchronize((cudaStream_t)page_stream_);
}
for (auto & st : chunks_) {
if (st.host_data) {
Expand Down Expand Up @@ -195,7 +198,7 @@ class KvFlashPager {
for (int b : free_blocks_) zero_block(b);
#ifdef KVFLASH_HAS_ASYNC_DMA
if (page_stream_) {
cudaStreamSynchronize(page_stream_);
cudaStreamSynchronize((cudaStream_t)page_stream_);
}
#endif
}
Expand Down Expand Up @@ -303,7 +306,7 @@ class KvFlashPager {
// tools) must call this first. No-op when async DMA is compiled out.
void synchronize_paging() {
#ifdef KVFLASH_HAS_ASYNC_DMA
cudaStreamSynchronize(page_stream_); // page_stream_==nullptr syncs the default stream
cudaStreamSynchronize((cudaStream_t)page_stream_); // nullptr syncs the default stream
#endif
has_pending_page_in_ = false;
}
Expand Down Expand Up @@ -367,6 +370,29 @@ class KvFlashPager {
int resident_blocks() const { return n_blocks_ - (int)free_blocks_.size(); }
int n_chunks() const { return (int)chunks_.size(); }

// Fully detach from backend-owned KV tensors. Use during backend shutdown
// before the cache/backend are freed; the destructor then becomes a no-op.
void detach() {
cleanup_();
attn_k_.clear();
attn_v_.clear();
chunks_.clear();
free_blocks_.clear();
zero_buf_.clear();
score_hook = nullptr;
cfg_ = {};
stats_ = {};
k_seg_bytes_ = 0;
v_seg_bytes_ = 0;
chunk_bytes_ = 0;
n_blocks_ = 0;
n_head_kv_ = 0;
cur_chunk_ = 0;
clock_ = 0;
epoch_ = 0;
has_pending_page_in_ = false;
}

// Bumped on every residency change (alloc / page_out / page_in).
// Callers cache the slot mask and refill only when the epoch moves.
uint64_t epoch() const { return epoch_; }
Expand Down Expand Up @@ -475,10 +501,10 @@ class KvFlashPager {
cudaError_t err;
if (to_host)
err = cudaMemcpyAsync(host_ptr, dev_ptr, seg,
cudaMemcpyDeviceToHost, page_stream_);
cudaMemcpyDeviceToHost, (cudaStream_t)page_stream_);
else
err = cudaMemcpyAsync(dev_ptr, host_ptr, seg,
cudaMemcpyHostToDevice, page_stream_);
cudaMemcpyHostToDevice, (cudaStream_t)page_stream_);
if (err != cudaSuccess) {
std::fprintf(stderr, "[kvflash] cudaMemcpyAsync(%s) failed: %s\n",
to_host ? "D2H" : "H2D", cudaGetErrorString(err));
Expand Down Expand Up @@ -517,7 +543,8 @@ class KvFlashPager {
for (int h = 0; h < n_head_kv_; h++) {
const size_t off = (size_t)block * cfg_.chunk_tokens * t->nb[1] + (size_t)h * t->nb[2];
#ifdef KVFLASH_HAS_ASYNC_DMA
cudaError_t err = cudaMemsetAsync((uint8_t *)t->data + off, 0, seg, page_stream_);
cudaError_t err = cudaMemsetAsync((uint8_t *)t->data + off, 0, seg,
(cudaStream_t)page_stream_);
if (err != cudaSuccess) {
std::fprintf(stderr, "[kvflash] cudaMemsetAsync failed: %s\n",
cudaGetErrorString(err));
Expand All @@ -536,9 +563,9 @@ class KvFlashPager {
// caller's responsibility (reset() or destructor via caller).
void cleanup_() {
#ifdef KVFLASH_HAS_ASYNC_DMA
cudaStreamSynchronize(page_stream_); // real stream, or default stream if create failed
cudaStreamSynchronize((cudaStream_t)page_stream_); // real stream, or default stream if create failed
if (page_stream_) {
cudaStreamDestroy(page_stream_);
cudaStreamDestroy((cudaStream_t)page_stream_);
page_stream_ = nullptr;
}
// Free pinned host buffers unconditionally: page_out() may have
Expand Down Expand Up @@ -572,9 +599,11 @@ class KvFlashPager {
uint64_t clock_ = 0;
uint64_t epoch_ = 0;

#ifdef KVFLASH_HAS_ASYNC_DMA
cudaStream_t page_stream_ = nullptr;
#endif
// Keep class layout independent of DFLASH27B_BACKEND_CUDA. LagunaBackend is
// instantiated by non-CUDA benchmark/server TUs but constructed in CUDA
// dflash_common TUs; a conditional member here changes sizeof(LagunaBackend)
// across TUs and corrupts the object at construction.
void * page_stream_ = nullptr;
bool has_pending_page_in_ = false;
};

Expand Down
9 changes: 9 additions & 0 deletions server/src/laguna/laguna_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3057,6 +3057,15 @@ void LagunaBackend::shutdown() {
drafter_loaded_ = false;
}
free_decode_draft();
kvflash_pager_.detach();
kvflash_scorer_.reset();
kvflash_scores_.clear();
stream_engine_.destroy();
moe_hybrid_.reset();
routing_stats_.reset();
layer_expert_bytes_.clear();
routing_stats_out_path_.clear();
kvflash_drafter_path_.clear();
if (!target_parked_) {
free_laguna_target_cache(cache_);
free_laguna_target_weights(w_);
Expand Down
24 changes: 23 additions & 1 deletion server/src/laguna/laguna_target_loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,7 +572,29 @@ bool load_target_gguf_laguna_partial(const std::string & path,
void free_laguna_target_weights(LagunaTargetWeights & w) {
if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; }
if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; }
// CpuEmbedder destructor handles mmap.
// Release the CPU embedding mmap now so repeated free/reload cycles and
// later member destruction cannot touch stale ownership fields.
#if defined(_WIN32)
if (w.embedder.mmap_addr) UnmapViewOfFile(w.embedder.mmap_addr);
if (w.embedder.mmap_hmap) CloseHandle(w.embedder.mmap_hmap);
if (w.embedder.mmap_hfile != INVALID_HANDLE_VALUE) CloseHandle(w.embedder.mmap_hfile);
w.embedder.mmap_addr = nullptr;
w.embedder.mmap_len = 0;
w.embedder.mmap_hfile = INVALID_HANDLE_VALUE;
w.embedder.mmap_hmap = nullptr;
#else
if (w.embedder.mmap_addr) ::munmap(w.embedder.mmap_addr, w.embedder.mmap_len);
if (w.embedder.mmap_fd >= 0) ::close(w.embedder.mmap_fd);
w.embedder.mmap_addr = nullptr;
w.embedder.mmap_len = 0;
w.embedder.mmap_fd = -1;
#endif
w.embedder.tok_embd_bytes = nullptr;
w.embedder.tok_embd_type = GGML_TYPE_COUNT;
w.embedder.n_embd = 0;
w.embedder.n_vocab = 0;
w.embedder.row_bytes = 0;
w.embedder.tok_embd_owned.clear();
w.layers.clear();
w.tok_embd = nullptr;
w.out_norm = nullptr;
Expand Down
8 changes: 8 additions & 0 deletions server/src/qwen35/gguf_target_loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,18 @@ CpuEmbedder::~CpuEmbedder() {
if (mmap_addr) UnmapViewOfFile(mmap_addr);
if (mmap_hmap) CloseHandle(mmap_hmap);
if (mmap_hfile != INVALID_HANDLE_VALUE) CloseHandle(mmap_hfile);
mmap_addr = nullptr;
mmap_len = 0;
mmap_hfile = INVALID_HANDLE_VALUE;
mmap_hmap = nullptr;
#else
if (mmap_addr) ::munmap(mmap_addr, mmap_len);
if (mmap_fd >= 0) ::close(mmap_fd);
mmap_addr = nullptr;
mmap_len = 0;
mmap_fd = -1;
#endif
tok_embd_bytes = nullptr;
}

bool CpuEmbedder::embed(const int32_t * ids, int n, float * out_f32) const {
Expand Down