common: force draft model to SPLIT_MODE_NONE when target uses SPLIT_MODE_TENSOR

common/common.cpp

@@ -1687,6 +1687,11 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.progress_callback_user_data = params.load_progress_callback_user_data;
     mparams.no_alloc                    = params.no_alloc;
     mparams.load_mtp                    = std::find(params.speculative.types.begin(), params.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
+    // DFlash/DSpark drafters need the target logits unsplit (MIRRORED output.weight) for their
+    // candidate selector (TOP_K / GET_ROWS); MTP and plain decoding don't, so keep the split there
+    // to save ~1.2 GB VRAM per GPU.
+    mparams.mirror_output_weight         = std::find(params.speculative.types.begin(), params.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params.speculative.types.end()
+                                        || std::find(params.speculative.types.begin(), params.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params.speculative.types.end();
 
     return mparams;
 }

common/speculative.cpp

@@ -2328,6 +2328,17 @@ common_params common_base_params_to_speculative(const common_params & params) {
         result.n_gpu_layers          = params_spec.n_gpu_layers;
         result.tensor_buft_overrides = params_spec.tensor_buft_overrides;
 
+        // SPLIT_MODE_TENSOR fix: the draft model must never inherit the target's
+        // tensor split mode. DFlash/DSpark drafters share tensors with the target
+        // (tok_embd, output) which may be split along axis 0 under tensor
+        // parallelism; per-row ops (GET_ROWS, TOP_K) in the draft graph cannot
+        // consume axis-0-split inputs. Force the draft context to a single-device
+        // (non-split) scheduler - shared tensors are then copied into the draft's
+        // own buffers instead of being split.
+        if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR) {
+            result.split_mode = LLAMA_SPLIT_MODE_NONE;
+        }
+
         if (params_spec.cpuparams.n_threads > 0) {
             result.cpuparams.n_threads       = params_spec.cpuparams.n_threads;
             result.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads;

include/llama.h

@@ -339,6 +339,7 @@ extern "C" {
         bool no_host;         // bypass host buffer allowing extra buffers to be used
         bool no_alloc;        // only load metadata and simulate memory allocations
         bool load_mtp;        // whether to load MTP layers
+        bool mirror_output_weight; // keep output.weight unsplit (mirrored on every device); required only by DFlash/DSpark drafters whose candidate selector consumes full target logits
     };
 
     struct llama_sampler_seq_config {

src/llama-model.cpp

@@ -339,6 +339,7 @@ llama_model * llama_model_create(llm_arch arch, const llama_model_params & param
 
     if (model != nullptr) {
         model->arch = arch;
+        model->mirror_output_weight = params.mirror_output_weight;
         if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR && !llm_arch_supports_sm_tensor(arch)) {
             throw std::runtime_error(std::string("LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '") + llm_arch_name(arch) + "'");
         }
@@ -516,6 +517,17 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
 
         // output
         if (std::regex_match(tensor_name, pattern_output_weight)) {
+            if (ud->model->mirror_output_weight) {
+                // MIRRORED (not AXIS_1): keep the output projection unsplit so the
+                // logits come out mirrored. This is required by DFlash/DSpark
+                // drafters, whose candidate selector runs TOP_K / GET_ROWS (per-row
+                // ops) on the shared target logits and cannot consume an
+                // axis-0-split result. The output weight (~2.4GB BF16) then lives
+                // on a single device instead of being vocab-split across GPUs.
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            }
+            // MTP / plain decoding: vocab-split the output projection across GPUs
+            // (half the rows per device), saving ~1.2 GB of VRAM per GPU vs MIRRORED.
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
         }
         if (std::regex_match(tensor_name, pattern_output_bias)) {
@@ -2494,6 +2506,7 @@ llama_model_params llama_model_default_params() {
         /*.no_host                     =*/ false,
         /*.no_alloc                    =*/ false,
         /*.load_mtp                    =*/ false,
+        /*.mirror_output_weight        =*/ false,
     };
 
     return result;

src/llama-model.h

@@ -588,6 +588,7 @@ struct llama_model {
     std::string name = "n/a";
 
     llama_hparams hparams = {};
+    bool mirror_output_weight = false; // output.weight stays unsplit (mirrored); set by DFlash/DSpark spec init
     llama_vocab   vocab;
 
     // for classifier models