Skip to content

Commit e0663be

Browse files
AnbeeldAnbeeld
committed
Improve DFlash multi-slot shared drafting
1 parent e536432 commit e0663be

8 files changed

Lines changed: 333 additions & 30 deletions

File tree

‎common/speculative.cpp‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4581,11 +4581,33 @@ void common_speculative_draft_batch(
45814581

45824582
llama_set_dflash_n_slots(ctx_dft, n_ready);
45834583

4584+
std::vector<llama_seq_id> ready_seq_ids;
4585+
ready_seq_ids.reserve(n_ready);
45844586
for (const auto & rs : ready) {
4587+
ready_seq_ids.push_back(rs.seq_id);
45854588
llama_memory_seq_rm(llama_get_memory(ctx_dft), rs.seq_id, rs.draft_pos_base, -1);
45864589
common_dflash_align_drafter_seq_or_clear(ctx_dft, rs.seq_id, rs.draft_pos_base, "batched draft");
45874590
}
45884591

4592+
if (n_ready > 1) {
4593+
const int ctx_window = params.dflash_cross_ctx > 0 ? params.dflash_cross_ctx : (int) LLAMA_DFLASH_PER_SLOT_CTX;
4594+
if (!llama_dflash_kv_cache_prepare_batch(ctx_dft, ready_seq_ids.data(), n_ready, ctx_window)) {
4595+
static bool warned_batch_kv = false;
4596+
if (!warned_batch_kv) {
4597+
LOG_WRN("dflash batch: batched K/V projection cache unavailable; falling back to per-slot cached drafting\n");
4598+
warned_batch_kv = true;
4599+
}
4600+
llama_set_dflash_n_slots(ctx_dft, 1);
4601+
return;
4602+
}
4603+
static bool logged_batch_kv = false;
4604+
if (!logged_batch_kv) {
4605+
LOG_INF("dflash: shared batched K/V projection cache enabled (%d slots, %d-token window)\n",
4606+
n_ready, ctx_window);
4607+
logged_batch_kv = true;
4608+
}
4609+
}
4610+
45894611
const int64_t t1 = ggml_time_us();
45904612

45914613
llama_batch batch = llama_batch_init(n_ready * batch_len, 0, 1);

‎docs/beellama-args.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -530,7 +530,7 @@ DFlash diagnostic environment variables:
530530
| `GGML_DFLASH_VERBOSE_CONTRACT=1` | Logs extra drafter/target contract details during DFlash setup. |
531531
| `GGML_DFLASH_FORCE_CPU_CROSS=1` | Force the CPU hidden-state cross path even when the GPU ring is available. |
532532
| `GGML_DFLASH_VERIFY_PAD=1` | Re-enable diagnostic verifier padding to the active draft depth. Default is off because padded rows consume target verify time but are not sampled or accepted. |
533-
| `GGML_DFLASH_SHARED_DRAFT_BATCH=1` | Opt into the experimental shared multi-slot DFlash drafter graph. The default keeps per-slot cached drafting because the shared graph pays the full multi-slot cross-attention window and disables the drafter K/V projection cache. |
533+
| `GGML_DFLASH_SHARED_DRAFT_BATCH=0` | Disable shared multi-slot DFlash drafter batching and fall back to per-slot cached drafting. By default, flat multi-slot DFlash uses a composite batched K/V projection cache when available. |
534534
| `GGML_DFLASH_GPU_RING=0` | Disable the GPU cross-attention ring and force the CPU ring path. |
535535
| `GGML_DFLASH_MULTI_GPU_TAPE=0` | Disable default-on multi-GPU DFlash GPU ring, hidden capture, tape, and replay. Use to force the CPU/eval-callback fallback for split target placement. |
536536
| `GGML_DFLASH_ALLOW_MULTI_GPU_TAPE=0` | Compatibility spelling for the same multi-GPU DFlash kill switch. |

‎docs/beellama-features.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ BeeLlama's practical advantage is combination. Public llama.cpp has the main run
4242
| Draft top-k for DFlash tree | No | No | `--draft-topk` | `--spec-draft-top-k` |
4343
| Adaptive DFlash depth | No | No | Speculative-code adaptive tracking | Server-side `profit` and `fringe` controllers with `--spec-dm-*` CLI |
4444
| Sampled DFlash CLI | No | No | No checked CLI surface for draft temperature | `--spec-draft-temp 0`, positive value, or `auto` |
45-
| Multi-slot DFlash | No | No | Shared DFlash slot machinery | Shared drafter context, slot cap, and per-slot cached drafting by default; experimental flat batched drafting behind `GGML_DFLASH_SHARED_DRAFT_BATCH=1` |
45+
| Multi-slot DFlash | No | No | Shared DFlash slot machinery | Shared drafter context, slot cap, and flat batched drafting with a composite K/V projection cache by default; `GGML_DFLASH_SHARED_DRAFT_BATCH=0` restores per-slot cached drafting |
4646
| Multimodal with speculation | Upstream rules | Upstream rules | Speculative decoding disabled under multimodal | Flat DFlash allowed; tree and non-DFlash spec disabled |
4747
| CopySpec | No checked match | No checked match | Yes | Yes as explicit `--spec-type copyspec` |
4848
| Suffix/recycle speculation | No checked match | No checked match | Yes | Yes |

‎include/llama.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1274,6 +1274,11 @@ extern "C" {
12741274
int ring_write_pos, int ring_filled,
12751275
int n_layers, int n_embd, int n_tokens,
12761276
llama_seq_id seq_id);
1277+
LLAMA_API bool llama_dflash_kv_cache_prepare_batch(
1278+
struct llama_context * ctx,
1279+
const llama_seq_id * seq_ids,
1280+
int n_seq,
1281+
int ctx_window);
12771282
LLAMA_API bool llama_dflash_target_kv_cache_update_from_ring(
12781283
struct llama_context * ctx, void * handle,
12791284
int ring_write_pos, int ring_filled,

‎src/llama-context.cpp‎

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4296,6 +4296,209 @@ bool llama_context::dflash_kv_cache_prepare(int ctx_window) {
42964296
return true;
42974297
}
42984298

4299+
bool llama_context::dflash_kv_cache_prepare_batch(const llama_seq_id * seq_ids, int n_seq, int ctx_window) {
4300+
cross.dflash_kv_cache = nullptr;
4301+
4302+
if (!seq_ids || n_seq < 2 || n_seq > (int) LLAMA_DFLASH_MAX_SLOTS ||
4303+
ctx_window <= 0 || !llm_arch_is_dflash_drafter(model.arch)) {
4304+
return false;
4305+
}
4306+
if (model.n_devices() > 1) {
4307+
return false;
4308+
}
4309+
4310+
ggml_backend_t gpu_backend = nullptr;
4311+
for (auto & backend : backends) {
4312+
auto * dev = ggml_backend_get_device(backend.get());
4313+
if (dev && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) {
4314+
gpu_backend = backend.get();
4315+
break;
4316+
}
4317+
}
4318+
if (!gpu_backend) {
4319+
return false;
4320+
}
4321+
4322+
struct batch_slot_cache {
4323+
dflash_kv_cache_data * cache = nullptr;
4324+
int n_copy = 0;
4325+
int src_offset = 0;
4326+
};
4327+
4328+
std::vector<batch_slot_cache> slot_caches;
4329+
slot_caches.reserve(n_seq);
4330+
4331+
int n_layers = 0;
4332+
int64_t n_embd_head = 0;
4333+
int64_t n_head_kv = 0;
4334+
int n_elem = 0;
4335+
4336+
for (int s = 0; s < n_seq; ++s) {
4337+
const llama_seq_id seq_id = seq_ids[s];
4338+
auto cache_it = dflash_kv_caches.find(seq_id);
4339+
if (cache_it == dflash_kv_caches.end() || !cache_it->second) {
4340+
return false;
4341+
}
4342+
auto * cache = cache_it->second.get();
4343+
4344+
auto cross_it = cross.v_embd_per_seq.find(seq_id);
4345+
if (cross_it == cross.v_embd_per_seq.end()) {
4346+
return false;
4347+
}
4348+
const int64_t cross_real = cross_it->second.v_embd_gpu
4349+
? cross_it->second.v_embd_gpu_n_enc_real
4350+
: cross_it->second.n_enc_real;
4351+
const int n_needed = (int) std::min<int64_t>(std::max<int64_t>(cross_real, 0), ctx_window);
4352+
4353+
if (n_needed <= 0 ||
4354+
cache->ring_size < ctx_window ||
4355+
cache->n_filled < n_needed ||
4356+
cache->write_pos != 0 ||
4357+
cache->n_elem <= 0 ||
4358+
!cache->fn_copy_d2d ||
4359+
cache->view.n_layers <= 0 ||
4360+
(int) cache->k_ring.size() < cache->view.n_layers ||
4361+
(int) cache->v_ring.size() < cache->view.n_layers) {
4362+
return false;
4363+
}
4364+
4365+
if (s == 0) {
4366+
n_layers = cache->view.n_layers;
4367+
n_embd_head = cache->view.n_embd_head;
4368+
n_head_kv = cache->view.n_head_kv;
4369+
n_elem = cache->n_elem;
4370+
} else if (n_layers != cache->view.n_layers ||
4371+
n_embd_head != cache->view.n_embd_head ||
4372+
n_head_kv != cache->view.n_head_kv ||
4373+
n_elem != cache->n_elem) {
4374+
return false;
4375+
}
4376+
4377+
for (int il = 0; il < n_layers; ++il) {
4378+
if (!cache->k_ring[il] || !cache->v_ring[il]) {
4379+
return false;
4380+
}
4381+
}
4382+
4383+
slot_caches.push_back({ cache, n_needed, cache->n_filled - n_needed });
4384+
}
4385+
4386+
if (n_layers <= 0 || n_elem <= 0 || n_embd_head <= 0 || n_head_kv <= 0) {
4387+
return false;
4388+
}
4389+
4390+
const int total_ctx = n_seq * ctx_window;
4391+
const bool reuse_batch =
4392+
dflash_kv_cache_batch.ctx != nullptr &&
4393+
dflash_kv_cache_batch.buf != nullptr &&
4394+
dflash_kv_cache_batch.n_slots == n_seq &&
4395+
dflash_kv_cache_batch.ctx_window == ctx_window &&
4396+
dflash_kv_cache_batch.n_elem == n_elem &&
4397+
dflash_kv_cache_batch.view.n_layers == n_layers &&
4398+
dflash_kv_cache_batch.view.n_embd_head == n_embd_head &&
4399+
dflash_kv_cache_batch.view.n_head_kv == n_head_kv &&
4400+
dflash_kv_cache_batch.view.ctx_len == total_ctx;
4401+
4402+
if (!reuse_batch) {
4403+
dflash_kv_cache_batch.reset();
4404+
4405+
const size_t ctx_mem = ggml_tensor_overhead() * ((size_t) n_layers * 2 + 8);
4406+
struct ggml_init_params ctx_params = { ctx_mem, nullptr, true };
4407+
ggml_context * batch_ctx = ggml_init(ctx_params);
4408+
if (!batch_ctx) {
4409+
return false;
4410+
}
4411+
4412+
dflash_kv_cache_batch.ctx = batch_ctx;
4413+
dflash_kv_cache_batch.n_slots = n_seq;
4414+
dflash_kv_cache_batch.ctx_window = ctx_window;
4415+
dflash_kv_cache_batch.n_elem = n_elem;
4416+
dflash_kv_cache_batch.view.n_layers = n_layers;
4417+
dflash_kv_cache_batch.view.n_embd_head = n_embd_head;
4418+
dflash_kv_cache_batch.view.n_head_kv = n_head_kv;
4419+
dflash_kv_cache_batch.view.ctx_len = total_ctx;
4420+
dflash_kv_cache_batch.view.ring_size = total_ctx;
4421+
dflash_kv_cache_batch.k_ring.resize(n_layers);
4422+
dflash_kv_cache_batch.v_ring.resize(n_layers);
4423+
dflash_kv_cache_batch.view.k_ring.resize(n_layers);
4424+
dflash_kv_cache_batch.view.v_ring.resize(n_layers);
4425+
4426+
for (int il = 0; il < n_layers; ++il) {
4427+
dflash_kv_cache_batch.k_ring[il] =
4428+
ggml_new_tensor_3d(batch_ctx, GGML_TYPE_F32, n_embd_head, n_head_kv, total_ctx);
4429+
dflash_kv_cache_batch.v_ring[il] =
4430+
ggml_new_tensor_3d(batch_ctx, GGML_TYPE_F32, n_embd_head, n_head_kv, total_ctx);
4431+
dflash_kv_cache_batch.view.k_ring[il] = dflash_kv_cache_batch.k_ring[il];
4432+
dflash_kv_cache_batch.view.v_ring[il] = dflash_kv_cache_batch.v_ring[il];
4433+
}
4434+
4435+
dflash_kv_cache_batch.buf = ggml_backend_alloc_ctx_tensors(batch_ctx, gpu_backend);
4436+
if (!dflash_kv_cache_batch.buf) {
4437+
dflash_kv_cache_batch.reset();
4438+
return false;
4439+
}
4440+
4441+
for (int il = 0; il < n_layers; ++il) {
4442+
ggml_backend_tensor_memset(dflash_kv_cache_batch.k_ring[il], 0, 0, ggml_nbytes(dflash_kv_cache_batch.k_ring[il]));
4443+
ggml_backend_tensor_memset(dflash_kv_cache_batch.v_ring[il], 0, 0, ggml_nbytes(dflash_kv_cache_batch.v_ring[il]));
4444+
}
4445+
ggml_backend_synchronize(gpu_backend);
4446+
4447+
const size_t total_size = ggml_backend_buffer_get_size(dflash_kv_cache_batch.buf);
4448+
LLAMA_LOG_INFO("%s: allocated DFlash batched K/V cache: %.1f MB (%d slots, %d tokens/slot, %d layers, %d elems/token)\n",
4449+
__func__, total_size / (1024.0 * 1024.0), n_seq, ctx_window, n_layers, n_elem);
4450+
}
4451+
4452+
const size_t stride = (size_t) n_elem * sizeof(float);
4453+
for (int s = 0; s < n_seq; ++s) {
4454+
auto * cache = slot_caches[s].cache;
4455+
const int n_copy = slot_caches[s].n_copy;
4456+
const int src_offset = slot_caches[s].src_offset;
4457+
const size_t bytes = (size_t) n_copy * stride;
4458+
const size_t dst_offset = (size_t) s * (size_t) ctx_window * stride;
4459+
4460+
for (int il = 0; il < n_layers; ++il) {
4461+
auto copy_tensor = [&](ggml_tensor * dst_tensor, const ggml_tensor * src_tensor) -> bool {
4462+
char * dst = (char *) dst_tensor->data + dst_offset;
4463+
const char * src = (const char *) src_tensor->data + (size_t) src_offset * stride;
4464+
const bool fast_copy =
4465+
cache->fn_copy_d2d_no_check &&
4466+
cache->fn_prepare_ptr &&
4467+
cache->fn_prepare_ptr(dst) &&
4468+
cache->fn_prepare_ptr(src);
4469+
const auto fn_copy = fast_copy ? cache->fn_copy_d2d_no_check : cache->fn_copy_d2d;
4470+
return fn_copy && fn_copy(dst, src, bytes);
4471+
};
4472+
4473+
if (!copy_tensor(dflash_kv_cache_batch.k_ring[il], cache->k_ring[il]) ||
4474+
!copy_tensor(dflash_kv_cache_batch.v_ring[il], cache->v_ring[il])) {
4475+
cross.dflash_kv_cache = nullptr;
4476+
return false;
4477+
}
4478+
}
4479+
}
4480+
4481+
auto * sync_cache = slot_caches.empty() ? nullptr : slot_caches[0].cache;
4482+
if (sync_cache && sync_cache->fn_wait_dflash_stream) {
4483+
if (!sync_cache->fn_wait_dflash_stream(gpu_backend)) {
4484+
cross.dflash_kv_cache = nullptr;
4485+
return false;
4486+
}
4487+
} else if (sync_cache && sync_cache->fn_sync_ptr && !dflash_kv_cache_batch.k_ring.empty()) {
4488+
if (!sync_cache->fn_sync_ptr(dflash_kv_cache_batch.k_ring[0]->data)) {
4489+
cross.dflash_kv_cache = nullptr;
4490+
return false;
4491+
}
4492+
} else {
4493+
ggml_backend_synchronize(gpu_backend);
4494+
}
4495+
4496+
dflash_kv_cache_batch.view.n_filled = total_ctx;
4497+
dflash_kv_cache_batch.view.write_pos = 0;
4498+
cross.dflash_kv_cache = &dflash_kv_cache_batch.view;
4499+
return true;
4500+
}
4501+
42994502
bool llama_context::dflash_kv_cache_update(int n_tokens) {
43004503
dflash_kv_cache_data * dflash_kv_cache = dflash_kv_cache_active();
43014504
if (!dflash_kv_cache || n_tokens <= 0 || !llm_arch_is_dflash_drafter(model.arch)) {
@@ -8755,6 +8958,15 @@ bool llama_dflash_kv_cache_update_from_ring_seq(
87558958
ctx, handle, ring_write_pos, ring_filled, n_layers, n_embd, n_tokens);
87568959
}
87578960

8961+
bool llama_dflash_kv_cache_prepare_batch(
8962+
llama_context * ctx,
8963+
const llama_seq_id * seq_ids,
8964+
int n_seq,
8965+
int ctx_window) {
8966+
if (!ctx) return false;
8967+
return ctx->dflash_kv_cache_prepare_batch(seq_ids, n_seq, ctx_window);
8968+
}
8969+
87588970
bool llama_dflash_target_kv_cache_update_from_ring(
87598971
llama_context * ctx, void * handle,
87608972
int ring_write_pos, int ring_filled,

‎src/llama-context.h‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,41 @@ struct dflash_kv_cache_data {
436436
}
437437
}
438438
};
439+
440+
struct dflash_kv_cache_batch_data {
441+
llama_dflash_kv_cache_view view;
442+
443+
ggml_context * ctx = nullptr;
444+
ggml_backend_buffer_t buf = nullptr;
445+
446+
std::vector<ggml_tensor *> k_ring;
447+
std::vector<ggml_tensor *> v_ring;
448+
449+
int n_slots = 0;
450+
int ctx_window = 0;
451+
int n_elem = 0;
452+
453+
void reset() {
454+
if (buf) {
455+
ggml_backend_buffer_free(buf);
456+
buf = nullptr;
457+
}
458+
if (ctx) {
459+
ggml_free(ctx);
460+
ctx = nullptr;
461+
}
462+
view = {};
463+
k_ring.clear();
464+
v_ring.clear();
465+
n_slots = 0;
466+
ctx_window = 0;
467+
n_elem = 0;
468+
}
469+
470+
~dflash_kv_cache_batch_data() {
471+
reset();
472+
}
473+
};
439474
struct llama_context {
440475
// init scheduler and compute buffers, reserve worst-case graphs
441476
llama_context(
@@ -740,6 +775,7 @@ struct llama_context {
740775
bool dflash_kv_cache_update_gpu(const void * d_hidden, int n_tokens, int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d);
741776
bool dflash_target_kv_cache_update_gpu(llama_seq_id seq_id, llama_pos start_pos, const void * d_hidden, int n_tokens, int n_layers, int n_embd_layer, set_tensor_d2d_fn_t fn_d2d);
742777
bool dflash_kv_cache_prepare(int ctx_window);
778+
bool dflash_kv_cache_prepare_batch(const llama_seq_id * seq_ids, int n_seq, int ctx_window);
743779
dflash_kv_cache_data * dflash_kv_cache_active();
744780
void dflash_kv_cache_set_active_seq(llama_seq_id seq_id);
745781
std::unique_ptr<dflash_kv_cache_data> & dflash_kv_cache_active_ref();
@@ -920,6 +956,7 @@ struct llama_context {
920956
bool dflash_kv_cache_multi_gpu_fallback_logged = false;
921957
llama_seq_id dflash_kv_cache_active_seq = -1;
922958
std::map<llama_seq_id, std::unique_ptr<dflash_kv_cache_data>> dflash_kv_caches;
959+
dflash_kv_cache_batch_data dflash_kv_cache_batch;
923960

924961
bool dflash_capture_valid_last_decode = true;
925962
std::string dflash_capture_invalid_reason;

0 commit comments

Comments
 (0)