Skip to content

perf(qwen35): retain speculation during paged prefill #686

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
Graffioh wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

perf(qwen35): retain speculation during paged prefill #686

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
8 changes: 4 additions & 4 deletions server/src/common/step_graph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ namespace dflash::common {

using TargetPagedTreeGraphKey = std::tuple<
const TargetWeights *, const TargetCache *, ggml_backend_t,
int, int, int, int, int, int>;
int, int, int, int, int, int, int, std::vector<int>>;

struct StepGraph {
ggml_context * ctx = nullptr;
Expand Down Expand Up @@ -81,14 +81,14 @@ struct StepGraph {
// inclusive causal position, [n_tokens] i32 each. Padding rows carry -1.
ggml_tensor * paged_query_seq_ids = nullptr;
ggml_tensor * paged_query_positions = nullptr;
// DFlash target-feature destination rows. Multi-slot replay maps each
// token to its slot-local ring; padding maps to the cache's dead row.
// DFlash target-feature destination rows. In packed-tree graphs this maps
// only the durable prefill/AR prefix; tree features remain scratch.
ggml_tensor * target_feat_rows = nullptr;
// Packed-tree direct-commit metadata uploaded after posterior selection.
ggml_tensor * accepted_prefixes = nullptr; // [n_tree_seqs] i32
ggml_tensor * commit_slot_ids = nullptr; // [n_tree_seqs] i32
ggml_tensor * commit_rows = nullptr; // [tree_width,n_tree_seqs] i64
ggml_tensor * feature_commit_rows = nullptr; // [n_tokens] i32
ggml_tensor * feature_commit_rows = nullptr; // [tree_width*n_tree_seqs] i32
// Multi-prompt steps: i32 row indices gathered from the final norm
// before the LM head (committing rows + decode rows).
ggml_tensor * logits_row_indices = nullptr;
Expand Down
20 changes: 11 additions & 9 deletions server/src/internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -840,12 +840,12 @@ struct QwenGraphInputs {
// for reading the much smaller conv-state slabs; writes use active ids.
ggml_tensor * state_slot_ids = nullptr;
// [n_tokens] i32 per-row block-table column for the ragged paged
// attention read (prefill rows carry their slot, decode rows theirs,
// padding -1). Non-null exactly when n_prefill_tokens > 0.
// attention read (prefill rows carry their slot, direct decode/tree rows
// theirs, padding -1). A packed tree also uses this map for every row.
ggml_tensor * paged_query_seq_ids = nullptr;
// [n_tokens] i32 per-row inclusive logical position; the kernel clamps
// each row's KV extent to position+1, which IS the causal mask. -1 on
// padding rows.
// each direct row's KV extent to position+1, which IS the causal mask.
// Packed-tree and padding rows carry -1.
ggml_tensor * paged_query_positions = nullptr;
// Optional [n_rows] i32 gather of final-norm rows before the LM head:
// multi-prompt steps sample scattered rows (each committing segment's
Expand All @@ -854,8 +854,9 @@ struct QwenGraphInputs {
ggml_tensor * logits_row_indices = nullptr;
// Optional replay-stable DFlash capture destinations. When present, all
// captured layers are concatenated once and written with ggml_set_rows.
// Multi-slot callers provide per-slot ring rows (padding uses dead row).
ggml_tensor * target_feat_rows = nullptr; // [n_tokens] i32
// In a packed-tree graph this covers only the durable prefill/AR prefix;
// the fixed-width tree suffix is returned separately for promotion.
ggml_tensor * target_feat_rows = nullptr; // [n_direct_tokens] i32
// Prefill segments on the leading token axis (see QwenPrefillSegment).
// n_prefill_tokens is their total row count. seq_slot is ignored when
// segments are present.
Expand Down Expand Up @@ -888,8 +889,9 @@ struct QwenGraphInputs {
// compact decode rows; logits_tail_rows remains the dense-path fallback.
int n_seqs = 1;
// Mixed direct-commit tree graphs place this many one-token mapped AR
// sequences before the fixed-width speculative tree segment. Their slot
// IDs share active_slot_ids/state_slot_ids with the tree lanes.
// sequences after any prefill segments and before the fixed-width
// speculative tree suffix. Their slot IDs share active_slot_ids/
// state_slot_ids with the tree lanes; prefills own direct slab views.
int mapped_ar_seqs = 0;
int seq_slot = 0;
int paged_max_kv_len = 0;
Expand Down Expand Up @@ -926,7 +928,7 @@ struct QwenGraphOutputs {
// views marked as ggml_set_output() so their data persists after
// graph_compute; the spec-decode loop reads them host-side for rollback.
std::vector<DeltaNetCapture> delta_captures;
// BF16 [n_capture_layers*n_embd, n_tokens], packed-tree only.
// BF16 [n_capture_layers*n_embd, tree_width*n_tree_seqs], packed-tree only.
ggml_tensor * tree_features = nullptr;
// One entry per target layer. Populated only when capture_moe_router is
// true; qwen35 dense layers and non-MoE models leave entries null.
Expand Down
185 changes: 155 additions & 30 deletions server/src/qwen35/concurrency/qwen35_seq_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,6 @@ bool Qwen35SeqEngine::chain_spec_input_capable(
std::vector<uint8_t>
Qwen35SeqEngine::select_chain_lanes(const StepPlan & plan) const {
std::vector<uint8_t> selected(plan.decode.size(), 0);
if (!plan.prefills.empty()) return selected;
for (size_t i = 0; i < plan.decode.size(); ++i) {
selected[i] = chain_spec_input_capable(plan.decode[i]) ? 1 : 0;
}
Expand Down Expand Up @@ -480,7 +479,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
PreparedChainRound && prepared_round) {
StepResult result;
const std::vector<StepInput> & inputs = plan.decode;
if (!plan.prefills.empty() || selected.size() != inputs.size()) {
if (selected.size() != inputs.size()) {
result.error = "invalid fixed chain round";
return result;
}
Expand Down Expand Up @@ -553,6 +552,39 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
return result;
}

std::vector<PrefillStage> prefills;
std::vector<QwenPrefillSegment> prefill_segments;
prefills.reserve(plan.prefills.size());
prefill_segments.reserve(plan.prefills.size());
result.prefills.reserve(plan.prefills.size());
int n_prefill = 0;
int n_prefill_commits = 0;
int max_prefix = 1;
for (const PrefillSlice & slice : plan.prefills) {
const size_t outputs_before = result.prefills.size();
PrefillStage prefill =
stage_prefill_chunk(slice.slot, slice.max_tokens,
result.prefills);
if (!prefill.ready) {
if (result.prefills.size() == outputs_before) {
fail_prefill(
slice.slot, result.prefills,
"prefill made no progress during fixed chain round",
"prefill scheduler made no progress");
}
result.prefills.clear();
result.error = "selected prefill work made no progress";
return result;
}
prefill_segments.push_back(
{n_prefill, prefill.chunk, slice.slot});
n_prefill += prefill.chunk;
n_prefill_commits += prefill.commit ? 1 : 0;
max_prefix = std::max(
max_prefix, prefill.kv_pos + prefill.chunk);
prefills.push_back(std::move(prefill));
}

struct StagedRound {
Qwen35SlotManager & slots;
std::vector<int32_t> & seq_lens;
Expand Down Expand Up @@ -623,9 +655,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
const int ar_count = static_cast<int>(ar_lanes.size());
const int tree_bucket = chain_decode_bucket_width(spec_count);
const int tree_rows_count = tree_width * tree_bucket;
const int total_rows = ar_count + tree_rows_count;
const int direct_row0 = n_prefill;
const int tree_row0 = direct_row0 + ar_count;
const int total_rows = tree_row0 + tree_rows_count;
const int posterior_ar_row0 = n_prefill_commits;
const int posterior_tree_row0 = posterior_ar_row0 + ar_count;
const int logits_rows_count = n_prefill > 0
? posterior_tree_row0 + tree_rows_count
: 0;
const int posterior_rows_count =
logits_rows_count > 0 ? logits_rows_count : total_rows;

int max_prefix = 1;
for (const Proposal & proposal : proposals) {
max_prefix = std::max(
max_prefix, slots_.slot(proposal.slot).cur_pos);
Expand All @@ -639,10 +679,19 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
graph, b_.w_, b_.cache_, b_.target_backend_,
tree_width, tree_bucket, max_prefix,
fixed_chain_.scratch_base, fixed_chain_.scratch_stride,
b_.cfg_.kq_stride_pad, ar_count)) {
b_.cfg_.kq_stride_pad, ar_count, n_prefill,
prefill_segments.data(),
static_cast<int>(prefill_segments.size()),
logits_rows_count)) {
result.error = "fixed chain target graph build failed";
return result;
}
if ((logits_rows_count > 0 && !graph.logits_row_indices) ||
(tree_row0 > 0 &&
(!graph.paged_query_positions || !graph.target_feat_rows))) {
result.error = "fixed chain mixed graph inputs are incomplete";
return result;
}

std::vector<int32_t> tokens(static_cast<size_t>(total_rows), 0);
std::vector<int32_t> parents(
Expand All @@ -663,22 +712,61 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
static_cast<size_t>(4) * total_rows, 0);
std::vector<float> embeddings(
static_cast<size_t>(hidden) * total_rows, 0.0f);
const int feature_cap = b_.cache_.target_feat_cap;
std::vector<int32_t> direct_feature_rows(
static_cast<size_t>(tree_row0), -1);
std::vector<int32_t> logits_rows;
logits_rows.reserve(static_cast<size_t>(logits_rows_count));
seq_lens_.assign(static_cast<size_t>(n_slots), 0);

int prefill_row = 0;
for (size_t i = 0; i < prefills.size(); ++i) {
const PrefillStage & prefill = prefills[i];
const int slot = plan.prefills[i].slot;
std::copy(
prefill.embeddings.begin(), prefill.embeddings.end(),
embeddings.begin() + static_cast<size_t>(hidden) * prefill_row);
fill_qwen35_mrope_positions(
positions.data(), total_rows, prefill_row,
prefill.kv_pos, prefill.chunk);
for (int row = 0; row < prefill.chunk; ++row) {
const int packed_row = prefill_row + row;
query_slots[static_cast<size_t>(packed_row)] = slot;
query_positions[static_cast<size_t>(packed_row)] =
prefill.kv_pos + row;
direct_feature_rows[static_cast<size_t>(packed_row)] =
slot * feature_cap +
(prefill.kv_pos + row) % feature_cap;
for (int head = 0; head < n_head_kv; ++head) {
write_rows[static_cast<size_t>(head) * total_rows +
packed_row] = prefill.rows[static_cast<size_t>(row)];
}
}
if (prefill.commit) {
logits_rows.push_back(prefill_row + prefill.chunk - 1);
}
seq_lens_[static_cast<size_t>(slot)] =
prefill.kv_pos + prefill.chunk;
prefill_row += prefill.chunk;
}

for (int lane_index = 0; lane_index < ar_count; ++lane_index) {
const ArLane & lane = ar_lanes[static_cast<size_t>(lane_index)];
const int row = direct_row0 + lane_index;
mapped_slots[static_cast<size_t>(lane_index)] = lane.slot;
state_slots[static_cast<size_t>(lane_index)] = lane.slot;
tokens[static_cast<size_t>(lane_index)] = lane.token;
query_slots[static_cast<size_t>(lane_index)] = lane.slot;
query_positions[static_cast<size_t>(lane_index)] = lane.position;
tokens[static_cast<size_t>(row)] = lane.token;
query_slots[static_cast<size_t>(row)] = lane.slot;
query_positions[static_cast<size_t>(row)] = lane.position;
direct_feature_rows[static_cast<size_t>(row)] =
lane.slot * feature_cap + lane.position % feature_cap;
seq_lens_[static_cast<size_t>(lane.slot)] = lane.position + 1;
for (int axis = 0; axis < 3; ++axis) {
positions[static_cast<size_t>(axis) * total_rows + lane_index] =
positions[static_cast<size_t>(axis) * total_rows + row] =
lane.position;
}
for (int head = 0; head < n_head_kv; ++head) {
write_rows[static_cast<size_t>(head) * total_rows + lane_index] =
write_rows[static_cast<size_t>(head) * total_rows + row] =
lane.physical_row;
}
}
Expand All @@ -687,7 +775,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
const Proposal & proposal =
proposals[static_cast<size_t>(lane_index)];
const int tree_base = lane_index * tree_width;
const int row_base = ar_count + tree_base;
const int row_base = tree_row0 + tree_base;
const int mapped_lane = ar_count + lane_index;
tree_sizes[static_cast<size_t>(lane_index)] = tree_width;
mapped_slots[static_cast<size_t>(mapped_lane)] = proposal.slot;
Expand Down Expand Up @@ -717,8 +805,20 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
}
}

if (!b_.w_.embedder.embed(
tokens.data(), total_rows, embeddings.data())) {
if (logits_rows_count > 0) {
for (int row = direct_row0; row < total_rows; ++row) {
logits_rows.push_back(row);
}
if (static_cast<int>(logits_rows.size()) != logits_rows_count) {
result.error = "fixed chain logits layout is invalid";
return result;
}
}

const int embedded_rows = total_rows - n_prefill;
if (embedded_rows > 0 && !b_.w_.embedder.embed(
tokens.data() + n_prefill, embedded_rows,
embeddings.data() + static_cast<size_t>(hidden) * n_prefill)) {
result.error = "fixed chain embedding failed";
return result;
}
Expand Down Expand Up @@ -750,6 +850,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
graph.paged_query_positions, query_positions.data(), 0,
sizeof(int32_t) * query_positions.size());
}
if (graph.logits_row_indices) {
ggml_backend_tensor_set(
graph.logits_row_indices, logits_rows.data(), 0,
sizeof(int32_t) * logits_rows.size());
}
if (graph.target_feat_rows) {
ggml_backend_tensor_set(
graph.target_feat_rows, direct_feature_rows.data(), 0,
sizeof(int32_t) * direct_feature_rows.size());
}
ggml_backend_tensor_set(
graph.kv_write_rows, write_rows.data(), 0,
sizeof(int64_t) * write_rows.size());
Expand All @@ -764,14 +874,15 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
}

std::vector<int32_t> posterior(
static_cast<size_t>(total_rows), -1);
static_cast<size_t>(posterior_rows_count), -1);
ggml_backend_tensor_get(
graph.argmax_tokens, posterior.data(), 0,
sizeof(int32_t) * posterior.size());

for (int lane_index = 0; lane_index < spec_count; ++lane_index) {
Proposal & proposal = proposals[static_cast<size_t>(lane_index)];
const int row_base = ar_count + lane_index * tree_width;
const int row_base =
posterior_tree_row0 + lane_index * tree_width;
const int32_t * lane_posterior =
posterior.data() + static_cast<size_t>(row_base);
size_t accepted = chain_verified_prefix(
Expand Down Expand Up @@ -800,14 +911,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
std::vector<int64_t> commit_rows(
static_cast<size_t>(tree_rows_count), -1);
std::vector<int32_t> feature_commit_rows(
static_cast<size_t>(total_rows), -1);
const int feature_cap = b_.cache_.target_feat_cap;

for (int lane_index = 0; lane_index < ar_count; ++lane_index) {
const ArLane & lane = ar_lanes[static_cast<size_t>(lane_index)];
feature_commit_rows[static_cast<size_t>(lane_index)] =
lane.slot * feature_cap + lane.position % feature_cap;
}
static_cast<size_t>(tree_rows_count), -1);

const auto block_table_delta_fits = [&](int slot, int first_block,
size_t count) {
Expand Down Expand Up @@ -845,10 +949,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
commit_slots[static_cast<size_t>(lane_index)] = proposal.slot;
for (size_t depth = 0; depth < proposal.accepted; ++depth) {
const int flat = lane_index * tree_width + static_cast<int>(depth);
const int graph_row = ar_count + flat;
commit_rows[static_cast<size_t>(flat)] =
append.physical_rows[depth];
feature_commit_rows[static_cast<size_t>(graph_row)] =
feature_commit_rows[static_cast<size_t>(flat)] =
proposal.slot * feature_cap +
(append.position + static_cast<int>(depth)) % feature_cap;
}
Expand Down Expand Up @@ -945,27 +1048,49 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
staged_round.committed = true;
for (int lane_index = 0; lane_index < spec_count; ++lane_index) {
Proposal & proposal = proposals[static_cast<size_t>(lane_index)];
const int graph_row = ar_count + lane_index * tree_width +
const int logits_row = posterior_tree_row0 + lane_index * tree_width +
static_cast<int>(proposal.accepted) - 1;
proposal.pending = sample_graph_row(
proposal.slot, graph_row,
&posterior[static_cast<size_t>(graph_row)], &logits_buf_);
proposal.slot, logits_row,
&posterior[static_cast<size_t>(logits_row)], &logits_buf_);
if (proposal.pending < 0) {
result.error = "fixed chain sampling failed";
return result;
}
}
for (int lane_index = 0; lane_index < ar_count; ++lane_index) {
ArLane & lane = ar_lanes[static_cast<size_t>(lane_index)];
const int logits_row = posterior_ar_row0 + lane_index;
lane.pending = sample_graph_row(
lane.slot, lane_index,
&posterior[static_cast<size_t>(lane_index)], &logits_buf_);
lane.slot, logits_row,
&posterior[static_cast<size_t>(logits_row)], &logits_buf_);
if (lane.pending < 0) {
result.error = "compact AR sampling failed";
return result;
}
}

int prefill_logits_row = 0;
for (size_t i = 0; i < prefills.size(); ++i) {
PrefillOutput output;
output.slot = plan.prefills[i].slot;
if (prefills[i].commit) {
output.status = PrefillOutput::Status::completed;
output.token = sample_graph_row(
output.slot, prefill_logits_row,
&posterior[static_cast<size_t>(prefill_logits_row)],
&logits_buf_);
if (output.token < 0) {
result.prefills.clear();
result.error = "fixed chain prefill sampling failed";
return result;
}
++prefill_logits_row;
slots_.commit_prefill(output.slot);
}
result.prefills.push_back(std::move(output));
}

result.decode.reserve(inputs.size());
for (size_t i = 0; i < inputs.size(); ++i) {
DecodeOutput output;
Expand Down
Loading
Loading