Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 82 additions & 0 deletions proposals/dsv4-fixes/fix-deepseek4-mtp-decode-stabilize.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
# PR #3: deepseek4: stabilize MTP decode path

> Branch suggestion: `fix/deepseek4-mtp-decode-stabilize`
> Target: **ggml-org/llama.cpp master** — most changes are in shared files (`src/llama-context.cpp`, `common/speculative.cpp`) that exist upstream. Some sub-changes are arch-gated to `LLM_ARCH_DEEPSEEK4_MTP_SUPPORT` so they're inert unless DSv4 MTP support is loaded — that gating is what makes this submittable independently of DSv4 support landing.
> Commit: `f01ad092d` (43 lines across 2 files)

## Description

Follow-up to the MTP sidecar work originally in nisparks's DSv4 MTP integration (`4a310073c` snapshot). That snapshot ran model load + sidecar attach + speculative context init successfully but crashed on **the first inference request** with three independent bugs in the MTP-enabled decode path. This PR fixes all three.

### Bug 1: XOR assert in `llama_context::decode()` is too strict for MTP sidecar batches

The existing precondition

```cpp
GGML_ASSERT((!batch_inp.token && batch_inp.embd) || (batch_inp.token && !batch_inp.embd))
```

requires a batch to provide *only* tokens or *only* embeddings, never both. DeepSeek V4 Flash's MTP sidecar graph legitimately needs **both** — conditioning tokens at one position plus the hidden state row from the target's last decode. The recursive `llama_decode(ctx_mtp, ...)` call in `handle_mtp_for_ubatch()` hits this assert on the first decode.

**Fix**: relax the XOR to permit `token && embd` only when the loaded arch is `LLM_ARCH_DEEPSEEK4_MTP_SUPPORT`. Other models retain the original constraint.

### Bug 2: `common_speculative_state_mtp` writes `batch.token[0]` without allocating it

`llama_batch_init(..., embd=n_embd, ...)` allocates only the embedding buffer; the speculative state then writes to `batch.token[0]`, dereferencing a NULL or unallocated pointer.

**Fix**: explicitly allocate `batch.token` alongside `batch.embd`.

### Bug 3: hidden-row offset reads use wrong tensor stride

`handle_mtp_for_ubatch()` and the MTP-draft hidden-copy path computed row offsets as `row * n_embd * sizeof(float)`, which assumes a tightly packed F32 tensor with stride `n_embd`. Real `ggml_tensor` strides come from `tensor->nb[1]`; with non-default layouts the wrong offset reads out-of-bounds or stale data. After partial accept + checkpoint restore, `last_n_accepted=1` can also exceed the target graph's actual visible row count.

**Fix**: read row offsets via `tensor->nb[1]`; clamp `last_n_accepted` to the actual graph row count; add type/shape/bounds guards that stop the current draft pass rather than crash the server if the source tensor is malformed.

## Reproduction (before fix)

```bash
./llama-server \
-m DeepSeek-V4-Flash-IQ2XXS-...-chat-v2.gguf \
--spec-type mtp \
--spec-draft-model DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf \
--spec-draft-ngl all \
--spec-draft-ctx-size 131072 \
--spec-draft-n-max 8 --spec-draft-n-min 1 \
-c 131072 -np 1 -ngl all -fa on \
--no-warmup --cont-batching --jinja

# Container starts, /health returns 200.
# First chat completion request → GGML_ASSERT crash, exit 139:
#
# GGML_ASSERT((!batch_inp.token && batch_inp.embd)
# || (batch_inp.token && !batch_inp.embd)) failed
# at src/llama-context.cpp:1566
#
# call stack:
# llama_decode
# llama_context::handle_mtp_for_ubatch
# llama_decode (recursive)
# llama_context::decode ← assert here
```

## After fix

Same command; first request completes with valid output. Speculative decoding emits draft tokens and target-side verification finishes cleanly. Probes 7/10 pass (same set as the no-MTP same build — the 3 failing probes are pre-existing nisparks-tree issues unrelated to this fix). MTP-enabled decode now reaches 11-12 tok/s in the naïve configuration — see `fix/deepseek4-mtp-adaptive-fallback` (separate PR) for the perf work needed to keep MTP from regressing throughput vs no-MTP.

## Hardware tested

- NVIDIA RTX PRO 6000 Blackwell Workstation Edition (sm_120, 96 GB VRAM)
- CUDA 13.0, driver 595.71.05
- Build: same as companion PRs

## Why upstream

This fixes an actual crash in any DSv4 MTP-enabled decode path. The XOR relaxation is arch-gated to `LLM_ARCH_DEEPSEEK4_MTP_SUPPORT` so it's a no-op for other models. The batch.token allocation and tensor-stride fixes are general correctness improvements on the speculative state and hidden-state read paths — also safe to land independently.

The relaxation precedent is important: future hybrid-architecture sidecars (any model that needs simultaneous conditioning tokens + hidden embeddings for a recursive decode) will need the same hook.

## Related

- nisparks PR #22378 (closed) — original DSv4 support
- Snapshot commit `4a310073c "wip: mtp patch snapshot"` — the WIP that this PR fixes
- Companion: `fix/deepseek4-mtp-adaptive-fallback` (perf work; this PR is correctness only)
129 changes: 129 additions & 0 deletions proposals/dsv4-fixes/fix-deepseek4-mtp-decode-stabilize.patch
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
From f01ad092dd030acdab2884ce49d2afc4917d03b7 Mon Sep 17 00:00:00 2001
From: Codex <codex@openai.com>
Date: Sun, 10 May 2026 14:40:01 +0000
Subject: [PATCH] fix: stabilize DeepSeek4 MTP decode path

Follow-up to 4a310073c. Allow DeepSeek4 MTP sidecar batches to provide both conditioning tokens and hidden embeddings, allocate the MTP speculative token buffer, and read hidden rows using tensor strides with bounded row selection after partial accept/checkpoint restore.
---
common/speculative.cpp | 30 ++++++++++++++++++++++++++----
src/llama-context.cpp | 13 ++++++++++---
2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/common/speculative.cpp b/common/speculative.cpp
index f3a53a4e4..8a7aa1c54 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -10,6 +10,7 @@
#include "sampling.h"

#include <algorithm>
+#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <map>
@@ -638,6 +639,9 @@ struct common_speculative_state_mtp : public common_speculative_state {
smpl = common_sampler_init(model_mtp, sparams);

batch = llama_batch_init(/*n_tokens=*/ 1, /*embd=*/ n_embd, /*n_seq_max=*/ 1);
+ batch.token = (llama_token *) std::malloc(sizeof(llama_token));
+ GGML_ASSERT(batch.token != nullptr);
+
batch.n_tokens = 1;
batch.n_seq_id[0] = 1;
batch.seq_id[0][0] = 0;
@@ -701,10 +705,10 @@ struct common_speculative_state_mtp : public common_speculative_state {

if (k == 0) {
src = llama_context_get_t_h_pre_norm(ctx_tgt);
- if (last_n_accepted < 0) {
- src_row = (src && src->ne[1] > 0) ? (int32_t) src->ne[1] - 1 : 0;
+ if (src && src->ne[1] > 0 && last_n_accepted >= 0) {
+ src_row = std::min<int32_t>(last_n_accepted, (int32_t) src->ne[1] - 1);
} else {
- src_row = last_n_accepted;
+ src_row = (src && src->ne[1] > 0) ? (int32_t) src->ne[1] - 1 : 0;
}
llama_synchronize(ctx_tgt);
} else {
@@ -718,7 +722,25 @@ struct common_speculative_state_mtp : public common_speculative_state {
return;
}

- ggml_backend_tensor_get(src, batch.embd, (size_t) src_row * row_bytes, row_bytes);
+ if (src->type != GGML_TYPE_F32 || src->ne[0] != n_embd || src->nb[0] != sizeof(float) ||
+ src_row < 0 || src_row >= src->ne[1]) {
+ LOG_WRN("%s: invalid source tensor at k=%d: type=%d ne=[%" PRId64 ", %" PRId64 "] "
+ "nb=[%zu, %zu] row=%d bytes=%zu; stopping MTP chain\n",
+ __func__, k, (int) src->type, src->ne[0], src->ne[1],
+ (size_t) src->nb[0], (size_t) src->nb[1], src_row,
+ ggml_nbytes(src));
+ return;
+ }
+
+ const size_t src_offset = (size_t) src_row * src->nb[1];
+ if (src_offset + row_bytes > ggml_nbytes(src)) {
+ LOG_WRN("%s: source tensor row out of bounds at k=%d: row=%d bytes=%zu offset=%zu row_bytes=%zu; "
+ "stopping MTP chain\n",
+ __func__, k, src_row, ggml_nbytes(src), src_offset, row_bytes);
+ return;
+ }
+
+ ggml_backend_tensor_get(src, batch.embd, src_offset, row_bytes);
batch.token[0] = cond_tok;
batch.pos[0] = pos;

diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 289a22b99..c62d87bc5 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -1563,7 +1563,12 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map<llama_s
}

int llama_context::decode(const llama_batch & batch_inp) {
- GGML_ASSERT((!batch_inp.token && batch_inp.embd) || (batch_inp.token && !batch_inp.embd)); // NOLINT
+ const bool has_token = batch_inp.token != nullptr;
+ const bool has_embd = batch_inp.embd != nullptr;
+
+ // DeepSeek4 MTP sidecar batches consume the conditioning token and target hidden state together.
+ const bool is_mtp_dual_input = model.arch == LLM_ARCH_DEEPSEEK4_MTP_SUPPORT && has_token && has_embd;
+ GGML_ASSERT(is_mtp_dual_input || (has_token != has_embd)); // NOLINT

if (!memory) {
LLAMA_LOG_DEBUG("%s: cannot decode batches with this context (calling encode() instead)\n", __func__);
@@ -3153,6 +3158,7 @@ void llama_context::set_mtp(llama_context * ctx_mtp_in) {

mtp.hook_batch = llama_batch_init(n_ub, n_embd, 1);
mtp.hook_batch.token = (llama_token *) malloc(sizeof(llama_token) * n_ub);
+ GGML_ASSERT(mtp.hook_batch.token != nullptr);
mtp.pending_h.assign(n_embd, 0.0f);

LLAMA_LOG_INFO("%s: MTP draft head registered (ctx_mtp=%p, n_ubatch=%d, n_embd=%d)\n",
@@ -3199,6 +3205,7 @@ void llama_context::handle_mtp_for_ubatch(
synchronize();

const size_t row_bytes = (size_t) n_embd * sizeof(float);
+ const size_t row_stride = t->nb[1];
const int n_out = (pending_continues ? 1 : 0) + (n_rows - 1);

if (n_out > 0) {
@@ -3218,7 +3225,7 @@ void llama_context::handle_mtp_for_ubatch(
for (int k = 0; k + 1 < n_rows; ++k) {
ggml_backend_tensor_get(t,
mtp.hook_batch.embd + (size_t) out_idx * n_embd,
- (size_t) k * row_bytes,
+ (size_t) k * row_stride,
row_bytes);
mtp.hook_batch.token[out_idx] = tokens[k + 1];
mtp.hook_batch.pos[out_idx] = positions[k + 1];
@@ -3239,7 +3246,7 @@ void llama_context::handle_mtp_for_ubatch(
}

ggml_backend_tensor_get(t, mtp.pending_h.data(),
- (size_t) (n_rows - 1) * row_bytes, row_bytes);
+ (size_t) (n_rows - 1) * row_stride, row_bytes);
mtp.pending_pos = pos_start + n_rows - 1;
}

--
2.52.0