perf(draft): pack dynamic-convolution projections across lanes

.github/workflows/ci.yml

@@ -213,6 +213,7 @@ jobs:
             test_ds4_mix_registry_teardown
             test_model_smoke
             test_batched_gdn
+            test_gdn_transition_journal
             test_concat_transpose
           )
           if [[ "$RUN_SPARSE_ATTENTION_TEST" == "true" ]]; then
@@ -233,7 +234,7 @@ jobs:
       - name: Run concurrent-serving kernel tests
         run: |
           ctest --test-dir server/build --output-on-failure \
-            -R '^(test_model_smoke\.PagedAttention\.|batched_gdn$|concat_transpose$)' --no-tests=error
+            -R '^(test_model_smoke\.PagedAttention\.|batched_gdn$|gdn_transition_journal_preflight$|concat_transpose$)' --no-tests=error
 
       # Optional model-backed end-to-end smoke (real spec-decode on the 3090),
       # disabled by default because it builds dflash_server and lazy-loads the

docs/handoffs/pr651-draft-batching-performance.md

@@ -0,0 +1,73 @@
+# PR 651 draft batching performance follow-ups
+
+## 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<DraftKvAppendLane> & 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.

server/CMakeLists.txt

@@ -455,6 +455,7 @@ add_library(dflash_common STATIC
     src/common/domino_head.cpp
     src/common/dspark_head.cpp
     src/common/dflash2_head.cpp
+    src/common/dflash2_batch.cpp
     src/common/target_shard_ipc.cpp
     src/common/target_shard_ipc_daemon.cpp
     src/common/dflash_feature_ring.cpp
@@ -464,6 +465,7 @@ add_library(dflash_common STATIC
     src/common/pflash_drafter_ipc.cpp
     src/common/dflash_draft_graph.cpp
     src/common/dflash_draft_kv.cpp
+    src/common/draft_swa.cpp
     src/common/dflash_spec_decode.cpp
     src/common/concurrency/paged_kv_pool.cpp
     src/qwen35/concurrency/qwen35_slot_manager.cpp
@@ -1435,6 +1437,51 @@ if(DFLASH27B_TESTS)
             ${CMAKE_CURRENT_SOURCE_DIR}/test)
         list(APPEND _raw_unit_test_targets test_seq_engine_contract)
     endif()
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ddtree_path.cpp")
+        # Pure host-side accepted-path/pending-token contract tests.
+        add_executable(test_ddtree_path
+            test/test_ddtree_path.cpp
+            src/common/ddtree.cpp)
+        target_include_directories(test_ddtree_path PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/src)
+        list(APPEND _raw_unit_test_targets test_ddtree_path)
+    endif()
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_chain_spec_shapes.cpp")
+        # Pure host-side fixed-chain prefix, bucket, and EOS contracts.
+        add_executable(test_chain_spec_shapes
+            test/test_chain_spec_shapes.cpp)
+        target_include_directories(test_chain_spec_shapes PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/src
+            ${CMAKE_CURRENT_SOURCE_DIR}/test)
+        list(APPEND _raw_unit_test_targets test_chain_spec_shapes)
+    endif()
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash2_selector_validation.cpp")
+        # Pure host-side selector metadata/layout validation: no GPU.
+        add_executable(test_dflash2_selector_validation
+            test/test_dflash2_selector_validation.cpp)
+        target_include_directories(test_dflash2_selector_validation PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/src
+            ${CMAKE_CURRENT_SOURCE_DIR}/test)
+        list(APPEND _raw_unit_test_targets test_dflash2_selector_validation)
+    endif()
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_swa.cpp")
+        # Pure host-side draft SWA metadata/override policy tests.
+        add_executable(test_draft_swa test/test_draft_swa.cpp)
+        target_include_directories(test_draft_swa PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/src)
+        target_link_libraries(test_draft_swa PRIVATE dflash_common)
+    endif()
+    if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR
+        DFLASH27B_GPU_BACKEND STREQUAL "hip") AND
+       EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_swa_multilane.cpp")
+        # Opt-in real-model proof for post-window lane masks and packed draft parity.
+        add_executable(test_draft_swa_multilane test/test_draft_swa_multilane.cpp)
+        target_include_directories(test_draft_swa_multilane PRIVATE
+            ${DFLASH27B_SRC_INCLUDE_DIRS})
+        target_link_libraries(test_draft_swa_multilane PRIVATE
+            dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET})
+        list(APPEND _raw_unit_test_targets test_draft_swa_multilane)
+    endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_batch_plan.cpp")
         # Pure-host tests for model-neutral token-budget/FIFO planning.
         add_executable(test_seq_batch_plan test/test_seq_batch_plan.cpp)
@@ -1774,6 +1821,7 @@ if(DFLASH27B_TESTS)
         test_model_smoke)
     set(_new_cppunit_test_targets
         test_platform_compat
+        test_draft_swa
         test_feature_gate
         test_qwen35_split_tree_guard
         test_recurrent_snapshot
@@ -1935,6 +1983,23 @@ if(DFLASH27B_TESTS)
             add_dependencies(check test_batched_gdn)
         endif()
     endif()
+    if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR
+        DFLASH27B_GPU_BACKEND STREQUAL "hip")
+       AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gdn_replay_log.cpp")
+        dflash_add_ggml_gpu_executable(
+            test_gdn_replay_log
+            test/test_gdn_replay_log.cpp)
+        add_test(NAME gdn_replay_log COMMAND test_gdn_replay_log)
+        add_test(
+            NAME gdn_replay_log_gallocr
+            COMMAND test_gdn_replay_log --gallocr-only)
+        add_test(
+            NAME gdn_replay_log_preflight
+            COMMAND test_gdn_replay_log --preflight-only)
+        if(TARGET check)
+            add_dependencies(check test_gdn_replay_log)
+        endif()
+    endif()
     if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR
         DFLASH27B_GPU_BACKEND STREQUAL "hip")
        AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_concat_transpose.cpp")

server/deps/llama.cpp/ggml/include/ggml-cuda.h

@@ -99,6 +99,50 @@ GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void);
 GGML_BACKEND_API bool ggml_backend_cuda_topk_rows(const struct ggml_tensor * logits, int k,
                                                   float * probs_out, int32_t * ids_out);
 
+// Batched concurrent-tree commit. Validation is fail-closed before any kernel
+// launches; all layer replay logs and convolution windows commit on one device
+// synchronization.
+GGML_BACKEND_API bool ggml_backend_cuda_gdn_replay_log_commit_many(
+        const struct ggml_tensor * const * replay_logs,
+        struct ggml_tensor * const * states,
+        const struct ggml_tensor * const * conv_inputs,
+        struct ggml_tensor * const * conv_states,
+        int n_layers,
+        const struct ggml_tensor * accepted_prefixes,
+        const struct ggml_tensor * active_slot_ids);
+
+// Promote accepted packed-tree K/V scratch rows into pager-owned rows.
+GGML_BACKEND_API bool ggml_backend_cuda_tree_cache_commit_many(
+        struct ggml_tensor * const * caches, int n_caches,
+        const struct ggml_tensor * commit_rows,
+        const struct ggml_tensor * active_slot_ids,
+        int tree_scratch_base, int tree_scratch_stride);
+
+// Promote accepted BF16 tree feature rows into slot-local feature rings.
+GGML_BACKEND_API bool ggml_backend_cuda_tree_feature_commit(
+        const struct ggml_tensor * source, struct ggml_tensor * destination,
+        const struct ggml_tensor * destination_rows);
+
+// Validate every packed-tree destination before changing any cache. Once the
+// first kernel launches, a device failure is fatal because fallback cannot
+// recover from a partially committed state.
+GGML_BACKEND_API bool ggml_backend_cuda_tree_commit_transaction(
+        struct ggml_tensor * const * caches,
+        int n_caches,
+        const struct ggml_tensor * feature_source,
+        struct ggml_tensor * feature_destination,
+        const struct ggml_tensor * feature_destination_rows,
+        const struct ggml_tensor * const * replay_logs,
+        struct ggml_tensor * const * states,
+        const struct ggml_tensor * const * conv_inputs,
+        struct ggml_tensor * const * conv_states,
+        int n_layers,
+        const struct ggml_tensor * commit_rows,
+        const struct ggml_tensor * accepted_prefixes,
+        const struct ggml_tensor * active_slot_ids,
+        int tree_scratch_base,
+        int tree_scratch_stride);
+
 // Attach learned per-expert decode tables to a mixed-precision tensor. The
 // host variants copy the tables to the device that owns `base`. Call the
 // matching unregister function before releasing the tensor's backing buffer.

server/deps/llama.cpp/ggml/include/ggml.h

@@ -2502,6 +2502,7 @@ extern "C" {
     // prefill chunks can attend the paged pool causally. A negative position
     // marks a padding row. NULL keeps the decode semantics (full cached
     // length per row).
+    //
     GGML_API struct ggml_tensor * ggml_paged_attn_ext(
             struct ggml_context * ctx,
             struct ggml_tensor  * q,
@@ -2515,6 +2516,39 @@ extern "C" {
             int                   block_size,
             int                   max_kv_seq_len);
 
+    // Packed tree verification over the same paged K/V pool.
+    // In a pure-tree batch, queries are flattened sequence-major and tree
+    // sequence s occupies rows [s*tree_width, (s+1)*tree_width). In a mixed
+    // AR/tree batch, the compact AR rows come first. Tree sequence s starts at
+    // ar_rows + s*tree_width, where
+    // ar_rows = q_rows - n_tree_seq*tree_width. parent_ids is tree-local,
+    // contiguous I32 [tree_width, n_tree_seq] (root parent -1), and tree_sizes
+    // is contiguous I32 [n_tree_seq]. active_slot_ids is required per query row;
+    // it selects the physical block-table column and scratch slab. Each live
+    // query attends its complete committed prefix from the block table plus
+    // its own candidate node and ancestors from physical K/V rows
+    // tree_scratch_base + slot*tree_scratch_stride + node. Siblings and rows
+    // at or beyond tree_sizes[s] are excluded. query_positions may describe
+    // compact autoregressive rows in a mixed AR/tree batch; tree rows ignore
+    // it and read the full committed prefix. Pure tree batches pass NULL.
+    // tree_width is derived from parent_ids.
+    GGML_API struct ggml_tensor * ggml_paged_attn_ext_tree(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * block_table,
+            struct ggml_tensor  * kv_seq_lens,
+            struct ggml_tensor  * active_slot_ids,
+            struct ggml_tensor  * query_positions,
+            float                 scale,
+            int                   block_size,
+            int                   max_kv_seq_len,
+            struct ggml_tensor  * parent_ids,
+            struct ggml_tensor  * tree_sizes,
+            int                   tree_scratch_base,
+            int                   tree_scratch_stride);
+
     // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse.
     // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse).
     // Used for KV cache rotation in TurboQuant quantization types (TQ3_0).
@@ -2938,6 +2972,15 @@ extern "C" {
             struct ggml_tensor * tensor,
             bool                 skip_intermediate);
 
+    // CUDA/HIP fixed-chain replay log in compact F32 [J,H,T,B] layout:
+    // scalar gate J=2*S_v+1 stores [g | k | delta], while KDA J=3*S_v
+    // stores [g[S_v] | k | delta]. Delta is captured after the
+    // state-dependent reduction. The returned tensor owns a compact copy so
+    // graph allocation can release the much larger GDN result after capture.
+    GGML_API struct ggml_tensor * ggml_gated_delta_net_capture_replay_log(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * tensor);
+
     // dflash extension: let the kernel derive the gates from the raw
     // projections instead of graph-side sigmoid/softplus ops:
     //   beta_val = sigmoid(beta_raw)

server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -478,11 +478,13 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
             // implements the original mode and asserts if one reaches it.
             return ggml_get_op_params_i32(op, 0) == 0;
         case GGML_OP_GATED_DELTA_NET:
-            // The Specla GDN variant (op param 2 == 1) is stateful via HLD and is
-            // only supported by CUDA/HIP. Raw-gate mode is also CUDA/HIP-only:
-            // the CPU kernel expects beta/g to have already been transformed.
+            // The CPU kernel supports in-place and active-slot recurrence,
+            // but not tree parents, persistent intermediate storage, raw
+            // gates or SpecLA state.
             return ggml_get_op_params_i32(op, 2) != 1 &&
-                   ggml_get_op_params_i32(op, 10) == 0 && op->src[9] == nullptr;
+                   ggml_get_op_params_i32(op, 10) == 0 &&
+                   op->src[6] == nullptr && op->src[7] == nullptr &&
+                   op->src[9] == nullptr;
         case GGML_OP_OUT_PROD:
             return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
                 src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;