From 1e8917fd5b6a3dbfb2de217e8f7af0a97bbff3cb Mon Sep 17 00:00:00 2001 From: sunagent Date: Thu, 20 Aug 2026 01:47:55 +0800 Subject: [PATCH 1/2] common: force draft model to SPLIT_MODE_NONE when target uses SPLIT_MODE_TENSOR DFlash/DSpark drafters share tensors with the target model (tok_embd, output). Under --split-mode tensor those tensors are split along axis 0, and the per-row ops in the draft graph (GET_ROWS, TOP_K, candidate selection) cannot consume an axis-0-split tensor, aborting with GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0). Force the draft context to a single-device (non-split) scheduler so the shared tensors are copied into the draft's own buffers instead of being split. Tested on 2x RTX 5060 Ti 16GB (tensor split): Qwen3.8-27B NVFP4 + Qwen3.8-27B-DFlash2-Q4_K_M drafter now starts and sustains 51.6-58 t/s with 86-90% draft acceptance (n_max=2, 150k ctx, q4_0 KV). --- common/speculative.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index ae55e357d51a..9db6672d7687 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; From eecc14592d5593d50455d73a069b0a3fb4055bb1 Mon Sep 17 00:00:00 2001 From: sunagent Date: Thu, 20 Aug 2026 01:52:42 +0800 Subject: [PATCH 2/2] model: keep output.weight MIRRORED only for DFlash/DSpark spec types DFlash/DSpark candidate selectors run per-row ops (TOP_K, GET_ROWS) on the shared target logits and cannot consume an axis-0-split (AXIS_1 vocab-split) result. Mirror output.weight onto every device when those spec types are enabled. For MTP and plain decoding the split is kept, saving ~1.2 GB VRAM per GPU (output.weight is ~2.4 GB; split = half per device). Verified: Qwen3.8-27B NVFP4 + DFlash2 on 2x RTX 5060 Ti tensor split - 51.6-58 t/s sustained, 86-90% draft acceptance; MTP mode unaffected (14.78 GB/GPU at 200k ctx, same as before this option existed). --- common/common.cpp | 5 +++++ include/llama.h | 1 + src/llama-model.cpp | 13 +++++++++++++ src/llama-model.h | 1 + 4 files changed, 20 insertions(+) diff --git a/common/common.cpp b/common/common.cpp index 0f2f01ad0ea8..31c02c819341 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/include/llama.h b/include/llama.h index 177fc10a9139..1855ff475465 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 0d74a2135b6e..19aba106d511 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 4412ef08e74b..18176086c0c0 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