fix(ds4): retain wide heterogeneous feature capture

server/src/deepseek4/deepseek4_backend.cpp

@@ -930,6 +930,21 @@ int DeepSeek4Backend::capture_safe_prefill_tokens(
     return safe_tokens;
 }
 
+bool DeepSeek4Backend::supports_batched_spec_feature_capture(
+        bool hybrid,
+        PrefillAttentionMode mode,
+        int n_tokens) {
+    if (mode == PrefillAttentionMode::Exact || n_tokens <= 4 ||
+        n_tokens > DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS) {
+        return false;
+    }
+    // The monolithic layer-major path reads only the requested token range.
+    // Sparse heterogeneous prefill returns every requested capture row; the
+    // caller then retains the final/snapshot window. Other hybrid modes are
+    // tokenwise and must still split at capture boundaries.
+    return !hybrid || mode == PrefillAttentionMode::Sparse;
+}
+
 bool DeepSeek4Backend::init() {
     // The shared MMVQ/MMQ crossover defaults to q=3 for NVIDIA. On gfx1151,
     // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged,
@@ -1761,10 +1776,8 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & tokens,
         }
         if (spec_enabled_ && spec_drafter_) {
             const bool batch_final_capture =
-                !w_.moe_hybrid &&
-                cache_.prefill_mode != PrefillAttentionMode::Exact &&
-                n_tok > 4 &&
-                n_tok <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS;
+                supports_batched_spec_feature_capture(
+                    w_.moe_hybrid, cache_.prefill_mode, n_tok);
             n_tok = capture_safe_prefill_tokens(
                 i, n_tok, spec_final_from, batch_final_capture,
                 save_snapshot && !snapshot_saved,

server/src/deepseek4/deepseek4_backend.h

@@ -114,8 +114,14 @@ class DeepSeek4Backend : public ModelBackend {
     void release_spec_drafter(bool mark_parked);
     void keep_spec_feature_tail(std::vector<float> & features,
                                 size_t max_rows) const;
+    // True when a wide prefill path returns per-token DSpark features and the
+    // caller can retain only the requested capture window without splitting.
+    static bool supports_batched_spec_feature_capture(
+        bool hybrid,
+        PrefillAttentionMode mode,
+        int n_tokens);
     // Limit a prefill batch to a region with a uniform DSpark capture policy.
-    // Layer-major prefill can capture a subrange without splitting the final
+    // Wide GPU paths can capture a subrange without splitting the final
     // feature window; other paths still stop exactly at capture boundaries.
     static int capture_safe_prefill_tokens(int token_offset,
                                            int requested_tokens,

server/tests/test_deepseek4_unit.cpp

@@ -2039,8 +2039,23 @@ static void test_dspark_prefill_capture_boundaries() {
     std::fprintf(stderr, "  test_dspark_prefill_capture_boundaries ...");
 
     using Backend = DeepSeek4Backend;
-    // Layer-major prefill captures only the requested tail from a wide graph,
-    // while generic paths still stop exactly at the final feature window.
+    // Both monolithic layer-major and sparse heterogeneous prefill return the
+    // per-token capture rows needed to retain a tail from one wide graph.
+    TEST_ASSERT(Backend::supports_batched_spec_feature_capture(
+                    false, PrefillAttentionMode::Sparse, 2048));
+    TEST_ASSERT(Backend::supports_batched_spec_feature_capture(
+                    true, PrefillAttentionMode::Sparse, 2048));
+    TEST_ASSERT(!Backend::supports_batched_spec_feature_capture(
+                    true, PrefillAttentionMode::Dense, 2048));
+    TEST_ASSERT(!Backend::supports_batched_spec_feature_capture(
+                    true, PrefillAttentionMode::Exact, 2048));
+    TEST_ASSERT(!Backend::supports_batched_spec_feature_capture(
+                    true, PrefillAttentionMode::Sparse, 4));
+    TEST_ASSERT(!Backend::supports_batched_spec_feature_capture(
+                    true, PrefillAttentionMode::Sparse,
+                    DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS + 1));
+
+    // Generic paths still stop exactly at the final feature window.
     TEST_ASSERT(Backend::capture_safe_prefill_tokens(
                     0, 2048, 1920, true, false, 0, 0) == 2048);
     TEST_ASSERT(Backend::capture_safe_prefill_tokens(