diff --git a/docs/handoffs/pr651-draft-batching-performance.md b/docs/handoffs/pr651-draft-batching-performance.md new file mode 100644 index 000000000..970fcaa84 --- /dev/null +++ b/docs/handoffs/pr651-draft-batching-performance.md @@ -0,0 +1,78 @@ +# PR 651 draft batching performance follow-ups + +> Historical planning note: PR 659 implements the batched append projections and +> packed dynamic-convolution coefficient projections described below. Keeping draft +> hidden states on the device was evaluated but is not part of PR 659; the existing +> host handoff remains intentional until a separately measured change replaces it. + +## Scope + +PR 651 packs the main drafter projections across concurrent lanes. It keeps each lane's cache writes, RoPE, masks, and attention separate. + +This note covers three follow-up optimizations that do not belong in the ownership cleanup: + +- Batch the append projections. +- Pack the dynamic-convolution coefficient projections. +- Keep draft hidden states on the device through chain selection. + +Treat each item as a measured change. Do not combine all three into one patch. + +## Keep these invariants + +- Keep `build_draft_kv_steps()` as the shared C1-C6 implementation. +- Preserve lane-local positions, masks, cache writes, RoPE, attention, and dynamic-convolution history. +- Rebuild a cached graph when the backend or the ordered lane-state pointers change. +- Keep C1 output equivalent to the existing single-lane graph. +- Keep dummy lanes after real lanes, and discard dummy proposals. + +## Batch the append projections + +`draft_kv_batch_build()` currently calls `build_draft_kv_append()` once per lane. Each call runs `draft_fuse_features()`, then the per-layer `wk` and `wv` projections. + +Pack every lane's `ap_feat` columns before those shared-weight matrix multiplications. Split the projected columns before RoPE and `ggml_set_rows()`, because positions, destination rows, and caches remain lane-local. + +The packed append path must preserve the fixed `a_step` width. Padded append rows must still write only to each lane's trash slot. + +Suggested shape: + +```cpp +bool build_draft_kv_appends( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & weights, + const std::vector & lanes); +``` + +Pass one lane from the normal graph and C1-C6 lanes from the batched graph. Delete the old singular builder after migrating both callers. + +Prove the change with the existing multilane output comparison. Add cases with zero, partial, and full append counts so padding and trash-slot writes are covered. Measure draft compute separately at C1, C2, C3, C5, and C6. + +## Pack dynamic-convolution coefficient projections + +`build_draft_kv_steps()` calls `draft_dyn_conv_kernel()` once per lane for both attention and MLP. The function projects normalized hidden columns with shared weights. The temporal convolution in `draft_dyn_conv_apply()` is lane-local and must remain separate. + +Pack the normalized columns before the coefficient projection. Slice the projected coefficients back into lanes, then call `draft_dyn_conv_apply()` per lane. + +Do not pack the convolution history or apply step. Adjacent packed columns belong to different requests and must not influence each other. + +Verify exact lane isolation with distinct inputs and histories. Run the existing single-lane-versus-packed comparison with dynamic convolution enabled. Record kernel count and draft compute time before and after the change. + +## Keep draft hidden states on the device + +The current boundary copies every lane's draft hidden block to the host in `draft_kv_batch_compute()`. `Qwen35SeqEngine::prepare_chain_drafts()` converts the host vectors to pointers. `dflash2_select_chains_batched()` then packs the candidates and uploads them to both the projection graph and the selector graph. + +Replace that round trip with a device-resident contract. The batch graph should expose a packed hidden tensor or stable per-lane tensor views. The batched selector should accept those device tensors on the same backend. + +Keep only token IDs, top-K scores, and final proposals as host results. Do not expose an unowned device pointer whose lifetime is shorter than either consumer graph. + +This change crosses the draft and selector APIs, so ship it separately from append or dynamic-convolution packing. It also needs an explicit fallback when the draft and selector backends differ. + +Verification must compare complete proposals, not only projected hidden values. Cover C1-C6, reordered active slots, bucket padding, selector-graph reuse, and graph rebuilds. Measure transfer bytes, synchronization count, selector time, and complete-round latency. + +## Suggested order + +1. Batch append projections. This extends the packing pattern already introduced by PR 651. +2. Pack dynamic-convolution coefficient projections. This is local to `build_draft_kv_steps()`. +3. Remove the host hidden-state handoff in its own PR. This changes ownership across the draft and selector graphs. + +For every step, retain the previous implementation long enough to run an A/B output comparison. Delete it before merging the step. diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-replay-log.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-replay-log.cu index d92ea9de0..fddbc066d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-replay-log.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-replay-log.cu @@ -66,6 +66,62 @@ __global__ void gdn_replay_log_commit_kernel( state[state_offset] = current; } +__global__ void gdn_replay_log_commit_128_scalar_tile_kernel( + const float * replay_log, + float * state, + const int32_t * accepted_prefixes, + const int32_t * active_slot_ids, + int n_heads, + int n_tokens, + int n_seqs, + int n_state_slots) { + constexpr int state_size = 128; + constexpr int max_tokens = 8; + constexpr int values_per_token = 1 + 32 + 8; + const int sequence = blockIdx.z; + const int head = blockIdx.y; + if (sequence >= n_seqs || head >= n_heads) return; + const int slot = active_slot_ids[sequence]; + const int accepted = accepted_prefixes[sequence]; + if (slot < 0 || slot >= n_state_slots || + accepted < 0 || accepted > n_tokens || accepted > max_tokens) { + return; + } + const int row_base = (blockIdx.x & 3)*32; + const int col_base = (blockIdx.x >> 2)*8; + const int row = row_base + (threadIdx.x & 31); + const int col = col_base + (threadIdx.x >> 5); + __shared__ float staged[max_tokens][values_per_token]; + for (int item = threadIdx.x; + item < accepted*values_per_token; + item += blockDim.x) { + const int token = item/values_per_token; + const int field = item % values_per_token; + const float * transition = replay_log + + (((size_t) sequence*n_tokens + token)*n_heads + head) * + (2*state_size + 1); + if (field == 0) { + staged[token][0] = transition[0]; + } else if (field <= 32) { + staged[token][field] = transition[row_base + field]; + } else { + staged[token][field] = + transition[1 + state_size + col_base + field - 33]; + } + } + __syncthreads(); + const size_t state_offset = + (((size_t) slot*n_heads + head)*state_size + col)*state_size + row; + float current = state[state_offset]; + for (int token = 0; token < accepted; ++token) { + current = fmaf( + staged[token][1 + (threadIdx.x & 31)], + staged[token][33 + (threadIdx.x >> 5)], + staged[token][0]*current); + } + state[state_offset] = current; +} + bool device_pointer(const void * pointer, int & device) { if (pointer == nullptr) return false; cudaPointerAttributes attributes{}; @@ -280,13 +336,22 @@ static bool gdn_replay_log_commit_many_impl( const dim3 state_grid( (unsigned int) ((state_elements + threads - 1)/threads), (unsigned int) heads, (unsigned int) n_seqs); - gdn_replay_log_commit_kernel<<>>( - (const float *) replay_log->data, (float *) state->data, - (const int32_t *) accepted_prefixes->data, - (const int32_t *) active_slot_ids->data, - state_size, heads, tokens, (int) n_seqs, - (int) state->ne[3], replay_log_width, - replay_log_width == 2*state_size + 1 ? 1 : state_size); + if (state_size == 128 && replay_log_width == 2*state_size + 1 && + tokens <= 8) { + gdn_replay_log_commit_128_scalar_tile_kernel<<>>( + (const float *) replay_log->data, (float *) state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + heads, tokens, (int) n_seqs, (int) state->ne[3]); + } else { + gdn_replay_log_commit_kernel<<>>( + (const float *) replay_log->data, (float *) state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + state_size, heads, tokens, (int) n_seqs, + (int) state->ne[3], replay_log_width, + replay_log_width == 2*state_size + 1 ? 1 : state_size); + } const ggml_tensor * conv_input = conv_inputs[layer]; ggml_tensor * conv_state = conv_states[layer]; diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h index 23b16e7fc..44da3d699 100644 --- a/server/src/common/concurrency/chain_spec_shapes.h +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -18,6 +18,12 @@ inline int chain_decode_bucket_width(int lanes) { return 64; } +inline int chain_draft_bucket_width(int lanes) { + // Exact C=5 is stable for the draft graph. Target verification keeps the + // shared decode bucket because changing its MMQ shape changes numerics. + return lanes == 5 ? 5 : chain_decode_bucket_width(lanes); +} + // draft_tokens[0] is the already-pending root. The target posterior at row N // verifies draft_tokens[N + 1], so accepted tokens always form a prefix. inline size_t chain_verified_prefix( diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index efa0a34bd..34557cefc 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -99,12 +99,10 @@ static bool draft_kv_init_impl(DraftKvState & st, if (!st.g_ctx) return false; st.gf = ggml_new_graph_custom(st.g_ctx, 4096, false); - DraftKvAppendInputs ai{}; - ai.n_rows = st.a_step; - ai.feat = st.ap_feat; - ai.positions = st.ap_pos; - ai.rows = st.ap_rows; - if (!build_draft_kv_append(st.g_ctx, st.gf, dw, st.cache, ai)) return false; + const std::vector append_lanes{ + {&st.cache, st.ap_feat, st.ap_pos, st.ap_rows}, + }; + if (!build_draft_kv_appends(st.g_ctx, st.gf, dw, append_lanes)) return false; DraftKvStepInputs si{}; si.noise_embed = st.inp_embed; @@ -221,8 +219,10 @@ static bool draft_kv_bulk_append(DraftKvState & st, ok = gctx != nullptr; if (ok) { ggml_cgraph * g = ggml_new_graph_custom(gctx, 4096, false); - DraftKvAppendInputs ai{c, feat, tpos, trow}; - ok = build_draft_kv_append(gctx, g, dw, st.cache, ai); + const std::vector append_lanes{ + {&st.cache, feat, tpos, trow}, + }; + ok = build_draft_kv_appends(gctx, g, dw, append_lanes); if (!ok) std::fprintf(stderr, "[draft-kv] bulk: append build failed\n"); if (ok) { ggml_gallocr_t ga = @@ -419,19 +419,13 @@ static bool draft_kv_batch_build( batch.g_ctx, graph_capacity, false); batch.hidden_by_lane.reserve(static_cast(n_lanes)); + std::vector append_lanes; + append_lanes.reserve(n_lanes_size); std::vector lanes; lanes.reserve(n_lanes_size); for (DraftKvState * state : lane_states) { - DraftKvAppendInputs append{}; - append.n_rows = state->a_step; - append.feat = state->ap_feat; - append.positions = state->ap_pos; - append.rows = state->ap_rows; - if (!build_draft_kv_append( - batch.g_ctx, batch.gf, dw, state->cache, append)) { - draft_kv_batch_free(batch); - return false; - } + append_lanes.push_back( + {&state->cache, state->ap_feat, state->ap_pos, state->ap_rows}); DraftKvStepInputs step{}; step.noise_embed = state->inp_embed; @@ -442,6 +436,11 @@ static bool draft_kv_batch_build( lanes.push_back({&state->cache, step}); } + if (!build_draft_kv_appends( + batch.g_ctx, batch.gf, dw, append_lanes)) { + draft_kv_batch_free(batch); + return false; + } const std::vector outputs = build_draft_kv_steps( batch.g_ctx, batch.gf, dw, lanes); if (outputs.size() != lane_states.size()) { diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index dbf342d60..b774cd5ed 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -1,9 +1,9 @@ // StepGraph — per-forward-call compute graph container. // -// Holds the ggml context, graph, allocator, and named tensor handles for one -// forward step (prefill chunk, verify batch, or replay). Rebuilt per call -// since kv_len varies, but the persistent CUDA allocator buffer is kept -// alive across steps to avoid cudaMalloc/cudaFree churn. +// Holds the ggml context, graph, allocator, and named tensor handles for a +// forward topology (prefill chunk, verify batch, or replay). Most paths +// rebuild as shapes vary; topology-stable paths can retain the graph. The +// persistent CUDA allocator buffer stays alive across rebuilds. #pragma once @@ -12,23 +12,33 @@ #include "ggml.h" #include "ggml-alloc.h" +#include +#include +#include #include namespace dflash::common { +using TargetPagedTreeGraphKey = std::tuple< + const TargetWeights *, const TargetCache *, ggml_backend_t, + int, int, int, int, int, int>; + struct StepGraph { ggml_context * ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_gallocr_t alloc = nullptr; ggml_context * commit_ctx = nullptr; ggml_backend_buffer_t commit_buffer = nullptr; + std::optional paged_tree_key; // Persistent metadata arena for the draft graph. Reusing the same arena - // across rebuilds keeps every ggml_tensor at a stable address, which is - // what the ggml-cuda graph cache keys on (nodes[0] pointer + src tensor - // pointers). A fresh malloc per step would defeat CUDA-graph replay. + // keeps ggml_tensor addresses stable for CUDA-graph replay. std::vector meta_arena; + // Paged-tree metadata is retained with the graph. Use uninitialized + // storage: vector::resize would zero 512 MiB on the first decode round. + std::shared_ptr paged_tree_meta_arena; + // The ctx_len last used for ggml_gallocr_reserve (draft only). // When the real ctx_len fits within this, alloc_graph is a no-op. int alloc_reserved_ctx = 0; @@ -143,12 +153,14 @@ inline void step_graph_free(StepGraph & sg) { sg.delta_captures.clear(); sg.tree_features = nullptr; sg.moe_selected.clear(); + sg.paged_tree_key.reset(); } // Full cleanup: release the persistent gallocr + its CUDA buffer. inline void step_graph_destroy(StepGraph & sg) { if (sg.alloc) { ggml_gallocr_free(sg.alloc); sg.alloc = nullptr; } step_graph_free(sg); + sg.paged_tree_meta_arena.reset(); sg.meta_arena.clear(); sg.meta_arena.shrink_to_fit(); sg.alloc_reserved_ctx = 0; diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 1efde29f2..aab193862 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -37,6 +37,7 @@ #include #include #include +#include namespace dflash::common { @@ -387,60 +388,6 @@ DraftGraphOutputs build_draft_graph( // legacy concat-then-normalize math exactly; the only numeric difference is // the F16 cache storage (legacy keeps K/V in F32 for the one shot). -// Per-layer ctx-side K/V rows in the head-major cache layout. -// K: wk @ tf_kv → per-head k_norm → RoPE(positions) → [head_dim*n_kv, n] rows. -// V: wv @ tf_kv, already row-shaped. -static void draft_ctx_kv_rows( - ggml_context * ctx, - const DraftWeights & w, - const DraftLayer & L, - ggml_tensor * target_feat, // [hidden, n] - ggml_tensor * positions, // [n] i32 absolute - int n, - ggml_tensor ** k_rows_out, - ggml_tensor ** v_rows_out) { - const float eps = DFLASH27B_RMS_EPS; - ggml_tensor * tf_kv = target_feat; - if (w.context_kv_layer_norm) { - tf_kv = ggml_rms_norm(ctx, tf_kv, eps); - tf_kv = ggml_mul(ctx, tf_kv, L.attn_norm); - } - ggml_tensor * K = ggml_mul_mat(ctx, L.wk, tf_kv); // [kv_dim, n] - K = ggml_reshape_3d(ctx, K, w.head_dim, w.n_head_kv, n); - K = ggml_rms_norm(ctx, K, eps); - K = ggml_mul (ctx, K, L.k_norm); - K = draft_rope(ctx, K, positions, w); - // rope output is contiguous [head_dim, n_kv, n] → head-major rows view - *k_rows_out = ggml_view_2d(ctx, K, (int64_t)w.head_dim * w.n_head_kv, n, - K->nb[2], 0); - *v_rows_out = ggml_mul_mat(ctx, L.wv, tf_kv); // [kv_dim, n] -} - -bool build_draft_kv_append( - ggml_context * ctx, - ggml_cgraph * gf, - const DraftWeights & w, - const DraftKvCacheRefs & cache, - const DraftKvAppendInputs & in) { - if (!in.feat || !in.positions || !in.rows || in.n_rows <= 0) return false; - static const bool disable_aux_hidden_norms = - std::getenv("DFLASH_DISABLE_DRAFT_AUX_NORMS") != nullptr; - - ggml_tensor * target_feat = draft_fuse_features( - ctx, w, in.feat, in.n_rows, disable_aux_hidden_norms); - ggml_set_name(target_feat, "draft_kv_append_feat"); - - for (int il = 0; il < w.n_layer; il++) { - ggml_tensor * Krows = nullptr; - ggml_tensor * Vrows = nullptr; - draft_ctx_kv_rows(ctx, w, w.layers[il], target_feat, in.positions, - in.n_rows, &Krows, &Vrows); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache.k[il], Krows, in.rows)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache.v[il], Vrows, in.rows)); - } - return true; -} - static ggml_tensor * draft_pack_columns( ggml_context * ctx, const std::vector & lanes) { @@ -459,6 +406,105 @@ static ggml_tensor * draft_lane_columns( lane * static_cast(q_len) * packed->nb[1]); } +static std::vector draft_dyn_conv_kernels( + ggml_context * ctx, const DraftConvWeights & weights, + const std::vector & lanes, int q_len) { + if (lanes.size() == 1) { + return {draft_dyn_conv_kernel(ctx, weights, lanes.front())}; + } + const DraftDynConv packed = draft_dyn_conv_kernel( + ctx, weights, draft_pack_columns(ctx, lanes)); + std::vector kernels(lanes.size()); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + kernels[lane].dyn = draft_lane_columns(ctx, packed.dyn, q_len, lane); + } + return kernels; +} + +bool build_draft_kv_appends( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & w, + const std::vector & lanes) { + if (!ctx || !gf || lanes.empty()) { + return false; + } + + const int64_t append_width = lanes.front().feat + ? lanes.front().feat->ne[1] + : 0; + if (append_width <= 0 || + lanes.size() > static_cast( + std::numeric_limits::max() / append_width)) { + return false; + } + const int packed_width = + static_cast(append_width) * static_cast(lanes.size()); + const int64_t feature_width = lanes.front().feat->ne[0]; + + std::vector lane_features; + lane_features.reserve(lanes.size()); + for (const DraftKvAppendLane & lane : lanes) { + if (!lane.cache || !lane.feat || !lane.positions || !lane.rows || + lane.feat->ne[1] != append_width || + lane.feat->ne[0] != feature_width || + lane.positions->ne[0] != append_width || + lane.rows->ne[0] != append_width || + lane.cache->k.size() != static_cast(w.n_layer) || + lane.cache->v.size() != static_cast(w.n_layer)) { + return false; + } + lane_features.push_back(lane.feat); + } + + const int width = static_cast(append_width); + static const bool disable_aux_hidden_norms = + std::getenv("DFLASH_DISABLE_DRAFT_AUX_NORMS") != nullptr; + ggml_tensor * packed_features = + draft_pack_columns(ctx, lane_features); + ggml_tensor * target_feat = draft_fuse_features( + ctx, w, packed_features, packed_width, + disable_aux_hidden_norms); + ggml_set_name(target_feat, "draft_kv_append_feat"); + + const float eps = DFLASH27B_RMS_EPS; + for (int il = 0; il < w.n_layer; ++il) { + const DraftLayer & layer = w.layers[il]; + ggml_tensor * tf_kv = target_feat; + if (w.context_kv_layer_norm) { + tf_kv = ggml_rms_norm(ctx, tf_kv, eps); + tf_kv = ggml_mul(ctx, tf_kv, layer.attn_norm); + } + ggml_tensor * packed_k = ggml_mul_mat(ctx, layer.wk, tf_kv); + ggml_tensor * packed_v = ggml_mul_mat(ctx, layer.wv, tf_kv); + + for (size_t lane_index = 0; + lane_index < lanes.size(); ++lane_index) { + const DraftKvAppendLane & lane = lanes[lane_index]; + ggml_tensor * K = draft_lane_columns( + ctx, packed_k, width, lane_index); + K = ggml_reshape_3d( + ctx, K, w.head_dim, w.n_head_kv, width); + K = ggml_rms_norm(ctx, K, eps); + K = ggml_mul(ctx, K, layer.k_norm); + K = draft_rope(ctx, K, lane.positions, w); + ggml_tensor * Krows = ggml_view_2d( + ctx, K, + static_cast(w.head_dim) * w.n_head_kv, + width, K->nb[2], 0); + ggml_tensor * Vrows = draft_lane_columns( + ctx, packed_v, width, lane_index); + ggml_build_forward_expand( + gf, ggml_set_rows( + ctx, lane.cache->k[il], Krows, lane.rows)); + ggml_build_forward_expand( + gf, ggml_set_rows( + ctx, lane.cache->v[il], Vrows, lane.rows)); + } + } + return true; +} + std::vector build_draft_kv_steps( ggml_context * ctx, ggml_cgraph * gf, @@ -508,9 +554,11 @@ std::vector build_draft_kv_steps( for (size_t lane = 0; lane < n_lanes; ++lane) { hn[lane] = ggml_rms_norm(ctx, h[lane], eps); hn[lane] = ggml_mul(ctx, hn[lane], layer.attn_norm); - if (dyn_conv) { - attn_dc[lane] = draft_dyn_conv_kernel( - ctx, layer.attn_conv, hn[lane]); + } + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernels( + ctx, layer.attn_conv, hn, q_len); + for (size_t lane = 0; lane < n_lanes; ++lane) { hn[lane] = draft_dyn_conv_apply( ctx, w, layer.attn_conv, attn_dc[lane], 0, hn[lane]); } @@ -601,9 +649,11 @@ std::vector build_draft_kv_steps( for (size_t lane = 0; lane < n_lanes; ++lane) { hf[lane] = ggml_rms_norm(ctx, h[lane], eps); hf[lane] = ggml_mul(ctx, hf[lane], layer.ffn_norm); - if (dyn_conv) { - mlp_dc[lane] = draft_dyn_conv_kernel( - ctx, layer.mlp_conv, hf[lane]); + } + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernels( + ctx, layer.mlp_conv, hf, q_len); + for (size_t lane = 0; lane < n_lanes; ++lane) { hf[lane] = draft_dyn_conv_apply( ctx, w, layer.mlp_conv, mlp_dc[lane], 0, hf[lane]); } diff --git a/server/src/draft/draft_graph.h b/server/src/draft/draft_graph.h index 48deba49b..d18f4eb4c 100644 --- a/server/src/draft/draft_graph.h +++ b/server/src/draft/draft_graph.h @@ -60,19 +60,18 @@ struct DraftKvCacheRefs { std::vector v; // per layer [head_dim*n_head_kv, kv_total] f16 }; -// Fuse `n_rows` feature rows and set_rows their per-layer K/V into the caches. -struct DraftKvAppendInputs { - int n_rows = 0; - ggml_tensor * feat = nullptr; // [n_target_layers*hidden, n_rows] f32 - ggml_tensor * positions = nullptr; // [n_rows] i32, absolute - ggml_tensor * rows = nullptr; // [n_rows] i32, destination cache slots +struct DraftKvAppendLane { + const DraftKvCacheRefs * cache = nullptr; + ggml_tensor * feat = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * rows = nullptr; }; -bool build_draft_kv_append( - ggml_context * ctx, - ggml_cgraph * gf, - const DraftWeights & w, - const DraftKvCacheRefs & cache, - const DraftKvAppendInputs & in); + +bool build_draft_kv_appends( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & w, + const std::vector & lanes); // One draft step over the cached context KV. Noise K/V are written into the // scratch slots and the flash attention reads the full kv_total span; the diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 92b937841..002cfdca8 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -263,7 +263,7 @@ Qwen35SeqEngine::prepare_chain_drafts( } if (lanes.empty()) return round; - const int bucket = chain_decode_bucket_width( + const int bucket = chain_draft_bucket_width( static_cast(lanes.size())); std::vector batch_states; std::vector roots; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 59df8bc1a..f59d4af31 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -818,16 +818,40 @@ bool build_target_step_paged_tree( int kq_stride_pad, int mapped_ar_seqs) { (void)kq_stride_pad; - step_graph_free(sg); if (mapped_ar_seqs < 0 || mapped_ar_seqs > cache.n_seq_slots) { + step_graph_free(sg); return false; } if (!detail::validate_target_paged_tree_layout( cache, tree_width, n_tree_seqs, paged_max_kv_len, tree_scratch_base, tree_scratch_stride)) { + step_graph_free(sg); + return false; + } + const int64_t table_capacity = + cache.paged_block_table->ne[0] * PAGED_BLOCK_SIZE; + const int64_t logical_capacity = + std::min(cache.max_ctx, table_capacity); + if (logical_capacity < 1 || logical_capacity > INT32_MAX) { + step_graph_free(sg); return false; } + const int64_t requested = std::min( + std::max(1, paged_max_kv_len), logical_capacity); + const int paged_launch_kv_len = static_cast( + std::min(((requested + 255) / 256) * 256, + logical_capacity)); + const TargetPagedTreeGraphKey graph_key{ + &w, &cache, backend, tree_width, n_tree_seqs, + paged_launch_kv_len, tree_scratch_base, tree_scratch_stride, + mapped_ar_seqs, + }; + if (sg.paged_tree_key && *sg.paged_tree_key == graph_key) { + return true; + } + step_graph_free(sg); + size_t graph_capacity = 0; if (!detail::target_paged_tree_graph_capacity( tree_width, n_tree_seqs, graph_capacity)) { @@ -838,13 +862,11 @@ bool build_target_step_paged_tree( ggml_init_params ip{}; ip.mem_size = 512 * 1024 * 1024; - static thread_local std::unique_ptr g_tree_arena; - static thread_local size_t g_tree_arena_size = 0; - if (g_tree_arena_size < ip.mem_size) { - g_tree_arena.reset(new uint8_t[ip.mem_size]); - g_tree_arena_size = ip.mem_size; + if (!sg.paged_tree_meta_arena) { + sg.paged_tree_meta_arena.reset( + new uint8_t[ip.mem_size], std::default_delete()); } - ip.mem_buffer = g_tree_arena.get(); + ip.mem_buffer = sg.paged_tree_meta_arena.get(); ip.no_alloc = true; sg.ctx = ggml_init(ip); if (!sg.ctx) return false; @@ -917,7 +939,7 @@ bool build_target_step_paged_tree( gi.paged_query_positions = sg.paged_query_positions; gi.n_seqs = n_tree_seqs; gi.mapped_ar_seqs = mapped_ar_seqs; - gi.paged_max_kv_len = paged_max_kv_len; + gi.paged_max_kv_len = paged_launch_kv_len; gi.tree_width = tree_width; gi.tree_scratch_base = tree_scratch_base; gi.tree_scratch_stride = tree_scratch_stride; @@ -963,7 +985,9 @@ bool build_target_step_paged_tree( ggml_set_name(sg.feature_commit_rows, "feature_commit_rows"); sg.commit_buffer = ggml_backend_alloc_ctx_tensors( sg.commit_ctx, backend); - return sg.commit_buffer != nullptr; + if (!sg.commit_buffer) return false; + sg.paged_tree_key = graph_key; + return true; } diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 70308fb39..20a4b4f1a 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -885,6 +885,7 @@ bool Qwen35Backend::park(ParkTarget target) { std::printf("[park] draft released\n"); std::fflush(stdout); } if (want_target_model && !target_parked_) { + step_graph_destroy(sg_); step_graph_destroy(proj_sg_); dflash2_selector_graph_invalidate(); free_target_weights(w_); @@ -956,6 +957,9 @@ bool Qwen35Backend::unpark(ParkTarget target) { dw_.block_size = cfg_.draft_block_size; } } + // A reloaded drafter can select different target capture layers. + // Rebuild any retained target graph against the refreshed mapping. + step_graph_free(sg_); draft_parked_ = false; std::printf("[unpark] draft restored\n"); std::fflush(stdout); } diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index f3b6b9f32..0e87c1a36 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -1939,12 +1939,11 @@ static ggml_tensor * build_delta_net_block( } // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout. - // The fused chain/tree gated_delta_net kernels broadcast heads themselves + // The fused active/chain/tree gated_delta_net kernels broadcast heads themselves // (v head h reads q/k head h % num_k_heads, the same tiling ggml_repeat - // produces); the chunked, compact-decode and SpecLA paths take the - // materialized copies. + // produces); the chunked and SpecLA paths take the materialized copies. if (num_k_heads != num_v_heads && - (use_chunked || seg_active || use_specla_factorized || use_specla_hld)) { + (use_chunked || use_specla_factorized || use_specla_hld)) { q_c = ggml_repeat_4d(ctx, q_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); k_c = ggml_repeat_4d(ctx, k_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); } diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index 54b64a285..63fd2e8c8 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -37,6 +37,10 @@ int main() { CHECK(chain_decode_bucket_width(bucket_inputs[i]) == bucket_expected[i]); } + CHECK(chain_draft_bucket_width(4) == 4); + CHECK(chain_draft_bucket_width(5) == 5); + CHECK(chain_draft_bucket_width(6) == 6); + CHECK(chain_draft_bucket_width(7) == 8); const auto eos = [](int32_t token) { return token == 2; }; const std::vector eos_first_child = {10, 2, 11}; diff --git a/server/test/test_draft_swa_multilane.cpp b/server/test/test_draft_swa_multilane.cpp index 5a53bc0e2..84735c868 100644 --- a/server/test/test_draft_swa_multilane.cpp +++ b/server/test/test_draft_swa_multilane.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -96,19 +97,329 @@ float max_abs_diff(const std::vector & lhs, return max_diff; } +struct AppendStateFamilies { + std::array reference; + std::array packed; + + ~AppendStateFamilies() { + for (DraftKvState & state : reference) { + draft_kv_free(state); + } + for (DraftKvState & state : packed) { + draft_kv_free(state); + } + } +}; + +bool compute_append_graph( + ggml_backend_t backend, + const DraftWeights & weights, + const std::vector & lanes) { + std::vector arena( + 16u * 1024 * 1024 * std::max(lanes.size(), 1)); + ggml_init_params params{}; + params.mem_size = arena.size(); + params.mem_buffer = arena.data(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return false; + } + + ggml_cgraph * graph = ggml_new_graph_custom( + ctx, 4096 * std::max(lanes.size(), 1), false); + bool ok = build_draft_kv_appends(ctx, graph, weights, lanes); + ggml_gallocr_t allocator = nullptr; + if (ok) { + allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + ok = allocator && + ggml_gallocr_alloc_graph(allocator, graph) && + ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS; + } + if (allocator) { + ggml_gallocr_free(allocator); + } + ggml_free(ctx); + return ok; +} + +void fill_append_inputs( + DraftKvState & state, + int lane, + int append_count) { + const size_t feature_elements = + static_cast(state.fc_in) * state.a_step; + std::vector features(feature_elements); + for (int column = 0; column < state.a_step; ++column) { + const int feature_column = + column < append_count ? column : append_count; + for (int row = 0; row < state.fc_in; ++row) { + const size_t index = + static_cast(column) * state.fc_in + row; + features[index] = + 0.025f * static_cast(lane + 1) + + 0.00001f * static_cast( + (row + 17 * feature_column) % 997); + } + } + + std::vector positions(static_cast(state.a_step)); + std::vector rows(static_cast(state.a_step)); + for (int column = 0; column < state.a_step; ++column) { + if (column < append_count) { + const int position = 11 + 37 * lane + column; + positions[static_cast(column)] = position; + rows[static_cast(column)] = position % state.cap; + } else { + positions[static_cast(column)] = 0; + rows[static_cast(column)] = state.trash_slot; + } + } + + ggml_backend_tensor_set( + state.ap_feat, features.data(), 0, + features.size() * sizeof(float)); + ggml_backend_tensor_set( + state.ap_pos, positions.data(), 0, + positions.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + state.ap_rows, rows.data(), 0, + rows.size() * sizeof(int32_t)); +} + +bool compare_cache_tensor( + const ggml_tensor * reference, + const ggml_tensor * packed, + int lane, + int layer, + int trash_row, + const char * kind) { + if (reference->ne[0] != packed->ne[0] || + reference->ne[1] != packed->ne[1]) { + return false; + } + + const size_t elements = static_cast(ggml_nelements(reference)); + std::vector reference_values(elements); + std::vector packed_values(elements); + ggml_backend_tensor_get( + reference, reference_values.data(), 0, + elements * sizeof(ggml_fp16_t)); + ggml_backend_tensor_get( + packed, packed_values.data(), 0, + elements * sizeof(ggml_fp16_t)); + + const size_t row_width = static_cast(reference->ne[0]); + if (trash_row < 0 || + trash_row >= reference->ne[1]) { + return false; + } + float max_error = 0.0f; + float trash_error = 0.0f; + size_t max_index = 0; + for (size_t index = 0; index < elements; ++index) { + const float reference_value = + ggml_fp16_to_fp32(reference_values[index]); + const float packed_value = + ggml_fp16_to_fp32(packed_values[index]); + if (!std::isfinite(reference_value) || + !std::isfinite(packed_value)) { + return false; + } + const float error = + std::fabs(reference_value - packed_value); + if (index / row_width == static_cast(trash_row)) { + trash_error = std::max(trash_error, error); + } + if (error > max_error) { + max_error = error; + max_index = index; + } + } + if (max_error > MAX_ABS_ERROR) { + std::fprintf( + stderr, + "draft append cache mismatch lane=%d layer=%d kind=%s " + "row=%zu trash=%d max_abs=%.6g trash_abs=%.6g\n", + lane, layer, kind, max_index / row_width, trash_row, + max_error, trash_error); + return false; + } + return true; +} + +bool check_trash_row( + const ggml_tensor * cache, + int trash_row, + bool expect_written, + int lane, + int layer, + const char * family, + const char * kind) { + const size_t row_width = static_cast(cache->ne[0]); + std::vector values(row_width); + ggml_backend_tensor_get( + cache, values.data(), + static_cast(trash_row) * row_width * sizeof(ggml_fp16_t), + row_width * sizeof(ggml_fp16_t)); + + bool written = false; + for (ggml_fp16_t value : values) { + const float converted = ggml_fp16_to_fp32(value); + if (!std::isfinite(converted)) { + return false; + } + written = written || converted != 0.0f; + } + if (written != expect_written) { + std::fprintf( + stderr, + "draft append trash mismatch family=%s lane=%d layer=%d " + "kind=%s expected_written=%d actual_written=%d\n", + family, lane, layer, kind, + expect_written ? 1 : 0, written ? 1 : 0); + return false; + } + return true; +} + +bool check_packed_append_caches( + ggml_backend_t backend, + const DraftWeights & weights) { + AppendStateFamilies families; + for (int lane = 0; lane < N_LANES; ++lane) { + if (!draft_kv_init_batched( + families.reference[static_cast(lane)], + weights, backend, CAPACITY) || + !draft_kv_init_batched( + families.packed[static_cast(lane)], + weights, backend, CAPACITY)) { + std::fprintf( + stderr, + "draft append qualification: lane %d state init failed\n", + lane); + return false; + } + } + + const int append_width = families.reference.front().a_step; + const std::array append_counts{ + 0, append_width / 2, append_width, + }; + for (int lane = 0; lane < N_LANES; ++lane) { + fill_append_inputs( + families.reference[static_cast(lane)], + lane, append_counts[static_cast(lane)]); + fill_append_inputs( + families.packed[static_cast(lane)], + lane, append_counts[static_cast(lane)]); + } + + for (DraftKvState & state : families.reference) { + const std::vector lane{ + {&state.cache, state.ap_feat, state.ap_pos, state.ap_rows}, + }; + if (!compute_append_graph(backend, weights, lane)) { + std::fprintf( + stderr, + "draft append qualification: reference graph failed\n"); + return false; + } + } + + std::vector packed_lanes; + packed_lanes.reserve(N_LANES); + for (DraftKvState & state : families.packed) { + packed_lanes.push_back( + {&state.cache, state.ap_feat, state.ap_pos, state.ap_rows}); + } + if (!compute_append_graph(backend, weights, packed_lanes)) { + std::fprintf( + stderr, + "draft append qualification: packed graph failed\n"); + return false; + } + + for (int lane = 0; lane < N_LANES; ++lane) { + const DraftKvState & reference = + families.reference[static_cast(lane)]; + const DraftKvState & packed = + families.packed[static_cast(lane)]; + const bool expect_trash_written = + append_counts[static_cast(lane)] < append_width; + for (int layer = 0; layer < weights.n_layer; ++layer) { + if (!compare_cache_tensor( + reference.cache.k[static_cast(layer)], + packed.cache.k[static_cast(layer)], + lane, layer, reference.trash_slot, "K") || + !compare_cache_tensor( + reference.cache.v[static_cast(layer)], + packed.cache.v[static_cast(layer)], + lane, layer, reference.trash_slot, "V")) { + return false; + } + if (!check_trash_row( + reference.cache.k[static_cast(layer)], + reference.trash_slot, expect_trash_written, + lane, layer, "reference", "K") || + !check_trash_row( + reference.cache.v[static_cast(layer)], + reference.trash_slot, expect_trash_written, + lane, layer, "reference", "V") || + !check_trash_row( + packed.cache.k[static_cast(layer)], + packed.trash_slot, expect_trash_written, + lane, layer, "packed", "K") || + !check_trash_row( + packed.cache.v[static_cast(layer)], + packed.trash_slot, expect_trash_written, + lane, layer, "packed", "V")) { + return false; + } + } + } + + std::printf( + "draft append packed cache qualification passed " + "counts=0,%d,%d\n", + append_width / 2, append_width); + return true; +} + } // namespace int main(int argc, char ** argv) { if (argc == 1) { std::fprintf(stderr, - "usage: %s [gpu]\n", argv[0]); + "usage: %s [gpu] " + "[--append-only|--dynconv-only]\n", argv[0]); return 77; } - if (argc > 3) { + if (argc > 4) { return 2; } - const int gpu = argc == 3 ? std::atoi(argv[2]) : 0; + int gpu = 0; + bool gpu_set = false; + bool append_only = false; + bool dynconv_only = false; + for (int arg = 2; arg < argc; ++arg) { + if (std::strcmp(argv[arg], "--append-only") == 0) { + append_only = true; + } else if (std::strcmp(argv[arg], "--dynconv-only") == 0) { + dynconv_only = true; + } else if (gpu_set) { + return 2; + } else { + gpu = std::atoi(argv[arg]); + gpu_set = true; + } + } + if (append_only && dynconv_only) { + return 2; + } Resources resources; resources.backend = ggml_backend_cuda_init(gpu); if (!resources.backend) { @@ -120,20 +431,40 @@ int main(int argc, char ** argv) { dflash27b_last_error()); return 1; } - - const std::vector trained_pattern = layer_pattern(resources.weights); - if (!resources.weights.swa_pattern_loaded) { - std::fprintf(stderr, "draft SWA qualification: GGUF has no SWA pattern\n"); + if (dynconv_only && + (resources.weights.conv_kernel_size <= 0 || + resources.weights.conv_group_size <= 0 || + !std::all_of(resources.weights.layers.begin(), resources.weights.layers.end(), + [](const DraftLayer & layer) { + return layer.attn_conv.present() && layer.mlp_conv.present(); + }))) { + std::fprintf(stderr, "draft SWA qualification: dynamic convolution absent\n"); return 1; } - const DraftSwaOverrideResult swa = - apply_draft_swa_window_override(resources.weights, SWA_WINDOW); - if (layer_pattern(resources.weights) != trained_pattern || - swa.effective_window != SWA_WINDOW || swa.swa_layers == 0) { - std::fprintf(stderr, - "draft SWA qualification: override changed the trained pattern\n"); + if (!check_packed_append_caches( + resources.backend, resources.weights)) { return 1; } + if (append_only) { + return 0; + } + + if (!dynconv_only) { + const std::vector trained_pattern = layer_pattern(resources.weights); + if (!resources.weights.swa_pattern_loaded) { + std::fprintf(stderr, + "draft SWA qualification: GGUF has no SWA pattern\n"); + return 1; + } + const DraftSwaOverrideResult swa = + apply_draft_swa_window_override(resources.weights, SWA_WINDOW); + if (layer_pattern(resources.weights) != trained_pattern || + swa.effective_window != SWA_WINDOW || swa.swa_layers == 0) { + std::fprintf(stderr, + "draft SWA qualification: override changed the trained pattern\n"); + return 1; + } + } DraftKvState batched_state; if (!draft_kv_init_batched( @@ -169,15 +500,19 @@ int main(int argc, char ** argv) { if (!draft_kv_begin_step( state, resources.weights, resources.backend, unused_ring, committed[static_cast(lane)]) || - !check_lane_mask(state, committed[static_cast(lane)])) { + (!dynconv_only && + !check_lane_mask(state, committed[static_cast(lane)]))) { return 1; } std::vector embedding(hidden_elements); for (size_t index = 0; index < embedding.size(); ++index) { + const size_t token = index / resources.weights.n_embd; + const size_t channel = index % resources.weights.n_embd; embedding[index] = 0.01f * static_cast(lane + 1) + - 0.0001f * static_cast(index % 31); + 0.0001f * static_cast( + (channel + 17 * (lane + 3) * (token + 1)) % 127); } ggml_backend_tensor_set(state.inp_embed, embedding.data(), 0, embedding.size() * sizeof(float)); @@ -225,6 +560,7 @@ int main(int argc, char ** argv) { } } - std::printf("draft SWA post-window three-lane qualification passed\n"); + std::printf("draft %s three-lane qualification passed\n", + dynconv_only ? "dynamic-convolution" : "SWA post-window"); return 0; } diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index dcf6f4628..ddb945450 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -12,6 +12,7 @@ using namespace CppUnitTestFramework; using dflash::common::TargetCache; +using dflash::common::TargetPagedTreeGraphKey; using dflash::common::StepGraph; using dflash::common::restore_ssm_state; using dflash::common::snapshot_ssm_state; @@ -144,6 +145,19 @@ TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) { } +TEST_CASE(RecurrentSnapshotFixture, invalidates_paged_tree_graph_cache_key) { + StepGraph graph; + graph.paged_tree_meta_arena.reset( + new uint8_t[1], std::default_delete()); + graph.paged_tree_key = TargetPagedTreeGraphKey{ + nullptr, nullptr, nullptr, 8, 4, 256, 4096, 16, 0}; + step_graph_free(graph); + CHECK(!graph.paged_tree_key); + CHECK(graph.paged_tree_meta_arena); + step_graph_destroy(graph); + CHECK(!graph.paged_tree_meta_arena); +} + TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { ggml_backend_t backend = ggml_backend_cpu_init(); CHECK(backend != nullptr);