From 312beb19c5b680111a0a21e2af9aa7dae49531ee Mon Sep 17 00:00:00 2001
From: berto <bregliaumberto.02@gmail.com>
Date: Mon, 31 Aug 2026 10:43:04 +0000
Subject: [PATCH] perf(qwen35): retain speculation during paged prefill

---
 server/src/common/step_graph.h                |   8 +-
 server/src/internal.h                         |  20 +-
 .../qwen35/concurrency/qwen35_seq_engine.cpp  | 185 +++++++++--
 server/src/qwen35/graph_builders.cpp          | 121 +++++++-
 server/src/qwen35/graph_builders.h            |  36 ++-
 server/src/qwen35/qwen35_target_graph.cpp     |  99 ++++--
 server/test/test_gdn_replay_log.cpp           |  34 ++-
 server/test/test_generate.cpp                 | 286 +++++++++++++++++-
 server/test/test_paged_attention.cpp          |  11 +-
 server/test/test_recurrent_snapshot.cpp       |  30 +-
 10 files changed, 722 insertions(+), 108 deletions(-)

diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h
index b774cd5ed..6f0416f31 100644
--- a/server/src/common/step_graph.h
+++ b/server/src/common/step_graph.h
@@ -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;
@@ -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;
diff --git a/server/src/internal.h b/server/src/internal.h
index 07d6b89a0..6b9e6cda5 100644
--- a/server/src/internal.h
+++ b/server/src/internal.h
@@ -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
@@ -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.
@@ -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;
@@ -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.
diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp
index 002cfdca8..61787a38d 100644
--- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp
+++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp
@@ -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;
     }
@@ -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;
     }
@@ -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;
@@ -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);
@@ -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(
@@ -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;
         }
     }
@@ -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;
@@ -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;
     }
@@ -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());
@@ -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(
@@ -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) {
@@ -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;
         }
@@ -945,11 +1048,11 @@ 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;
@@ -957,15 +1060,37 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec(
     }
     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;
diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp
index f59d4af31..a7e588484 100644
--- a/server/src/qwen35/graph_builders.cpp
+++ b/server/src/qwen35/graph_builders.cpp
@@ -46,18 +46,56 @@ bool detail::target_graph_capacity_for_parallel_segments(
 bool detail::target_paged_tree_graph_capacity(
         int tree_width,
         int n_tree_seqs,
-        size_t & capacity) {
+        size_t & capacity,
+        int n_recurrent_segments) {
     static constexpr int tree_buckets[] = {
         1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64,
     };
     if (tree_width < 2 || tree_width > 16 ||
         std::find(std::begin(tree_buckets), std::end(tree_buckets),
                   n_tree_seqs) == std::end(tree_buckets) ||
+        n_recurrent_segments < 1 || n_recurrent_segments > 64 ||
         (int64_t)tree_width * n_tree_seqs > INT32_MAX) {
         return false;
     }
     return target_graph_capacity_for_parallel_segments(
-        n_tree_seqs, capacity);
+        std::max(n_tree_seqs, n_recurrent_segments), capacity);
+}
+
+bool detail::validate_target_paged_tree_prefix(
+        const TargetCache & cache,
+        int n_prefill_tokens,
+        const QwenPrefillSegment * prefill_segments,
+        int n_prefill_segments,
+        int mapped_ar_seqs,
+        int & n_direct_rows) {
+    n_direct_rows = 0;
+    if (n_prefill_tokens < 0 || n_prefill_segments < 0 ||
+        n_prefill_segments > 64 ||
+        mapped_ar_seqs < 0 || mapped_ar_seqs > cache.n_seq_slots ||
+        (n_prefill_tokens > 0) !=
+            (prefill_segments != nullptr && n_prefill_segments > 0)) {
+        return false;
+    }
+
+    int64_t prefill_total = 0;
+    for (int i = 0; i < n_prefill_segments; ++i) {
+        const QwenPrefillSegment & segment = prefill_segments[i];
+        if (segment.n_tokens < 1 ||
+            segment.token_offset != prefill_total ||
+            segment.seq_slot < 0 ||
+            segment.seq_slot >= cache.n_seq_slots) {
+            return false;
+        }
+        prefill_total += segment.n_tokens;
+        if (prefill_total > INT32_MAX) return false;
+    }
+    if (prefill_total != n_prefill_tokens ||
+        prefill_total + mapped_ar_seqs > INT32_MAX) {
+        return false;
+    }
+    n_direct_rows = static_cast<int>(prefill_total) + mapped_ar_seqs;
+    return true;
 }
 
 bool detail::validate_target_paged_tree_layout(
@@ -816,10 +854,17 @@ bool build_target_step_paged_tree(
     int tree_scratch_base,
     int tree_scratch_stride,
     int kq_stride_pad,
-    int mapped_ar_seqs) {
+    int mapped_ar_seqs,
+    int n_prefill_tokens,
+    const QwenPrefillSegment * prefill_segments,
+    int n_prefill_segments,
+    int n_logits_rows) {
     (void)kq_stride_pad;
 
-    if (mapped_ar_seqs < 0 || mapped_ar_seqs > cache.n_seq_slots) {
+    int n_direct_rows = 0;
+    if (!detail::validate_target_paged_tree_prefix(
+            cache, n_prefill_tokens, prefill_segments,
+            n_prefill_segments, mapped_ar_seqs, n_direct_rows)) {
         step_graph_free(sg);
         return false;
     }
@@ -842,10 +887,32 @@ bool build_target_step_paged_tree(
     const int paged_launch_kv_len = static_cast<int>(
         std::min<int64_t>(((requested + 255) / 256) * 256,
                           logical_capacity));
+    const int64_t tree_rows64 = (int64_t)tree_width * n_tree_seqs;
+    if (tree_rows64 > INT32_MAX - n_direct_rows) {
+        step_graph_free(sg);
+        return false;
+    }
+    const int tree_rows = static_cast<int>(tree_rows64);
+    const int n_tokens = n_direct_rows + tree_rows;
+    const int fixed_logits_rows = mapped_ar_seqs + tree_rows;
+    if (n_logits_rows < 0 || n_logits_rows > n_tokens ||
+        (n_logits_rows > 0 &&
+         (n_logits_rows < fixed_logits_rows ||
+          n_logits_rows > fixed_logits_rows + n_prefill_segments))) {
+        step_graph_free(sg);
+        return false;
+    }
+
+    std::vector<int> prefill_shape;
+    prefill_shape.reserve(static_cast<size_t>(2 * n_prefill_segments));
+    for (int i = 0; i < n_prefill_segments; ++i) {
+        prefill_shape.push_back(prefill_segments[i].seq_slot);
+        prefill_shape.push_back(prefill_segments[i].n_tokens);
+    }
     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,
+        mapped_ar_seqs, n_logits_rows, prefill_shape,
     };
     if (sg.paged_tree_key && *sg.paged_tree_key == graph_key) {
         return true;
@@ -853,11 +920,16 @@ bool build_target_step_paged_tree(
     step_graph_free(sg);
 
     size_t graph_capacity = 0;
+    const int n_recurrent_segments =
+        n_prefill_segments + (mapped_ar_seqs > 0 ? 1 : 0) + 1;
     if (!detail::target_paged_tree_graph_capacity(
-            tree_width, n_tree_seqs, graph_capacity)) {
+            tree_width, n_tree_seqs, graph_capacity,
+            n_recurrent_segments)) {
         return false;
     }
-    const int n_tokens = mapped_ar_seqs + tree_width * n_tree_seqs;
+    if (n_prefill_segments > 0) {
+        graph_capacity = std::max<size_t>(graph_capacity, 32768);
+    }
     const int n_mapped_seqs = mapped_ar_seqs + n_tree_seqs;
 
     ggml_init_params ip{};
@@ -873,8 +945,19 @@ bool build_target_step_paged_tree(
 
     // Salt graph addresses by the stable bucket shape so captured graphs for
     // different T/S buckets never alias in ggml-cuda's topology cache.
+    int shape_salt = 0;
+    if (!prefill_shape.empty() || n_logits_rows > 0) {
+        uint64_t shape_hash = 1469598103934665603ull;
+        const auto hash_shape = [&](int value) {
+            shape_hash ^= static_cast<uint32_t>(value);
+            shape_hash *= 1099511628211ull;
+        };
+        hash_shape(n_logits_rows);
+        for (int value : prefill_shape) hash_shape(value);
+        shape_salt = static_cast<int>(shape_hash % 64);
+    }
     for (int i = 0; i < tree_width + n_tree_seqs +
-                        mapped_ar_seqs + n_tokens; ++i) {
+                        mapped_ar_seqs + n_tokens + shape_salt; ++i) {
         (void)ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1);
     }
 
@@ -892,12 +975,20 @@ bool build_target_step_paged_tree(
         ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_mapped_seqs);
     sg.paged_query_seq_ids =
         ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens);
-    if (mapped_ar_seqs > 0) {
+    if (n_direct_rows > 0) {
         sg.paged_query_positions =
             ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens);
     }
     sg.kv_write_rows = ggml_new_tensor_2d(
         sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv);
+    if (n_logits_rows > 0) {
+        sg.logits_row_indices =
+            ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_logits_rows);
+    }
+    if (n_direct_rows > 0 && cache.target_feat) {
+        sg.target_feat_rows =
+            ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_direct_rows);
+    }
 
     const struct NamedInput {
         ggml_tensor * tensor;
@@ -912,6 +1003,8 @@ bool build_target_step_paged_tree(
         {sg.paged_query_seq_ids, "paged_query_seq_ids"},
         {sg.paged_query_positions, "paged_query_positions"},
         {sg.kv_write_rows, "kv_write_rows"},
+        {sg.logits_row_indices, "logits_row_indices"},
+        {sg.target_feat_rows, "target_feat_rows"},
     };
     for (const NamedInput & input : inputs) {
         if (!input.tensor) continue;
@@ -937,6 +1030,11 @@ bool build_target_step_paged_tree(
     gi.state_slot_ids = sg.state_slot_ids;
     gi.paged_query_seq_ids = sg.paged_query_seq_ids;
     gi.paged_query_positions = sg.paged_query_positions;
+    gi.logits_row_indices = sg.logits_row_indices;
+    gi.target_feat_rows = sg.target_feat_rows;
+    gi.prefill_segments = prefill_segments;
+    gi.n_prefill_segments = n_prefill_segments;
+    gi.n_prefill_tokens = n_prefill_tokens;
     gi.n_seqs = n_tree_seqs;
     gi.mapped_ar_seqs = mapped_ar_seqs;
     gi.paged_max_kv_len = paged_launch_kv_len;
@@ -949,7 +1047,8 @@ bool build_target_step_paged_tree(
     sg.logits = go.logits;
     sg.delta_captures = std::move(go.delta_captures);
     sg.tree_features = go.tree_features;
-    if (!sg.tree_features || sg.delta_captures.empty()) {
+    if (!sg.tree_features || sg.tree_features->ne[1] != tree_rows ||
+        sg.delta_captures.empty()) {
         return false;
     }
     ggml_set_output(sg.logits);
@@ -978,7 +1077,7 @@ bool build_target_step_paged_tree(
     sg.commit_rows = ggml_new_tensor_2d(
         sg.commit_ctx, GGML_TYPE_I64, tree_width, n_tree_seqs);
     sg.feature_commit_rows = ggml_new_tensor_1d(
-        sg.commit_ctx, GGML_TYPE_I32, n_tokens);
+        sg.commit_ctx, GGML_TYPE_I32, tree_rows);
     ggml_set_name(sg.accepted_prefixes, "accepted_prefixes");
     ggml_set_name(sg.commit_slot_ids, "commit_slot_ids");
     ggml_set_name(sg.commit_rows, "commit_rows");
diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h
index 4ab9f2266..30a796984 100644
--- a/server/src/qwen35/graph_builders.h
+++ b/server/src/qwen35/graph_builders.h
@@ -38,7 +38,20 @@ bool target_graph_capacity_for_parallel_segments(
 bool target_paged_tree_graph_capacity(
     int tree_width,
     int n_tree_seqs,
-    size_t & capacity);
+    size_t & capacity,
+    int n_recurrent_segments = 1);
+
+// Validate the positioned durable prefix of a packed-tree graph. Prefill
+// segments are dense and ordered; mapped AR rows follow them and the tree is
+// the fixed-width suffix. The returned direct row count is useful to callers
+// that allocate positioned-row metadata.
+bool validate_target_paged_tree_prefix(
+    const TargetCache & cache,
+    int n_prefill_tokens,
+    const QwenPrefillSegment * prefill_segments,
+    int n_prefill_segments,
+    int mapped_ar_seqs,
+    int & n_direct_rows);
 
 // Model-free validation shared by the packed-tree builder and its shape
 // tests. paged_max_kv_len is a logical launch bound and may exceed the
@@ -68,6 +81,8 @@ inline bool target_paged_tree_uploads_ready(const StepGraph & sg) {
            allocated(sg.paged_query_seq_ids) &&
            (!sg.paged_query_positions ||
             allocated(sg.paged_query_positions)) &&
+           (!sg.logits_row_indices || allocated(sg.logits_row_indices)) &&
+           (!sg.target_feat_rows || allocated(sg.target_feat_rows)) &&
            allocated(sg.kv_write_rows);
 }
 
@@ -207,13 +222,12 @@ bool build_target_step_tree(
     const SpecLAHLDSchedule * specla_hld = nullptr);
 
 // Packed fixed-chain verify over a paged multi-slot cache. Tokens are
-// flattened after an optional compact one-token AR prefix as
-// [mapped_ar_seqs + tree_width*n_tree_seqs]. n_tree_seqs is a stable graph-
-// bucket width; inactive trees use tree_size=0 and dead/safe row mappings. In
-// particular, state_slot_ids padding must map to a valid harmless slot
-// (normally 0), while active/paged sequence IDs may use -1. Speculative K/V is
-// written into per-slot scratch slabs; recurrent transitions and target
-// features are exposed for post-verification promotion.
+// flattened as [prefill segments][mapped AR rows][tree rows]. n_tree_seqs is
+// a stable graph-bucket width; inactive trees use tree_size=0 and dead/safe
+// row mappings. state_slot_ids padding must map to a valid harmless slot
+// (normally 0), while active/paged sequence IDs may use -1. Prefill/AR rows
+// update durable state directly. Tree K/V, recurrence, and target features
+// remain scratch until accepted-prefix promotion.
 bool build_target_step_paged_tree(
     StepGraph & sg,
     const TargetWeights & w,
@@ -225,7 +239,11 @@ bool build_target_step_paged_tree(
     int tree_scratch_base,
     int tree_scratch_stride,
     int kq_stride_pad = KQ_MASK_PAD,
-    int mapped_ar_seqs = 0);
+    int mapped_ar_seqs = 0,
+    int n_prefill_tokens = 0,
+    const QwenPrefillSegment * prefill_segments = nullptr,
+    int n_prefill_segments = 0,
+    int n_logits_rows = 0);
 
 // LM-head projection: project draft hidden states through the target output matrix.
 bool build_lm_head_projection_step(
diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp
index 7e7c3b613..68a23add8 100644
--- a/server/src/qwen35/qwen35_target_graph.cpp
+++ b/server/src/qwen35/qwen35_target_graph.cpp
@@ -1319,6 +1319,7 @@ static ggml_tensor * build_full_attn_block(
     auto paged_read = [&](ggml_tensor * q, int launch_kv_len,
                           ggml_tensor * row_seq_ids,
                           ggml_tensor * row_positions,
+                          bool tree_read,
                           bool dense_token_layout) {
         // max_kv_seq_len sizes the logical partition grid. Paged serving can
         // map that logical range onto a much smaller physical K/V pool, so
@@ -1339,7 +1340,7 @@ static ggml_tensor * build_full_attn_block(
         const int64_t padded = ((requested + 255) / 256) * 256;
         const int launch_len =
             (int)std::min<int64_t>(padded, logical_capacity);
-        ggml_tensor * out = paged_tree
+        ggml_tensor * out = tree_read
             ? ggml_paged_attn_ext_tree(
                 ctx, q, cache_k, cache_v, paged_block_table,
                 paged_kv_seq_lens, row_seq_ids, row_positions, kq_scale,
@@ -1361,15 +1362,53 @@ static ggml_tensor * build_full_attn_block(
         // ── Packed concurrent tree verify. Every query row selects its
         // physical sequence/scratch slab. The paged kernel combines the
         // committed block-table prefix with only this node's ancestor chain.
-        // A mixed graph uses causal positions for the compact AR prefix and
-        // -1 for the tree tail; a pure tree keeps positions absent.
-        ggml_tensor * Qfa = q_segment(0, n_tokens);
-        if (q_fa_out) *q_fa_out = Qfa;
         const int launch_kv_len = paged_max_kv_len > 0
             ? paged_max_kv_len : kv_start + n_tokens;
-        attn = paged_read(Qfa, launch_kv_len,
-                          paged_query_seq_ids, paged_query_positions,
-                          /*dense_token_layout=*/n_tokens > 1);
+        const int64_t tree_rows64 =
+            paged_tree_parent_ids->ne[0] * paged_tree_parent_ids->ne[1];
+        GGML_ASSERT(tree_rows64 > 0 && tree_rows64 <= n_tokens);
+        const int tree_rows = static_cast<int>(tree_rows64);
+        const int direct_rows = n_tokens - tree_rows;
+
+        if (direct_rows > 0) {
+            // Keep durable prefill/AR rows on the ordinary positioned-row
+            // operator. Folding them into the tree op lets the virtual tree
+            // tail change its partition topology at context boundaries, which
+            // changes the durable forward before any tree transaction runs.
+            GGML_ASSERT(paged_query_positions);
+            ggml_tensor * direct_seq_ids = ggml_view_1d(
+                ctx, paged_query_seq_ids, direct_rows, 0);
+            ggml_tensor * direct_positions = ggml_view_1d(
+                ctx, paged_query_positions, direct_rows, 0);
+            ggml_tensor * Qdirect = q_segment(0, direct_rows);
+            ggml_tensor * direct_attn = paged_read(
+                Qdirect, launch_kv_len, direct_seq_ids, direct_positions,
+                /*tree_read=*/false, /*dense_token_layout=*/true);
+
+            const size_t tree_row_offset =
+                static_cast<size_t>(direct_rows) *
+                paged_query_seq_ids->nb[0];
+            ggml_tensor * tree_seq_ids = ggml_view_1d(
+                ctx, paged_query_seq_ids, tree_rows, tree_row_offset);
+            ggml_tensor * Qtree = q_segment(direct_rows, tree_rows);
+            ggml_tensor * tree_attn = paged_read(
+                Qtree, launch_kv_len, tree_seq_ids,
+                /*row_positions=*/nullptr,
+                /*tree_read=*/true, /*dense_token_layout=*/true);
+
+            attn = ggml_concat(ctx, direct_attn, tree_attn, 2);
+            if (q_fa_out) {
+                *q_fa_out = ggml_concat(ctx, Qdirect, Qtree, 1);
+            }
+        } else {
+            ggml_tensor * Qfa = q_segment(0, n_tokens);
+            if (q_fa_out) *q_fa_out = Qfa;
+            attn = paged_read(
+                Qfa, launch_kv_len, paged_query_seq_ids,
+                /*row_positions=*/nullptr,
+                /*tree_read=*/true,
+                /*dense_token_layout=*/n_tokens > 1);
+        }
     } else if (ragged) {
         // ── Ragged concurrent step: prefill chunk rows and decode rows all
         // read the pool through one call, each row clamped to its own
@@ -1383,6 +1422,7 @@ static ggml_tensor * build_full_attn_block(
                                                        : kv_start + n_tokens;
         attn = paged_read(Qfa, launch_kv_len,
                           paged_query_seq_ids, paged_query_positions,
+                          /*tree_read=*/false,
                           /*dense_token_layout=*/n_tokens > 1);
     } else if (paged_block_table) {
         ggml_tensor * Qfa = q_segment(0, n_tokens);
@@ -1405,6 +1445,7 @@ static ggml_tensor * build_full_attn_block(
         const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_len;
         attn = paged_read(
             Qfa, launch_kv_len, active_slot_ids, /*row_positions=*/nullptr,
+            /*tree_read=*/false,
             /*dense_token_layout=*/active_slot_ids && n_tokens > 1);
         if (!active_slot_ids) {
             // The only non-mapped paged caller is classic single-token AR.
@@ -1545,9 +1586,9 @@ static ggml_tensor * build_delta_net_block(
                 (!cap->ssm_intermediate_states && !cap->conv_input));
     GGML_ASSERT(!active_slot_ids ||
                 (mapped_tree
-                     ? (!ragged && prefill_total == 0 &&
-                        n_tokens >= mapped_ar_seqs &&
-                        (n_tokens - mapped_ar_seqs) % n_seqs == 0 &&
+                     ? (n_tokens > prefill_total + mapped_ar_seqs &&
+                        (n_tokens - prefill_total - mapped_ar_seqs) %
+                                n_seqs == 0 &&
                         active_slot_ids->ne[0] ==
                             mapped_ar_seqs + n_seqs &&
                         state_slot_ids->ne[0] ==
@@ -1558,7 +1599,8 @@ static ggml_tensor * build_delta_net_block(
         GGML_ASSERT(n_seqs == 1);
         GGML_ASSERT(prefill_total == 0 || prefill_total == n_tokens);
     }
-    GGML_ASSERT(!ragged || (!cap && !parent_ids));
+    GGML_ASSERT(!ragged || !cap || mapped_tree);
+    GGML_ASSERT(!ragged || !parent_ids || mapped_tree);
 
     // Row slices of stacked projections are strided for multi-token inputs.
     // Materialize only the small beta/alpha slices; qkv keeps its explicit
@@ -1647,7 +1689,7 @@ static ggml_tensor * build_delta_net_block(
         ggml_tensor * state_ids;
     };
     std::vector<DeltaSeg> segs;
-    segs.reserve((size_t)n_prefill_segments + 1);
+    segs.reserve((size_t)n_prefill_segments + 2);
     for (int i = 0; i < n_prefill_segments; ++i) {
         const QwenPrefillSegment & pf = prefill_segments[i];
         GGML_ASSERT(pf.seq_slot >= 0 &&
@@ -1669,11 +1711,11 @@ static ggml_tensor * build_delta_net_block(
                 ctx, active_slot_ids, mapped_ar_seqs, 0);
             ggml_tensor * ar_state = ggml_view_1d(
                 ctx, state_slot_ids, mapped_ar_seqs, 0);
-            segs.push_back({0, 1, mapped_ar_seqs, true, false,
+            segs.push_back({prefill_total, 1, mapped_ar_seqs, true, false,
                             conv_state, ssm_state, ar_active, ar_state});
         }
         const int tree_tokens = mapped_tree
-            ? (n_tokens - mapped_ar_seqs) / n_seqs : 1;
+            ? (n_tokens - prefill_total - mapped_ar_seqs) / n_seqs : 1;
         const size_t slot_offset =
             (size_t)mapped_ar_seqs * active_slot_ids->nb[0];
         ggml_tensor * segment_active = mapped_ar_seqs > 0
@@ -2584,14 +2626,35 @@ QwenGraphOutputs build_qwen35_graph(
                 ctx, feat_cat, capture_slices[(size_t)k], 0);
         }
         feat_cat = ggml_cont(ctx, feat_cat);
+        int direct_feature_rows = n_tokens;
+        int tree_feature_rows = 0;
         if (capture_tree_features) {
-            og_early.tree_features = ggml_cast(ctx, feat_cat, GGML_TYPE_BF16);
+            tree_feature_rows = in.tree_width * in.n_seqs;
+            GGML_ASSERT(tree_feature_rows > 0 &&
+                        tree_feature_rows <= n_tokens);
+            direct_feature_rows = n_tokens - tree_feature_rows;
+            ggml_tensor * tree_source = direct_feature_rows == 0
+                ? feat_cat
+                : ggml_view_2d(
+                      ctx, feat_cat, feat_cat->ne[0], tree_feature_rows,
+                      feat_cat->nb[1],
+                      (size_t)direct_feature_rows * feat_cat->nb[1]);
+            og_early.tree_features =
+                ggml_cast(ctx, tree_source, GGML_TYPE_BF16);
             ggml_set_output(og_early.tree_features);
             ggml_build_forward_expand(gf, og_early.tree_features);
-        } else {
+        }
+        if (capture_with_rows) {
+            GGML_ASSERT(direct_feature_rows > 0 &&
+                        in.target_feat_rows->ne[0] == direct_feature_rows);
+            ggml_tensor * direct_source = direct_feature_rows == n_tokens
+                ? feat_cat
+                : ggml_view_2d(
+                      ctx, feat_cat, feat_cat->ne[0], direct_feature_rows,
+                      feat_cat->nb[1], 0);
             ggml_build_forward_expand(
                 gf, ggml_set_rows(
-                        ctx, cache.target_feat, feat_cat,
+                        ctx, cache.target_feat, direct_source,
                         in.target_feat_rows));
         }
     }
diff --git a/server/test/test_gdn_replay_log.cpp b/server/test/test_gdn_replay_log.cpp
index 3a98df02c..d7f0afe43 100644
--- a/server/test/test_gdn_replay_log.cpp
+++ b/server/test/test_gdn_replay_log.cpp
@@ -812,6 +812,7 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
     constexpr int state_slots = 1;
     constexpr int conv_window = 2;
     constexpr int conv_channels = 3;
+    constexpr int feature_rows_count = tokens * sequences;
 
     ggml_init_params params{};
     params.mem_size = 256*1024;
@@ -822,15 +823,15 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
     ggml_tensor * cache =
         ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 4, 8, 1, 1);
     ggml_tensor * commit_rows =
-        ggml_new_tensor_2d(ctx, GGML_TYPE_I64, 1, sequences);
+        ggml_new_tensor_2d(ctx, GGML_TYPE_I64, tokens, sequences);
     ggml_tensor * active_slots =
         ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sequences);
     ggml_tensor * feature_source =
-        ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 4, sequences);
+        ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 4, feature_rows_count);
     ggml_tensor * feature_destination =
         ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 4, 4);
     ggml_tensor * feature_rows =
-        ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sequences);
+        ggml_new_tensor_1d(ctx, GGML_TYPE_I32, feature_rows_count);
     ggml_tensor * replay_log = ggml_new_tensor_4d(
         ctx, GGML_TYPE_F32, 2*state_size + 1,
         heads, tokens, sequences);
@@ -856,6 +857,9 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
         (size_t) ggml_nelements(cache), 1.25f);
     std::vector<uint16_t> feature_before(
         (size_t) ggml_nelements(feature_destination), 0x3f80);
+    // Model an already-durable direct-row feature. Tree promotion must leave
+    // it alone while copying only the accepted tree suffix into row zero.
+    std::fill(feature_before.begin() + 4, feature_before.begin() + 8, 0x4100);
     std::vector<float> state_before(
         (size_t) ggml_nelements(state), -2.5f);
     std::vector<float> conv_before(
@@ -866,9 +870,13 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
         0.0f);
     std::vector<uint16_t> feature_source_values(
         (size_t) ggml_nelements(feature_source), 0x4000);
-    const int64_t destination_row = 0;
+    std::fill(
+        feature_source_values.begin(),
+        feature_source_values.begin() + 4,
+        0x4200);
+    const int64_t destination_rows[] = {0, -1};
     const int32_t active_slot = 0;
-    const int32_t feature_row = 0;
+    const int32_t feature_row_map[] = {0, -1};
     const int32_t valid_accepted = 1;
     const int32_t invalid_accepted = tokens + 1;
 
@@ -894,11 +902,11 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
         conv_input, zeros.data(), 0,
         (size_t) ggml_nelements(conv_input)*sizeof(float));
     ggml_backend_tensor_set(
-        commit_rows, &destination_row, 0, sizeof(destination_row));
+        commit_rows, destination_rows, 0, sizeof(destination_rows));
     ggml_backend_tensor_set(
         active_slots, &active_slot, 0, sizeof(active_slot));
     ggml_backend_tensor_set(
-        feature_rows, &feature_row, 0, sizeof(feature_row));
+        feature_rows, feature_row_map, 0, sizeof(feature_row_map));
     ggml_backend_tensor_set(
         accepted, &valid_accepted, 0, sizeof(valid_accepted));
 
@@ -911,7 +919,7 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
     const bool committed = ggml_backend_cuda_tree_commit_transaction(
         caches, 1, feature_source, feature_destination, feature_rows,
         replay_logs, states, conv_inputs, conv_states, 1,
-        commit_rows, accepted, active_slots, 4, 1);
+        commit_rows, accepted, active_slots, 4, 2);
     std::vector<uint16_t> committed_feature(feature_before.size());
     std::vector<float> committed_state(state_before.size());
     std::vector<float> committed_conv(conv_before.size());
@@ -927,7 +935,13 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
     const bool success_path_ok = committed &&
         std::all_of(
             committed_feature.begin(), committed_feature.begin() + 4,
-            [](uint16_t value) { return value == 0x4000; }) &&
+            [](uint16_t value) { return value == 0x4200; }) &&
+        std::all_of(
+            committed_feature.begin() + 4, committed_feature.begin() + 8,
+            [](uint16_t value) { return value == 0x4100; }) &&
+        std::all_of(
+            committed_feature.begin() + 8, committed_feature.end(),
+            [](uint16_t value) { return value == 0x3f80; }) &&
         std::all_of(
             committed_state.begin(), committed_state.end(),
             [](float value) { return value == 0.0f; }) &&
@@ -952,7 +966,7 @@ bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) {
     const bool rejected = !ggml_backend_cuda_tree_commit_transaction(
         caches, 1, feature_source, feature_destination, feature_rows,
         replay_logs, states, conv_inputs, conv_states, 1,
-        commit_rows, accepted, active_slots, 4, 1);
+        commit_rows, accepted, active_slots, 4, 2);
 
     std::vector<float> cache_after(cache_before.size());
     std::vector<uint16_t> feature_after(feature_before.size());
diff --git a/server/test/test_generate.cpp b/server/test/test_generate.cpp
index 8d05860aa..57cc80918 100644
--- a/server/test/test_generate.cpp
+++ b/server/test/test_generate.cpp
@@ -10,6 +10,8 @@
 // Usage:
 //   test_generate <qwen35.gguf> <prompt_ids.bin> <n_gen> <out_ids.bin>
 //   test_generate --seq-engine-contract <qwen35.gguf> [slots]
+//   test_generate --seq-engine-mixed-spec-contract <qwen35.gguf>
+//                 <draft.gguf|safetensors> [slots]
 
 #include "dflash27b.h"
 #include "internal.h"
@@ -60,6 +62,14 @@ struct GenerateStepGraph {
     ggml_tensor *     logits    = nullptr;
 };
 
+class ContractQwen35Backend final : public Qwen35Backend {
+public:
+    using Qwen35Backend::Qwen35Backend;
+
+    const StepGraph & step_graph() const { return target_step_graph(); }
+    const TargetCache & cache() const { return target_cache(); }
+};
+
 // Build a fresh single-token forward graph. We rebuild per step so that
 // `kv_start` updates drive the correct KV cache slot. The graph is cheap to
 // rebuild — all the weights + KV cache stay persistent.
@@ -125,9 +135,236 @@ static bool write_generate_tokens(const std::string & path,
     return (bool)f;
 }
 
-static int run_seq_engine_contract(const char * gguf_path, int slots) {
-    if (slots < 2 || slots > 64) {
-        std::fprintf(stderr, "slots must be in [2, 64], got %d\n", slots);
+static std::vector<std::string> check_mixed_spec_prompt_tail(
+        ContractQwen35Backend & backend, SeqEngine & engine) {
+    std::vector<std::string> violations;
+    const auto require = [&](bool ok, const char * message) {
+        if (!ok) violations.emplace_back(message);
+        return ok;
+    };
+    const auto capture_recurrent_state = [&](int slot,
+                                             std::vector<uint8_t> & bytes) {
+        const TargetCache & cache = backend.cache();
+        if (slot < 0 || slot >= cache.n_seq_slots) return false;
+        bytes.clear();
+        const auto append_slot = [&](ggml_tensor * tensor) {
+            if (!tensor || ggml_nbytes(tensor) % cache.n_seq_slots != 0) {
+                return false;
+            }
+            const size_t slab = ggml_nbytes(tensor) / cache.n_seq_slots;
+            const size_t old_size = bytes.size();
+            bytes.resize(old_size + slab);
+            ggml_backend_tensor_get(
+                tensor, bytes.data() + old_size,
+                static_cast<size_t>(slot) * slab, slab);
+            return true;
+        };
+        for (ggml_tensor * tensor : cache.conv_state) {
+            if (!append_slot(tensor)) return false;
+        }
+        for (ggml_tensor * tensor : cache.ssm_state) {
+            if (!append_slot(tensor)) return false;
+        }
+        return !bytes.empty();
+    };
+
+    std::vector<int32_t> completing_prompt(1024);
+    std::vector<int32_t> remaining_prompt(1025);
+    for (size_t i = 0; i < completing_prompt.size(); ++i) {
+        completing_prompt[i] = 11 + static_cast<int32_t>(i % 4);
+    }
+    for (size_t i = 0; i < remaining_prompt.size(); ++i) {
+        remaining_prompt[i] = 21 + static_cast<int32_t>(i % 5);
+    }
+
+    struct CaseResult {
+        int32_t prompt_token = -1;
+        std::vector<uint8_t> recurrent_state;
+    };
+    uint64_t next_request_id = 100;
+    const auto run_case = [&](bool allow_speculation, CaseResult & output) {
+        const SamplerCfg greedy{};
+        std::vector<int> live_slots;
+        const auto retire_all = [&]() {
+            for (int slot : live_slots) engine.retire(slot);
+            live_slots.clear();
+        };
+        const auto admit = [&](const std::vector<int32_t> & prompt) {
+            const SeqEngine::AdmitResult admitted =
+                engine.admit(next_request_id++, prompt, greedy);
+            if (admitted.status !=
+                SeqEngine::AdmitResult::Status::admitted) {
+                return -1;
+            }
+            live_slots.push_back(admitted.slot);
+            return admitted.slot;
+        };
+        const auto decode_token = [](
+                const SeqEngine::StepResult & result, int slot,
+                int32_t & token) {
+            const auto it = std::find_if(
+                result.decode.begin(), result.decode.end(),
+                [slot](const SeqEngine::DecodeOutput & row) {
+                    return row.slot == slot;
+                });
+            if (it == result.decode.end() || it->failed) return false;
+            token = it->token;
+            return true;
+        };
+
+        const int established_a = admit({31, 32});
+        const int established_b = admit({41, 42});
+        if (!require(established_a >= 0 && established_b >= 0,
+                     "mixed-spec boundary fixture could not admit decoders")) {
+            retire_all();
+            return false;
+        }
+        SeqEngine::StepPlan establish_plan;
+        establish_plan.prefills.push_back({established_a, 2});
+        establish_plan.prefills.push_back({established_b, 2});
+        const SeqEngine::StepResult established = engine.step(establish_plan);
+        if (!require(established.ok() && established.prefills.size() == 2,
+                     "mixed-spec boundary fixture could not establish decoders")) {
+            retire_all();
+            return false;
+        }
+        int32_t token_a = -1;
+        int32_t token_b = -1;
+        for (const SeqEngine::PrefillOutput & row : established.prefills) {
+            if (row.slot == established_a) token_a = row.token;
+            if (row.slot == established_b) token_b = row.token;
+        }
+        if (!require(token_a >= 0 && token_b >= 0,
+                     "mixed-spec boundary decoder prefill did not complete")) {
+            retire_all();
+            return false;
+        }
+
+        const int completing = admit(completing_prompt);
+        const int remaining = admit(remaining_prompt);
+        if (!require(completing >= 0 && remaining >= 0,
+                     "mixed-spec boundary fixture could not admit prefills")) {
+            retire_all();
+            return false;
+        }
+
+        SeqEngine::StepPlan first_plan;
+        first_plan.decode.push_back(
+            {established_a, token_a, allow_speculation});
+        first_plan.decode.push_back(
+            {established_b, token_b, allow_speculation});
+        first_plan.prefills.push_back({completing, 512});
+        first_plan.prefills.push_back({remaining, 512});
+        const SeqEngine::StepResult first = engine.step(first_plan);
+        if (!require(first.ok() && first.prefills.size() == 2,
+                     "mixed-spec boundary first prefill chunks failed") ||
+            !require(decode_token(first, established_a, token_a) &&
+                         decode_token(first, established_b, token_b),
+                     "mixed-spec boundary first decode rows failed")) {
+            retire_all();
+            return false;
+        }
+
+        SeqEngine::StepPlan boundary_plan;
+        boundary_plan.decode.push_back(
+            {established_a, token_a, allow_speculation});
+        boundary_plan.decode.push_back(
+            {established_b, token_b, allow_speculation});
+        boundary_plan.prefills.push_back({completing, 512});
+        boundary_plan.prefills.push_back({remaining, 512});
+        const SeqEngine::StepResult boundary = engine.step(boundary_plan);
+        if (!require(boundary.ok(),
+                     "mixed-spec 1024-token boundary step failed")) {
+            retire_all();
+            return false;
+        }
+        const auto completing_row = std::find_if(
+            boundary.prefills.begin(), boundary.prefills.end(),
+            [completing](const SeqEngine::PrefillOutput & row) {
+                return row.slot == completing;
+            });
+        const auto remaining_row = std::find_if(
+            boundary.prefills.begin(), boundary.prefills.end(),
+            [remaining](const SeqEngine::PrefillOutput & row) {
+                return row.slot == remaining;
+            });
+        if (!require(
+                completing_row != boundary.prefills.end() &&
+                    completing_row->status ==
+                        SeqEngine::PrefillOutput::Status::completed,
+                "mixed-spec boundary completing prefill did not complete") ||
+            !require(
+                remaining_row != boundary.prefills.end() &&
+                    remaining_row->status ==
+                        SeqEngine::PrefillOutput::Status::advanced,
+                "mixed-spec boundary peer did not remain in prefill")) {
+            retire_all();
+            return false;
+        }
+        output.prompt_token = completing_row->token;
+        if (!require(capture_recurrent_state(
+                         completing, output.recurrent_state),
+                     "mixed-spec boundary recurrent-state capture failed")) {
+            retire_all();
+            return false;
+        }
+
+        if (allow_speculation) {
+            const StepGraph & graph = backend.step_graph();
+            const bool tree_graph = require(
+                graph.parent_ids && graph.logits_row_indices &&
+                    graph.argmax_tokens,
+                "mixed prefill did not retain the fixed-chain target graph");
+            if (tree_graph) {
+                const int tree_width =
+                    static_cast<int>(graph.parent_ids->ne[0]);
+                const int tree_lanes =
+                    static_cast<int>(graph.parent_ids->ne[1]);
+                const int logits_count =
+                    static_cast<int>(graph.logits_row_indices->ne[0]);
+                require(tree_width == 8 && tree_lanes == 2,
+                        "mixed prefill used the wrong fixed-chain W8 bucket");
+                require(logits_count == tree_width * tree_lanes + 1,
+                        "mixed prefill produced the wrong compact logits shape");
+                if (logits_count == tree_width * tree_lanes + 1) {
+                    std::vector<int32_t> logits_rows(
+                        static_cast<size_t>(logits_count), -1);
+                    ggml_backend_tensor_get(
+                        graph.logits_row_indices, logits_rows.data(), 0,
+                        sizeof(int32_t) * logits_rows.size());
+                    require(logits_rows[0] == 511,
+                            "mixed speculative prompt tail did not gather its final row");
+                    for (int row = 1; row < logits_count; ++row) {
+                        require(
+                            logits_rows[static_cast<size_t>(row)] ==
+                                1023 + row,
+                            "mixed speculative tree logits were not gathered "
+                            "after the direct prefix");
+                    }
+                }
+            }
+        }
+        retire_all();
+        return true;
+    };
+
+    CaseResult ordinary;
+    CaseResult speculative;
+    if (run_case(false, ordinary) && run_case(true, speculative)) {
+        require(speculative.prompt_token == ordinary.prompt_token,
+                "mixed speculation changed the 1024-token prompt-tail result");
+        require(speculative.recurrent_state == ordinary.recurrent_state,
+                "mixed speculation changed durable convolution/Gated DeltaNet state");
+    }
+    return violations;
+}
+
+static int run_seq_engine_contract(
+        const char * gguf_path, const char * draft_path, int slots) {
+    const int min_slots = draft_path ? 4 : 2;
+    if (slots < min_slots || slots > 64) {
+        std::fprintf(stderr, "slots must be in [%d, 64], got %d\n",
+                     min_slots, slots);
         return 2;
     }
 
@@ -140,13 +377,18 @@ static int run_seq_engine_contract(const char * gguf_path, int slots) {
 
     Qwen35Config cfg;
     cfg.target_path = gguf_path;
+    cfg.draft_path = draft_path;
+    cfg.draft_block_size = draft_path ? 8 : 0;
     cfg.device.gpu = 0;
-    cfg.device.max_ctx = 256;
+    cfg.device.max_ctx = draft_path ? 1152 : 256;
     cfg.draft_gpu = 0;
     cfg.paged_attention = true;
     cfg.max_concurrency = slots;
+    cfg.kv_pool_tokens = draft_path
+        ? static_cast<int64_t>(cfg.device.max_ctx) * slots
+        : 0;
 
-    Qwen35Backend backend(cfg);
+    ContractQwen35Backend backend(cfg);
     if (!backend.init()) {
         std::fprintf(stderr, "seq-engine backend init failed: %s\n",
                      dflash27b_last_error());
@@ -160,15 +402,21 @@ static int run_seq_engine_contract(const char * gguf_path, int slots) {
         return 1;
     }
 
-    const std::vector<std::string> violations =
-        check_seq_engine_contract(*engine);
+    std::vector<std::string> violations = check_seq_engine_contract(*engine);
+    if (violations.empty() && draft_path) {
+        std::vector<std::string> mixed_violations =
+            check_mixed_spec_prompt_tail(backend, *engine);
+        violations.insert(
+            violations.end(), mixed_violations.begin(),
+            mixed_violations.end());
+    }
     for (const std::string & violation : violations) {
         std::fprintf(stderr, "seq-engine contract: %s\n", violation.c_str());
     }
     if (!violations.empty()) return 1;
 
-    std::printf("seq-engine contract passed against Qwen35Backend (%d slots)\n",
-                slots);
+    std::printf("seq-engine %scontract passed against Qwen35Backend (%d slots)\n",
+                draft_path ? "mixed-spec " : "", slots);
     return 0;
 }
 
@@ -182,14 +430,28 @@ int main(int argc, char ** argv) {
             return 2;
         }
         const int slots = argc == 4 ? std::atoi(argv[3]) : 4;
-        return run_seq_engine_contract(argv[2], slots);
+        return run_seq_engine_contract(argv[2], nullptr, slots);
+    }
+    if (argc >= 2 &&
+        std::strcmp(argv[1], "--seq-engine-mixed-spec-contract") == 0) {
+        if (argc < 4 || argc > 5) {
+            std::fprintf(stderr,
+                "usage: %s --seq-engine-mixed-spec-contract "
+                "<qwen35.gguf> <draft.gguf|safetensors> [slots]\n",
+                argv[0]);
+            return 2;
+        }
+        const int slots = argc == 5 ? std::atoi(argv[4]) : 4;
+        return run_seq_engine_contract(argv[2], argv[3], slots);
     }
 
     if (argc < 5) {
         std::fprintf(stderr,
             "usage: %s <qwen35.gguf> <prompt_ids.bin> <n_gen> <out_ids.bin>\n"
-            "       %s --seq-engine-contract <qwen35.gguf> [slots]\n",
-            argv[0], argv[0]);
+            "       %s --seq-engine-contract <qwen35.gguf> [slots]\n"
+            "       %s --seq-engine-mixed-spec-contract <qwen35.gguf> "
+            "<draft.gguf|safetensors> [slots]\n",
+            argv[0], argv[0], argv[0]);
         return 2;
     }
     const char * gguf_path   = argv[1];
diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp
index 449a52ac5..be8650bdf 100644
--- a/server/test/test_paged_attention.cpp
+++ b/server/test/test_paged_attention.cpp
@@ -562,11 +562,14 @@ void run_mixed_tree_case() {
             7, 8, 9, 10, 11, 12, 13, 14,
         },
         {16},
-        2,
+        5,
     };
-    std::vector<int32_t> query_slots{1, 0};
+    // Two positioned ragged segments precede the speculative tree suffix.
+    // The operator treats every leading positioned row as durable causal
+    // work, regardless of whether the row came from prefill or one-token AR.
+    std::vector<int32_t> query_slots{1, 1, 1, 0, 0};
     query_slots.insert(query_slots.end(), 16, 2);
-    std::vector<int32_t> query_positions{7, 15};
+    std::vector<int32_t> query_positions{5, 6, 7, 14, 15};
     query_positions.insert(query_positions.end(), 16, -1);
     CHECK(run_case(backend, mixed_case, GGML_TYPE_F16, GGML_TYPE_F16,
                    &query_slots, &query_positions, &tree_metadata));
@@ -654,7 +657,7 @@ TEST_CASE(PagedAttention, PackedTreesMatchReference) {
     run_tree_case();
 }
 
-TEST_CASE(PagedAttention, CompactArAndFixedChainMatchReference) {
+TEST_CASE(PagedAttention, RaggedDirectPrefixAndFixedChainMatchReference) {
     run_mixed_tree_case();
 }
 
diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp
index ddb945450..6111feeca 100644
--- a/server/test/test_recurrent_snapshot.cpp
+++ b/server/test/test_recurrent_snapshot.cpp
@@ -54,6 +54,8 @@ TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_capacity_and_uploads) {
                    65, graph_capacity));
     CHECK(dflash::common::detail::target_paged_tree_graph_capacity(
               16, 16, graph_capacity) && graph_capacity == 32768);
+    CHECK(dflash::common::detail::target_paged_tree_graph_capacity(
+              8, 1, graph_capacity, 10) && graph_capacity == 32768);
     CHECK(!dflash::common::detail::target_paged_tree_graph_capacity(
                17, 16, graph_capacity));
 
@@ -139,6 +141,23 @@ TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) {
                 shape_cache, 8, 2, 4096, 48, 16));
             CHECK(!dflash::common::detail::validate_target_paged_tree_layout(
                 shape_cache, 8, 5, 4096, 32, 16));
+
+            const dflash::common::QwenPrefillSegment segments[] = {
+                {0, 3, 1},
+                {3, 2, 0},
+            };
+            int direct_rows = -1;
+            CHECK(dflash::common::detail::validate_target_paged_tree_prefix(
+                shape_cache, 5, segments, 2, 1, direct_rows));
+            CHECK(direct_rows == 6);
+            const dflash::common::QwenPrefillSegment gap[] = {
+                {0, 3, 1},
+                {4, 2, 0},
+            };
+            CHECK(!dflash::common::detail::validate_target_paged_tree_prefix(
+                shape_cache, 5, gap, 2, 1, direct_rows));
+            CHECK(!dflash::common::detail::validate_target_paged_tree_prefix(
+                shape_cache, 5, segments, 2, 3, direct_rows));
             ggml_free(shape_ctx);
         }
     }
@@ -150,7 +169,16 @@ TEST_CASE(RecurrentSnapshotFixture, invalidates_paged_tree_graph_cache_key) {
     graph.paged_tree_meta_arena.reset(
         new uint8_t[1], std::default_delete<uint8_t[]>());
     graph.paged_tree_key = TargetPagedTreeGraphKey{
-        nullptr, nullptr, nullptr, 8, 4, 256, 4096, 16, 0};
+        nullptr, nullptr, nullptr, 8, 4, 256, 4096, 16, 0, 18,
+        {1, 3, 0, 2}};
+    const TargetPagedTreeGraphKey other_slot{
+        nullptr, nullptr, nullptr, 8, 4, 256, 4096, 16, 0, 18,
+        {0, 3, 1, 2}};
+    const TargetPagedTreeGraphKey other_logits{
+        nullptr, nullptr, nullptr, 8, 4, 256, 4096, 16, 0, 19,
+        {1, 3, 0, 2}};
+    CHECK(*graph.paged_tree_key != other_slot);
+    CHECK(*graph.paged_tree_key != other_logits);
     step_graph_free(graph);
     CHECK(!graph.paged_tree_key);
     CHECK(graph.paged_tree_meta_arena);