Skip to content

speculative: keep DFlash draft positions contiguous across mtmd vision chunks #5

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

Open
wangcc57 wants to merge 3 commits into
base: dflash2
Choose a base branch
from
Open

speculative: keep DFlash draft positions contiguous across mtmd vision chunks #5

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8396658
speculative: keep DFlash draft positions contiguous across mtmd
wangcc57 Aug 23, 2026
fbe93ca
style: format DFlash mtmd position fix
wangcc57 Aug 23, 2026
a789917
Fix formatting in speculative.cpp
wangcc57 Aug 23, 2026
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
15 changes: 13 additions & 2 deletions common/speculative.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1096,6 +1096,16 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

const int32_t n_ubatch = (int32_t) llama_n_ubatch(ctx_dft);

// The target context may advance over a vision chunk using positions that the
// independent DFlash context cannot reproduce exactly. Keep the draft in its
// own consecutive position space so a later decode always satisfies the
// one-axis RoPE invariant (Y = X + 1).
auto * mem_dft = llama_get_memory(ctx_dft);
std::vector<llama_pos> dnext(n_seq, 0);
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
dnext[seq_id] = llama_memory_seq_pos_max(mem_dft, seq_id) + 1;
}

// Flatten token-wise encoder work into shared chunks while preserving each row's position and sequence.
for (int32_t offset = 0; offset < n_tokens; offset += n_ubatch) {
const int32_t n_chunk = std::min(n_ubatch, n_tokens - offset);
Expand Down Expand Up @@ -1139,7 +1149,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
GGML_ASSERT(batch_in.n_seq_id[j] == 1);
const llama_seq_id seq_id = batch_in.seq_id[j][0];
GGML_ASSERT(seq_id >= 0 && seq_id < (llama_seq_id) n_seq);
batch_inject.pos[i] = batch_in.pos[j];
batch_inject.pos[i] = dnext[seq_id]++;
batch_inject.n_seq_id[i] = 1;
batch_inject.seq_id[i][0] = seq_id;
batch_inject.logits[i] = false;
Expand All @@ -1163,6 +1173,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

common_batch_clear(batch);

auto * mem_dft = llama_get_memory(ctx_dft);
// build one batch holding every drafting sequence's noise block into a single decode)
// record where each block starts and its size
std::vector<int32_t> i_block_beg(n_seq, -1);
Expand All @@ -1176,7 +1187,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

common_sampler_reset(smpls[seq_id].get());

const int32_t n = (int32_t) dp.n_past;
const int32_t n = llama_memory_seq_pos_max(mem_dft, seq_id) + 1;

const int32_t n_draft = params.n_max;

Expand Down