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

[art-den]

Overview

Assertion is failed on llama-server startup if using dflash2 draft model and --split-mode tensor flag:

ggml\src\ggml-backend-meta.cpp:543: GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) failed

The changes in this PR address two issues:

• Crash in GGML_ASSERT (root cause)
• Reducing the VRAM footprint of the DFlash2 model for --split-mode tensor

Additional information

Issue for this case: #27819

Root cause

The DFlash2 candidate selector runs ggml_top_k() on the lm_head output (t_logits) inside the graph. Under tensor parallelism output.weight is split along the vocab axis, so t_logits is split along axis 0, which GGML_OP_TOP_K (a per-row op) cannot handle in the meta backend.

Fix: When a DFlash2 draft is used with tensor parallelism, the lm_head must be replicated in full on every device (GGML_BACKEND_SPLIT_AXIS_MIRRORED) so the whole vocabulary is available in-graph.

Reducing the VRAM footprint

DFlash2's in-graph selector (ggml_top_k over the lm_head output) cannot run on a vocab-split lm_head under --split-mode tensor. Rather than replicating the lm_head on every device (MIRRORED, ~+500 MiB/device), the selector now runs on the CPU from the gathered logits + decoder hidden states, so the lm_head stays split.

VRAM savings for https://huggingface.co/z-lab/Qwen3.8-27B-DFlash2-GGUF/blob/main/Qwen3.8-27B-DFlash2-Q4_K_M.gguf:

• lm_head stays vocab-split instead of replicated: ~-500 MiB/device.
• Selector codebook tables are TENSOR_SKIP'd under tensor parallelism and read from the GGUF at init: ~-68 MiB/device.
• Draft ubatch is capped (floor 64) so the reserved [vocab x ubatch] compute buffer drops from ~539 MiB to ~67 MiB/device.

Without these modifications, the amount of occupied VRAM was x4.4 of the gguf file size.

NOTE: No ggml/ changes; the non-tensor path (single GPU / layer split) keeps the original in-graph selector.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, Deepseek 4 Flash was used for investigation, fixing and reviewing

[ggml-gh-bot]

Hi @art-den, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[art-den]

Bad news. VRAM usage for draft model is x4.4 of original Qwen3.8-27B-DFlash2-Q4_K_M.gguf size for my two RTX 5060 Ti 16Gb. I have no idea how to fix it without large changes in sources.

I'm not sure yet if this MR is necessary ((

[mb1986]

Bad news. VRAM usage for draft model is x4.4 of original Qwen3.8-27B-DFlash2-Q4_K_M.gguf size for my two RTX 5060 Ti 16Gb. I have no idea how to fix it without large changes in sources.

I'm not sure yet if this MR is necessary ((

Have you tried to set --parallel 1?

It is possible for me to run --split-mode tensor --spec-type draft-dflash --parallel 1 on two P40s.

b10679 + #25592 + #26636 + #27858

[art-den]

I think it's ready. I managed to reduce video memory consumption without a much rewriting.

The current implementation produces ~73 t/s with dflash2 vs ~63 t/s with MTP for my x2 RTX 5060 Ti 16Gb.

Please check that I haven't broken dflash2 and dspark in single-GPU mode.

[pwilkin]

@ruixiang63 could you maybe take a look at this one? The diagnosis seems correct, but the code could use a look from a specdec expert.

[ruixiang63]

I don't have a good feeling about this PR, there are too much unnecessary changes.

Fix: When a DFlash2/DSpark draft is used with tensor parallelism, the lm_head must be replicated in full on every device (GGML_BACKEND_SPLIT_AXIS_MIRRORED) so the whole vocabulary is available in-graph.

This is the fix proposed in the description but there are so much irrlevant changes in the PR. Could you limit this PR to the changes required for the fix?

[art-den]

@ruixiang63

I also have a bad feeling about my PR. I asked the agent to make as few changes as possible, but they ended up making far more than I expected.

The minimum change that would solve the problem is the first commit. But it has some unresolved issues (too much VRAM usage for the draft model).

This PR needs to be reviewed by someone familiar with llama.cpp. If these changes don't solve the problem correctly, I'll simply close the PR.

[ruixiang63]

The PR authors should be responsible for validating their own changes before requesting review. If you’re not confident in the implementation, please investigate and self-review it first. Asking an agent to write the code and then leaving reviewers or maintainers to determine whether it’s correct doesn't make any sense.

[pwilkin]

/bot review

[ggml-gh-bot]

Automated code review

Reviewed statically against the repo conventions; the PR touches common/, src/llama-model.*, src/models/dflash.cpp and adds a public llama.h field, so I ran the scope gate, security pass, approach/design pass, and the public-API checklist. The issue link (27819) and AI disclosure are in place, and the core approach is sound: I verified the CPU-side selector math (unary + sel_next[cand,k] * sel_prev[pred,k] * gate[k], gate = hidden^T selector_hidden) matches the removed in-graph version, including the anchor/predecessor chain and the p_min early-break. The ubatch cap is load-bearing: with every block position now requesting logits, an uncapped decode would trip GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max) in output_reserve, so having it in both the real init and the fit estimate is correct.

Blocking

• (1) common/speculative.cpp load_dflash2_selector() (~lines 1077-1117): every failure path on GGUF-file-derived values is a GGML_ASSERT (open failure, missing tensor, unsupported type, seek/read failure, and the rank-consistency checks at 1114-1116). These run at speculative-init on a user-supplied draft GGUF, and elsewhere this exact class of problem throws std::runtime_error so callers can report it cleanly (the old path did: create_tensor with flags 0 threw on a missing or mis-shaped selector tensor). Worse, TENSOR_SKIP in src/models/dflash.cpp:144-145 also removes the loader's shape validation, so a malformed draft GGUF now passes model load and aborts the whole process later. Please validate the tensor meta explicitly (ne == {rank, n_vocab}, type check) and throw instead of asserting. Related: the comment says the tables "are i8" but the code only checks tt->to_float exists, so any type is silently accepted; either check the expected type or drop the i8 claim.

Will slow the review

• (2) src/models/dflash.cpp:153-157: output_replicated = true fires for selector_meta || markov_meta, but the comment (and the PR description) justify replication only for the in-graph DSpark markov head. DFlash2 now ranks on the CPU from gathered logits, so per this PR's own design it does not need a replicated head. If selector_meta is actually required (e.g. the d2t ggml_set_rows scatter on axis-0-split logits), document that reason in the comment; otherwise drop selector_meta so DFlash2 drafts keep the split lm_head and the VRAM saving applies to them too.

• (3) common_speculative_draft_replicated_lm_head() (~line 2507): this replicates the target lm_head for any DSpark draft, even when the draft ships its own output.weight (drafts may, see dflash.cpp load) and thus never touches the target's projection in-graph - that wastes n_vocab * n_embd bytes per device. Consider also requiring that the draft GGUF has no output.weight before replicating the target's. Also, this duplicates common_speculative_types_from_gguf(): same GGUF open, same markov_w1.weight sniff. Reusing it (check for COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) would avoid adding an exported helper and a third full GGUF metadata parse at startup.

• (4) Perf data: the PR moves a GPU in-graph selector to a host-side O(n_tokens * n_vocab) scan plus a full logits copy-out per draft decode, and adds ubatch caps with a magic floor of 64 - yet the commits claim "without reducing the speed" and one is literally a speed-regression fix. Please add before/after measurements (draft latency, prompt processing, VRAM/compute buffer) for at least one realistic draft, and confirm the h_nextn extraction under tensor parallelism yields full hidden rows (a partial row would silently corrupt the gate scores, not crash).

• (5) common/speculative.cpp:1066: the comment "DFlash2 reads its selector lattice from h_nextn and never consumes raw logits" is now stale - this PR makes DFlash2 consume raw logits. Update it.

Nits

• (6) The ubatch-cap formula now exists twice (fit estimate in common/common.cpp ~1322, real init in common/speculative.cpp ~2742) kept in sync by comments; a tiny shared helper would prevent drift.
• (7) src/llama-model.h:679-680: dflash_selector_prev/next members are now write-only (TENSOR_SKIP returns nullptr, no reads remain). Drop the members or don't assign. Also note TENSOR_SKIP logs "model has unused tensor ... -- ignoring" as a WARN for tensors that are in fact used on the CPU.
• (8) common/speculative.cpp:1093-1100: the hand-rolled fseeko/_fseeki64 + fread block could use llama_file (the file already includes ../src/llama-ext.h, so a src include has precedent) and drop the platform ifdef.
• (9) include/llama.h:350: the new field's comment should note who sets it (common/, for speculative drafts) so API users don't enable it casually - it silently multiplies lm_head VRAM per device. Otherwise the public-API addition itself looks justified: the target model genuinely cannot know at load time that a DSpark draft will consume its output tensor in-graph.

This review was generated automatically by pi coding agent using zai-org/GLM-5.3. It may contain mistakes. Maintainers make the final call.

[art-den]

/bot review

[ggml-gh-bot]

Sorry @art-den, only members of the ggml-org organization can trigger /bot commands.

[art-den]

@ruixiang63, @pwilkin new changes for this PR:

• Updated comment for PR
• Fixed build for GCC
• Fixed comments from review bot
• Restored the non-tensor path (single GPU / layer split) to minimize PR changes
• Reverted fix for DSpark to minimize PR changes

[art-den]

Hm... --no-mmproj-offload is broken in this PR if DFlash2 used

[AiwendilInTheWoods]

Tested this branch on ROCm — it fixes the assert there too.

Environment: 2× AMD Radeon AI PRO R9700 (gfx1201, 32 GiB each), ROCm 7.2.3, kernel 6.8, HIP_VISIBLE_DEVICES=0,1. Branch fix/dflash2-tensor-split @ 3286b033c built with -DGGML_HIP=ON -DAMDGPU_TARGETS=gfx1201 -DCMAKE_HIP_ARCHITECTURES=gfx1201 -DGGML_HIP_GRAPHS=ON -DCMAKE_BUILD_TYPE=Release (/opt/rocm/llvm/bin/clang 22). Target unsloth/Qwen3.8-27B-GGUF Qwen3.8-27B-UD-Q6_K_XL.gguf; drafter z-lab/Qwen3.8-27B-DFlash2-GGUF Qwen3.8-27B-DFlash2-Q8_0.gguf (the reconverted one with dflash.rope.dimension_sections).

Before (master-based build without this branch): at n_max 3, 5 and 7 under --split-mode tensor:

common_speculative_impl_draft_dflash: - n_max=7, n_min=0, p_min=0.00
common_speculative_impl_draft_dflash: - block_size=8, mask_token_id=248070, n_extract=5, sample_from_anchor=true
ggml/src/ggml-backend-meta.cpp:543: GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) failed

With this branch: loads and runs at all three depths; ~15 loads, no NaNs, no crashes.

llama-server -m Qwen3.8-27B-UD-Q6_K_XL.gguf -md Qwen3.8-27B-DFlash2-Q8_0.gguf \
  --spec-type draft-dflash --spec-draft-n-max 3 -sm tensor -fa on -c 131072 -ctk q8_0 -ctv q8_0 \
  -ngl 999 --jinja --parallel 1 --no-context-shift --host 127.0.0.1 --port 8091

Results (greedy, temperature 0, 4 distinct ~30-token prose prompts × 256 completion tokens × 3 reps, fresh server per config, same binary):

config	tok/s (grand median)	mean accepted len	VRAM by PID, 2 GPUs
no speculation	30.5	—	30.8 GB
draft-mtp n3	56.5	2.4–2.9	32.7 GB
draft-dflash n3	53.3	2.2–3.0	33.9 GB
draft-dflash n5	51.0	2.5–3.2	34.2 GB
draft-dflash n7	45.7	2.5–3.3	34.5 GB

Output is deterministic across reps and byte-identical (SHA-256 of the completion) at n3, n5 and n7, and identical to draft-mtp's on all four prompts. So the CPU-side selector is producing the same tokens as the target here.

draft acceptance in the slot timings: 0.41–0.65 at n3, 0.21–0.33 at n7 — accepted length doesn't grow with depth on this pairing, which is why n7 is slowest; that's a separate topic from this PR.

Happy to re-run anything useful (f16 KV, other depths, --no-mmproj-offload). Result JSON and server logs available.