fix assertion DFlash2 with --split-mode tensor on CUDA

common/common.cpp

@@ -1310,6 +1310,12 @@ common_init_result::common_init_result(common_params & params, bool model_only)
         }
         cparams_dft.n_rs_seq = 0;
 
+        // keep the fit estimate in sync with common_speculative_init_result (see there)
+        if (common_speculative_is_block_draft(params.speculative.types)) {
+            cparams_dft.n_ubatch = std::min(cparams_dft.n_ubatch,
+                    common_speculative_block_draft_n_ubatch(params_dft.n_parallel, params_dft.speculative.draft.n_max));
+        }
+
         const common_fit_extra_model extra = {
             /*.path_model   =*/ params_dft.model.path.c_str(),
             /*.mparams      =*/ &mparams_dft,

common/speculative.h

@@ -3,6 +3,8 @@
 #include "llama.h"
 #include "common.h"
 
+#include <algorithm>
+
 struct common_speculative;
 
 // comma separated list the provided types
@@ -46,6 +48,20 @@ struct common_speculative_output_limits {
 common_speculative_output_limits common_speculative_get_output_limits(
         int32_t n_batch, int32_t n_parallel, int32_t n_draft);
 
+// true when any enabled type is a block draft (DFlash/DSpark)
+inline bool common_speculative_is_block_draft(const std::vector<common_speculative_type> & types) {
+    return std::any_of(types.begin(), types.end(), [](common_speculative_type t) {
+        return t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
+    });
+}
+
+// max draft-context ubatch for block drafts (DFlash/DSpark): their largest decode batch is
+// one noise block per sequence (n_max + 1 tokens), but the target-feature prefill (encoder +
+// KV injection) is chunked by the same ubatch, so keep a floor to avoid launch-bound chunks
+inline uint32_t common_speculative_block_draft_n_ubatch(uint32_t n_parallel, int32_t n_max) {
+    return std::max<uint32_t>(n_parallel * (uint32_t) std::max(1, n_max + 1), 64);
+}
+
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq);
 
 void common_speculative_free(common_speculative * spec);

src/llama-ext.h

@@ -122,6 +122,9 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
 
 LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model);
 
+// returns how the model is split across devices
+LLAMA_API enum llama_split_mode llama_model_get_split_mode(const struct llama_model * model);
+
 // returns pointer to the target-model layer indices
 LLAMA_API const int32_t * llama_model_target_layer_ids  (const struct llama_model * model);
 // returns the number of extracted layers from target model

src/llama-model.cpp

@@ -2738,6 +2738,10 @@ int32_t llama_model_dflash_selector_top_k(const llama_model * model) {
     return model->hparams.dflash_selector_top_k;
 }
 
+llama_split_mode llama_model_get_split_mode(const llama_model * model) {
+    return model->split_mode();
+}
+
 int32_t llama_model_n_head(const llama_model * model) {
     return model->hparams.n_head();
 }

src/llama-model.h

@@ -676,6 +676,8 @@ struct llama_model {
     struct ggml_tensor * dspark_conf_proj   = nullptr;
     struct ggml_tensor * dspark_conf_proj_b = nullptr;
 
+    // DFlash2 selector: the tables feed either the in-graph selector (no tensor parallelism)
+    // or the CPU-side selector that reads them straight from the GGUF file (tensor parallelism)
     struct ggml_tensor * dflash_selector_prev   = nullptr;
     struct ggml_tensor * dflash_selector_next   = nullptr;
     struct ggml_tensor * dflash_selector_hidden = nullptr;

src/models/dflash.cpp

@@ -139,8 +139,12 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
             throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice");
         }
 
-        dflash_selector_prev   = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV,   "weight"), { rank, n_vocab }, 0);
-        dflash_selector_next   = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT,   "weight"), { rank, n_vocab }, 0);
+        // Under tensor parallelism the lm_head is split along the vocabulary, which the
+        // in-graph selector cannot consume, so DFlash2 ranks on the CPU and reads these tables
+        // straight from the GGUF file; otherwise the in-graph selector uses them in the model
+        const bool use_cpu_selector = params.split_mode == LLAMA_SPLIT_MODE_TENSOR;
+        dflash_selector_prev   = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV,   "weight"), { rank, n_vocab }, use_cpu_selector ? TENSOR_SKIP : 0);
+        dflash_selector_next   = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT,   "weight"), { rank, n_vocab }, use_cpu_selector ? TENSOR_SKIP : 0);
         dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0);
 
         LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__,
@@ -469,8 +473,11 @@ static ggml_tensor * build_dflash2_conv(
     return result;
 }
 
-// DFlash2 selector: top-k candidates per block position plus the pairwise
-// transition scores, packed into the nextn output slot for the CPU-side walk.
+// DFlash2 selector (no tensor parallelism): top-k candidates per block position plus the
+// pairwise transition scores, packed into the nextn output slot for the CPU-side walk.
+// Under tensor parallelism the lm_head is split, so the same candidate math runs on the host
+// via load_dflash2_selector()/build_dflash2_selector_cpu() in common/speculative.cpp;
+// keep the two implementations in sync.
 static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
     ggml_context * ctx0 = g.ctx0;
     auto         & res  = g.res;
@@ -762,6 +769,13 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
 
     res->t_embd = cur;
 
+    // DFlash2 under tensor parallelism ranks on the CPU from the gathered lm_head output, so
+    // the decoder hidden states ride the nextn slot; without tensor parallelism the in-graph
+    // selector packs its lattice there instead (see below)
+    if (model.dflash_selector_hidden && model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
+        res->t_h_nextn = cur;
+    }
+
     // lm_head from the target model (shared via ctx_other)
     auto * output   = model.output;
     auto * output_s = model.output_s;
@@ -775,8 +789,9 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
 
     cur = build_lora_mm(output, cur, output_s);
 
-    // DFlash2 feeds these logits to the selector, so they need the target's output
-    // transforms; DFlash1 and DSpark read them through the sampler instead
+    // DFlash2 feeds these logits to its candidate selector (in-graph or on the CPU after
+    // gathering), so they need the target's output transforms; DFlash1 and DSpark read them
+    // through the sampler instead
     if (model.dflash_selector_hidden) {
         if (hparams.f_logit_scale != 0.0f) {
             cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
@@ -813,7 +828,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
         build_dspark_markov_head(*this, model, inp_tokens);
     }
 
-    if (model.dflash_selector_hidden) {
+    // DFlash2 ranks its candidates in-graph when the lm_head is not split along the vocabulary
+    if (model.dflash_selector_hidden && model.split_mode() != LLAMA_SPLIT_MODE_TENSOR) {
         build_dflash2_selector(*this, model, inp_tokens);
     }
 }