Skip to content

Commit 5724f1a

Browse files
kmbandyclaude
kmbandy
and
claude
committed
fix(merge): restore device-lifecycle refcounting + fix speculative-decode compile breaks
Upstream sync (305 commits) merged cleanly at the text level in several spots but broke semantically: - ggml-cuda.cu: the active_count/device_mutex per-device refcounting (used by ggml_backend_cuda_device_get_memory to safely cudaDeviceReset an idle device instead of leaking a lazily-created CUDA context) was silently dropped across all 7 call sites during the merge, since each drop was a non-conflicting textual change. Restored via two forward- declared helpers (ggml_backend_cuda_device_active_count_inc/dec) since most call sites precede the struct's new (upstream-relocated) definition point and can't do direct member access there. - server-context.cpp: fixed compile errors introduced while porting our MAD-120 paged-KV admission control onto upstream's new iterate()/ handle_last_sampled_token() API: - continue; inside the iterate() lambda (invalid — lambdas aren't loops) changed to return;, matching the surrounding skip-slot convention already used there. - stale batch.n_tokens field reference (type is now a class exposing .size()) in a log line. - a broken speculative-decode telemetry block referencing three undefined variables (id, tok_idx, t_current) copy-pasted from a different function's scope — fixed to use the correct in-scope equivalents (ids[i], t_now) and dropped the populate_token_probs call that had no valid index in this scope (the existing "TODO: set result.probs" gap is left as-is, not newly introduced). Verified: build-army (CUDA) llama-server + test-backend-ops clean, 1134/1134 MUL_MAT pass on CUDA0. build-vk (Vulkan+CUDA) test-backend-ops clean, 956/956 MUL_MAT pass on Vulkan0. test-paged-attn-vk (SP2 dual- backend oracle) all cases pass. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014ZRfPpL8XFzk1hep9MMg9P
1 parent 9d916fe commit 5724f1a

2 files changed

Lines changed: 68 additions & 18 deletions

File tree

‎ggml/src/ggml-cuda/ggml-cuda.cu‎

Lines changed: 62 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -665,6 +665,17 @@ struct ggml_backend_cuda_buffer_context {
665665
}
666666
};
667667

668+
// ggml_backend_cuda_device_context (defined later in this file, alongside the
669+
// other device-management code) tracks how many buffers/backends are active
670+
// per device so ggml_backend_cuda_device_get_memory can safely cudaDeviceReset
671+
// a device nothing is using instead of leaking a lazily-created CUDA context.
672+
// These sites run before that struct's definition, so use forward-declared
673+
// helpers rather than accessing dev_ctx members directly here.
674+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
675+
static void ggml_backend_cuda_device_active_count_inc(ggml_backend_dev_t dev);
676+
static void ggml_backend_cuda_device_active_count_dec(ggml_backend_dev_t dev);
677+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
678+
668679
static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) {
669680
ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
670681
// The ml8 / ml8-fp8 weight-repack caches key on a weight's device pointer.
@@ -675,9 +686,7 @@ static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) {
675686
ggml_cuda_ml8_clear_cache();
676687

677688
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
678-
ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) buffer->buft->device->context;
679-
std::lock_guard<std::mutex> lock(dev_ctx->device_mutex);
680-
dev_ctx->active_count--;
689+
ggml_backend_cuda_device_active_count_dec(buffer->buft->device);
681690
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
682691

683692
delete ctx;
@@ -849,6 +858,10 @@ static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_bac
849858

850859
ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr);
851860

861+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
862+
ggml_backend_cuda_device_active_count_inc(buft->device);
863+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
864+
852865
return ggml_backend_buffer_init(buft, ggml_backend_cuda_buffer_interface, ctx, size);
853866
}
854867

@@ -1571,6 +1584,10 @@ static bool ggml_backend_buft_is_cuda_host(ggml_backend_buffer_type_t buft) {
15711584
}
15721585

15731586
static void ggml_backend_cuda_host_buffer_free_buffer(ggml_backend_buffer_t buffer) {
1587+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
1588+
ggml_backend_cuda_device_active_count_dec(buffer->buft->device);
1589+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
1590+
15741591
CUDA_CHECK(cudaFreeHost(buffer->context));
15751592
}
15761593

@@ -1604,6 +1621,10 @@ static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggm
16041621
buffer->buft = buft;
16051622
buffer->iface.free_buffer = ggml_backend_cuda_host_buffer_free_buffer;
16061623

1624+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
1625+
ggml_backend_cuda_device_active_count_inc(buft->device);
1626+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
1627+
16071628
return buffer;
16081629
}
16091630

@@ -3475,6 +3496,10 @@ static const char * ggml_backend_cuda_get_name(ggml_backend_t backend) {
34753496
static void ggml_backend_cuda_free(ggml_backend_t backend) {
34763497
ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
34773498

3499+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
3500+
ggml_backend_cuda_device_active_count_dec(backend->device);
3501+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
3502+
34783503
delete cuda_ctx;
34793504
delete backend;
34803505
}
@@ -5340,8 +5365,26 @@ struct ggml_backend_cuda_device_context {
53405365
std::string description;
53415366
std::string pci_bus_id;
53425367
int op_offload_min_batch_size;
5368+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5369+
std::mutex device_mutex;
5370+
int active_count = 0;
5371+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
53435372
};
53445373

5374+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5375+
static void ggml_backend_cuda_device_active_count_inc(ggml_backend_dev_t dev) {
5376+
ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context;
5377+
std::lock_guard<std::mutex> lock(dev_ctx->device_mutex);
5378+
dev_ctx->active_count++;
5379+
}
5380+
5381+
static void ggml_backend_cuda_device_active_count_dec(ggml_backend_dev_t dev) {
5382+
ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context;
5383+
std::lock_guard<std::mutex> lock(dev_ctx->device_mutex);
5384+
dev_ctx->active_count--;
5385+
}
5386+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5387+
53455388
static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) {
53465389
ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context;
53475390
return ctx->name.c_str();
@@ -5430,6 +5473,11 @@ static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_k
54305473

54315474
static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
54325475
ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context;
5476+
5477+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5478+
std::lock_guard<std::mutex> lock(ctx->device_mutex);
5479+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5480+
54335481
ggml_cuda_set_device(ctx->device);
54345482
CUDA_CHECK(cudaMemGetInfo(free, total));
54355483

@@ -5456,6 +5504,13 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t *
54565504
}
54575505
#endif // defined(__linux__)
54585506

5507+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5508+
// If no backends or buffers are active, the cudaMemGetInfo call above lazily created a CUDA
5509+
// context that permanently consumes VRAM. Reset the device to free it.
5510+
if (ctx->active_count == 0) {
5511+
CUDA_CHECK(cudaDeviceReset());
5512+
}
5513+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
54595514
}
54605515

54615516
static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) {
@@ -6336,6 +6391,10 @@ ggml_backend_t ggml_backend_cuda_init(int device) {
63366391
/* .context = */ ctx,
63376392
};
63386393

6394+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
6395+
ggml_backend_cuda_device_active_count_inc(cuda_backend->device);
6396+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
6397+
63396398
return cuda_backend;
63406399
}
63416400

‎tools/server/server-context.cpp‎

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3564,10 +3564,10 @@ struct server_context_impl {
35643564
// abort doesn't fire on this same iteration just
35653565
// because we haven't built a token batch yet.
35663566
n_empty_consecutive = 0;
3567-
continue;
3567+
return;
35683568
}
35693569
}
3570-
continue;
3570+
return;
35713571
}
35723572
// Slot was admitted — any prior no-progress streak ends here.
35733573
slot.paged_preempt_no_progress_count = 0;
@@ -4081,7 +4081,7 @@ struct server_context_impl {
40814081
slot.i_batch = batch.size() - 1;
40824082

40834083
slot.init_sampler();
4084-
SLT_INF(slot, "prompt processing done, n_tokens = %d, batch.n_tokens = %d\n", slot.prompt.n_tokens(), batch.n_tokens);
4084+
SLT_INF(slot, "prompt processing done, n_tokens = %d, batch.n_tokens = %d\n", slot.prompt.n_tokens(), batch.size());
40854085

40864086
// MAD-129: prefill-time semantic fingerprint trigger.
40874087
// Walk the seq's COMPLETE blocks (skip the partial last
@@ -4531,12 +4531,12 @@ struct server_context_impl {
45314531
slot.n_decoded += 1;
45324532

45334533
if (slot.n_decoded == 1) {
4534-
slot.t_start_generation = t_current;
4534+
slot.t_start_generation = t_now;
45354535
slot.t_prompt_processing = (slot.t_start_generation - slot.t_start_process_prompt) / 1e3;
45364536
metrics.on_prompt_eval(slot);
45374537
}
45384538

4539-
slot.t_token_generation = std::max<int64_t>(1, t_current - slot.t_start_generation) / 1e3;
4539+
slot.t_token_generation = std::max<int64_t>(1, t_now - slot.t_start_generation) / 1e3;
45404540

45414541
if (slot.n_decoded % 64 == 0 || slot.n_decoded == 1) {
45424542
const float tps = slot.t_token_generation > 0
@@ -4545,16 +4545,7 @@ struct server_context_impl {
45454545
const int32_t n_think = common_reasoning_budget_get_n_thinking(rbudget);
45464546
const auto rb_st = common_reasoning_budget_get_state(rbudget);
45474547
SLT_INF(slot, "generate: n_decoded=%d n_past=%d t/s=%.2f rbudget_state=%d n_thinking=%d token_id=%d\n",
4548-
slot.n_decoded, slot.prompt.n_tokens(), tps, (int)rb_st, n_think, (int)id);
4549-
}
4550-
4551-
completion_token_output result;
4552-
result.tok = id;
4553-
result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok));
4554-
result.prob = 1.0f; // TODO: set it here instead of doing inside populate_token_probs
4555-
4556-
if (slot.task->params.sampling.n_probs > 0) {
4557-
populate_token_probs(slot, result, slot.task->params.post_sampling_probs, params_base.special, tok_idx);
4548+
slot.n_decoded, slot.prompt.n_tokens(), tps, (int)rb_st, n_think, (int)ids[i]);
45584549
}
45594550

45604551
if (!process_token(result, slot)) {

0 commit comments

Comments
 (0)