nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix dspark: seed draft block at id_last's true position (+1 off-by-one) (#2296) 

The draft block was seeded at last_target_pos (the newest committed
feature row = id_last's predecessor), placing the whole block one
position early vs mainline's [id_last @ n_past, ...] convention and
colliding the seed with the newest cross-KV row. Shift the common-side
batch and the graph-side SWA mask base coherently.
This commit is contained in:
rumas77 2026-08-12 12:21:30 -04:00 committed by GitHub
parent 1dede1d79e
commit c46ffaa566
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
common/speculative-dflash-impl.h
View File

@ -299,11 +299,13 @@ struct common_speculative_state_dflash : public common_speculative_state {
llama_kv_cache_clear(ctx_dft);
batch.n_tokens = 0;
const int32_t batch_len = is_dspark ? n_keep : n_keep + 1;
// id_last's true position is one past the newest committed feature row
// (last_target_pos): seed there, masks follow. Mirrors mainline's
// [id_last @ n_past, mask @ n_past+1, ...] block geometry.
const llama_pos draft_pos_base = last_target_pos >= 0 ? last_target_pos + 1 : (llama_pos) target_window_rows;
const llama_pos seed_pos = last_target_pos >= 0 ? last_target_pos : draft_pos_base - 1;
common_batch_add(batch, id_last, seed_pos, { 0 }, is_dspark);
common_batch_add(batch, id_last, draft_pos_base, { 0 }, is_dspark);
for (int32_t i = 1; i < batch_len; ++i) {
common_batch_add(batch, mask_token_id, draft_pos_base + (i - 1), { 0 }, true);
common_batch_add(batch, mask_token_id, draft_pos_base + i, { 0 }, true);
}
if (llama_decode(ctx_dft, batch) != 0) {

4
src/llama-dflash.cpp
View File

@ -671,7 +671,9 @@ bool llama_prepare_dflash_graph_inputs(
if (kq_mask_swa != nullptr) {
const int32_t swa_window = (int32_t) lctx.model.hparams.n_swa;
const int32_t draft_pos_base = (int32_t) last_target_pos;
// keep in sync with the draft batch geometry: block starts one past
// the newest committed feature row
const int32_t draft_pos_base = (int32_t) last_target_pos + 1;
if (kq_mask_swa->type == GGML_TYPE_F16) {
const ggml_fp16_t h_inf = ggml_fp32_to_fp16(-INFINITY);