diff --git a/common/common.cpp b/common/common.cpp index 0f2f01ad0ea..31c02c81934 100644 --- a/common/common.cpp +++ b/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; } diff --git a/common/speculative.cpp b/common/speculative.cpp index ae55e357d51..9db6672d768 100644 --- a/common/speculative.cpp +++ b/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; diff --git a/include/llama.h b/include/llama.h index 177fc10a913..1855ff47546 100644 --- a/include/llama.h +++ b/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 { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d74a2135b6..19aba106d51 100644 --- a/src/llama-model.cpp +++ b/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; diff --git a/src/llama-model.h b/src/llama-model.h index 4412ef08e74..18176086c0c 100644 --- a/src/llama-model.h +++ b/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