From ba4f6913ab96e888ac9a7febed8360bb123ab21a Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 7 Jul 2026 13:59:25 +0200 Subject: [PATCH 01/12] feat(dflash): bridge DeepSeek4 DSpark aux heads --- server/docs/DS4.md | 36 +++++ server/scripts/convert_dflash_to_gguf.py | 167 +++++++++++++++-------- 2 files changed, 143 insertions(+), 60 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index ca83f8555..c1d8ea264 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -138,6 +138,42 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. DeepSeek4 no longer uses the old expert-split environment variables or expert-worker tuning knobs. Those retired knobs were removed from the codebase rather than left behind as unsupported debug switches. +## DSpark Aux Draft Heads + +The current DSpark runtime lives in the shared DFlash speculative stack and is +used by the Laguna backend when the draft GGUF carries `dflash.dspark.*` +tensors. DeepSeek4 DSpark work is still a draft bridge: the DeepSeek4/MTP +artifact stores compatible aux heads under the `mtp.2.*` namespace, and the +converter can now map those names into the existing GGUF contract. + +Supported DeepSeek4/MTP input tensors: + +| DeepSeek4/MTP tensor | GGUF tensor | +|----------------------|-------------| +| `mtp.2.markov_head.markov_w1.weight` | `dflash.dspark.markov.w1` | +| `mtp.2.markov_head.markov_w2.weight` | `dflash.dspark.markov.w2` | +| `mtp.2.confidence_head.proj.weight` | `dflash.dspark.confidence.weight` | +| `mtp.2.confidence_head.proj.bias` | `dflash.dspark.confidence.bias` | + +If the MTP confidence projection is bias-less, the converter writes a zero +bias so the GGUF loader still sees the pair it expects. The Markov head alone +is enough for DSpark greedy-chain correction; confidence gating remains +optional. + +Example conversion with the DS4 MTP shard that contains the DSpark heads: + +```bash +python server/scripts/convert_dflash_to_gguf.py \ + /path/to/dflash-draft/model.safetensors \ + /path/to/dflash-draft.gguf \ + --aux-heads /path/to/hf-ds4-flash-dspark/model-00048-of-00048.safetensors +``` + +This does not by itself make the production DeepSeek4 layer-split backend use +DSpark. It makes the DeepSeek4 DSpark aux artifact consumable by the existing +`dflash.dspark.*` GGUF loader/runtime so the follow-up runtime PR can wire it +into the DS4 decode path with a clean tensor contract. + ## Example: CUDA + Halo Layer Split Automatic split (CUDA prefix chosen from free memory, optional manual override via `DFLASH_DS4_CUDA_LAYERS`): diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index 9c31f7531..6c1c46291 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -249,36 +249,83 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: DSPARK_TENSOR_MAP = { - "dspark_markov_head.markov_w1.weight": ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), - "dspark_markov_head.markov_w2.weight": ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), + ("dspark_markov_head.markov_w1.weight", + "mtp.2.markov_head.markov_w1.weight"): ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), + ("dspark_markov_head.markov_w2.weight", + "mtp.2.markov_head.markov_w2.weight"): ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), } DSPARK_CONFIDENCE_TENSOR_MAP = { - "dspark_confidence_head.weight": ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), - "dspark_confidence_head.bias": ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), + ("dspark_confidence_head.weight", + "mtp.2.confidence_head.proj.weight"): ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), + ("dspark_confidence_head.bias", + "mtp.2.confidence_head.proj.bias"): ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), } -def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): - if aux_path is None: - return - if not aux_path.exists(): - print(f"[warn] Domino aux heads not found: {aux_path}") - return +def load_aux_state(aux_path: Path, label: str, wanted_names: set[str] | None = None): + if aux_path.suffix == ".safetensors": + header_size, header = load_safetensors_header(aux_path) + state = {} + for name, info in header.items(): + if name == "__metadata__": + continue + if wanted_names is not None and name not in wanted_names: + continue + raw = read_tensor_bytes(aux_path, header_size, info) + state[name] = bytes_to_np(raw, info["dtype"], info["shape"]) + return state try: import torch except ImportError as exc: - print(f"[error] --aux-heads requires torch: {exc}", file=sys.stderr) + print(f"[error] --aux-heads requires torch for {label} .pt files: {exc}", file=sys.stderr) sys.exit(1) - print(f"[info] reading Domino aux heads from {aux_path}") state = torch.load(aux_path, map_location="cpu") if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict): state = state["state_dict"] if not isinstance(state, dict): - print(f"[error] Domino aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) + print(f"[error] {label} aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) sys.exit(1) + return state + + +def tensor_to_np(t, raw_dtype): + if hasattr(t, "detach"): + arr = t.detach().cpu().float().numpy() + else: + arr = np.asarray(t) + if arr.dtype.kind not in ("f", "i", "u"): + arr = arr.astype(np.float32) + if raw_dtype == gguf.GGMLQuantizationType.F16: + return arr.astype(" set[str]: + out = set() + for names in alias_map: + out.update(names) + return out + + +def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): + if aux_path is None: + return + if not aux_path.exists(): + print(f"[warn] Domino aux heads not found: {aux_path}") + return + + print(f"[info] reading Domino aux heads from {aux_path}") + state = load_aux_state(aux_path, "Domino", set(DOMINO_TENSOR_MAP)) if not any(k in state for k in DOMINO_TENSOR_MAP): return @@ -300,14 +347,7 @@ def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): writer.add_uint32(f"{arch}.dflash.domino.vocab_size", vocab) for st_name, (gguf_name, raw_dtype) in DOMINO_TENSOR_MAP.items(): - t = state[st_name] - if hasattr(t, "detach"): - t = t.detach().cpu() - arr = t.float().numpy() - if raw_dtype == gguf.GGMLQuantizationType.F16: - arr = arr.astype("{raw_dtype.name:4s} {tuple(arr.shape)}") @@ -318,26 +358,22 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): if not aux_path.exists(): return - try: - import torch - except ImportError as exc: - print(f"[error] --aux-heads requires torch: {exc}", file=sys.stderr) - sys.exit(1) - - state = torch.load(aux_path, map_location="cpu") - if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict): - state = state["state_dict"] - if not isinstance(state, dict): - print(f"[error] DSpark aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) - sys.exit(1) - - missing = [k for k in DSPARK_TENSOR_MAP if k not in state] + wanted_names = flat_alias_names(DSPARK_TENSOR_MAP) | flat_alias_names(DSPARK_CONFIDENCE_TENSOR_MAP) + state = load_aux_state(aux_path, "DSpark", wanted_names) + resolved = {} + missing = [] + for names, spec in DSPARK_TENSOR_MAP.items(): + found_name, tensor = find_state_tensor(state, names) + if tensor is None: + missing.append("/".join(names)) + continue + resolved[names] = (found_name, tensor, spec) if missing: return print(f"[info] reading DSpark aux heads from {aux_path}") - w1 = state["dspark_markov_head.markov_w1.weight"] - w2 = state["dspark_markov_head.markov_w2.weight"] + w1 = resolved[("dspark_markov_head.markov_w1.weight", "mtp.2.markov_head.markov_w1.weight")][1] + w2 = resolved[("dspark_markov_head.markov_w2.weight", "mtp.2.markov_head.markov_w2.weight")][1] vocab = int(w1.shape[0]) rank = int(w1.shape[1]) if tuple(w2.shape) != (vocab, rank): @@ -348,21 +384,39 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): writer.add_uint32(f"{arch}.dflash.dspark.markov_rank", rank) writer.add_uint32(f"{arch}.dflash.dspark.vocab_size", vocab) - for st_name, (gguf_name, raw_dtype) in DSPARK_TENSOR_MAP.items(): - t = state[st_name] - if hasattr(t, "detach"): - t = t.detach().cpu() - arr = t.float().numpy().astype("{raw_dtype.name:4s} {tuple(arr.shape)}") + print(f"[tensor] {gguf_name:50s} aux:{st_name} ->{raw_dtype.name:4s} {tuple(arr.shape)}") + + conf_resolved = {} + conf_missing = [] + for names, spec in DSPARK_CONFIDENCE_TENSOR_MAP.items(): + found_name, tensor = find_state_tensor(state, names) + if tensor is None: + conf_missing.append(names) + continue + conf_resolved[names] = (found_name, tensor, spec) + weight_names = ("dspark_confidence_head.weight", "mtp.2.confidence_head.proj.weight") + bias_names = ("dspark_confidence_head.bias", "mtp.2.confidence_head.proj.bias") + if weight_names not in conf_resolved: + if conf_missing: + print("[warn] incomplete DSpark confidence head; Markov head will still load") + return + if bias_names not in conf_resolved: + conf_resolved[bias_names] = ( + "", + np.zeros((1,), dtype=np.float32), + DSPARK_CONFIDENCE_TENSOR_MAP[bias_names], + ) + print("[warn] DSpark confidence head has no bias tensor; writing a zero bias") - conf_missing = [k for k in DSPARK_CONFIDENCE_TENSOR_MAP if k not in state] - if conf_missing: - print(f"[warn] incomplete DSpark confidence head; missing {conf_missing}; Markov head will still load") + if conf_missing and bias_names in conf_missing and len(conf_missing) > 1: + print("[warn] incomplete DSpark confidence head; Markov head will still load") return - conf_w = state["dspark_confidence_head.weight"] - conf_b = state["dspark_confidence_head.bias"] + conf_w = conf_resolved[weight_names][1] + conf_b = conf_resolved[bias_names][1] confidence_dim = int(conf_w.shape[1]) if int(conf_w.shape[0]) != 1 or tuple(conf_b.shape) != (1,): print( @@ -372,17 +426,10 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): sys.exit(1) writer.add_uint32(f"{arch}.dflash.dspark.confidence_dim", confidence_dim) writer.add_uint32(f"{arch}.dflash.dspark.confidence.enabled", 1) - for st_name, (gguf_name, raw_dtype) in DSPARK_CONFIDENCE_TENSOR_MAP.items(): - t = state[st_name] - if hasattr(t, "detach"): - t = t.detach().cpu() - arr = t.float().numpy() - if raw_dtype == gguf.GGMLQuantizationType.F16: - arr = arr.astype("{raw_dtype.name:4s} {tuple(arr.shape)}") + print(f"[tensor] {gguf_name:50s} aux:{st_name} ->{raw_dtype.name:4s} {tuple(arr.shape)}") # ────────────────────────────────────────────────────────────────────── @@ -394,9 +441,9 @@ def main(): ap.add_argument("safetensors", type=Path) ap.add_argument("out_gguf", type=Path) ap.add_argument("--aux-heads", type=Path, default=None, - help="optional Domino aux-head .pt file; defaults to dflash_aux_heads.pt next to the safetensors") + help="optional Domino/DSpark aux-head .pt or DS4 MTP .safetensors file; defaults to dflash_aux_heads.pt next to the safetensors") ap.add_argument("--no-aux-heads", action="store_true", - help="do not auto-embed Domino aux-head tensors") + help="do not auto-embed Domino/DSpark aux-head tensors") args = ap.parse_args() if not args.safetensors.exists(): From 0ea1be6c3f92ef97f207121b6cecaad64589f27d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:36:45 +0200 Subject: [PATCH 02/12] feat(deepseek4): add DSpark speculative runtime --- server/CMakeLists.txt | 18 + server/src/deepseek4/deepseek4_backend.cpp | 84 +- server/src/deepseek4/deepseek4_backend.h | 6 + server/src/deepseek4/deepseek4_dspark.cpp | 300 +++++++ server/src/deepseek4/deepseek4_dspark.h | 141 ++++ .../src/deepseek4/deepseek4_dspark_spec.cpp | 686 ++++++++++++++++ .../src/deepseek4/deepseek4_fused_verify.inc | 517 ++++++++++++ server/src/deepseek4/deepseek4_graph.cpp | 743 +++++++++++++++++- server/src/deepseek4/deepseek4_internal.h | 19 +- server/test/test_ds4_dspark_load.cpp | 153 ++++ 10 files changed, 2630 insertions(+), 37 deletions(-) create mode 100644 server/src/deepseek4/deepseek4_dspark.cpp create mode 100644 server/src/deepseek4/deepseek4_dspark.h create mode 100644 server/src/deepseek4/deepseek4_dspark_spec.cpp create mode 100644 server/src/deepseek4/deepseek4_fused_verify.inc create mode 100644 server/test/test_ds4_dspark_load.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 853cf571e..b2d5f833f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -258,6 +258,8 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_daemon.cpp src/deepseek4/deepseek4_layer_split_adapter.cpp src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp + src/deepseek4/deepseek4_dspark.cpp + src/deepseek4/deepseek4_dspark_spec.cpp src/flashprefill_q8.cpp src/kv_cache.cpp src/kv_quant.cpp @@ -892,6 +894,22 @@ if(DFLASH27B_TESTS) endif() add_test(NAME deepseek4_unit COMMAND test_deepseek4_unit) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ds4_dspark_load.cpp") + add_executable(test_ds4_dspark_load test/test_ds4_dspark_load.cpp) + target_include_directories(test_ds4_dspark_load PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/include) + if(DFLASH27B_GPU_BACKEND STREQUAL "hip") + target_compile_definitions(test_ds4_dspark_load PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) + else() + target_compile_definitions(test_ds4_dspark_load PRIVATE DFLASH27B_BACKEND_CUDA=1 GGML_USE_CUDA) + endif() + target_link_libraries(test_ds4_dspark_load PRIVATE dflash_common ggml ggml-cpu ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + find_package(CUDAToolkit REQUIRED) + target_link_libraries(test_ds4_dspark_load PRIVATE CUDA::cudart) + else() + target_link_libraries(test_ds4_dspark_load PRIVATE hip::host) + endif() + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_load_draft.cpp") add_executable(smoke_load_draft test/smoke_load_draft.cpp) target_include_directories(smoke_load_draft PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index eb6799b58..b3c1a14f6 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -403,6 +403,27 @@ bool DeepSeek4Backend::init() { w_.n_layer, max_ctx, w_.n_expert, active_experts, w_.n_expert_used, w_.fused_decode ? "on" : "off", moe_hybrid_ ? " [hybrid]" : ""); + + if (env_flag_enabled("DFLASH_DS4_SPEC") && moe_hybrid_) { + std::fprintf(stderr, + "[deepseek4] DSpark spec-decode requires monolithic model placement; " + "disabled for hybrid expert placement\n"); + } else if (env_flag_enabled("DFLASH_DS4_SPEC")) { + const char * dp = std::getenv("DFLASH_DS4_DRAFT"); + if (dp && *dp) { + spec_drafter_ = std::make_unique(); + if (load_deepseek4_dspark_drafter(dp, backend_, *spec_drafter_)) { + spec_enabled_ = true; + std::fprintf(stderr, "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", dp); + } else { + std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", + deepseek4_dspark_last_error()); + spec_drafter_.reset(); + } + } else { + std::fprintf(stderr, "[deepseek4] DFLASH_DS4_SPEC set but DFLASH_DS4_DRAFT gguf missing\n"); + } + } return true; } @@ -578,6 +599,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, reset_deepseek4_cache(cache_); } last_logits_.clear(); + if (spec_enabled_ && kv_offset == 0) spec_feat_window_.clear(); const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); DeepSeek4StepTelemetry tel_acc; @@ -596,11 +618,30 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); std::vector logits; - if (!deepseek4_step(backend_, w_, cache_, embed.data(), n_tok, pos, logits, - moe_hybrid_.get(), tokens.data() + i, - moe_hybrid_ ? &stream_engine_ : nullptr, - timing ? &step_tel : nullptr, - routing_stats_.get())) { + Ds4VerifyHooks spec_hooks; + std::vector spec_cap; + Ds4VerifyHooks * hp = nullptr; + if (spec_enabled_ && spec_drafter_) { + spec_hooks.capture_layer_ids = &spec_drafter_->capture_layer_ids; + spec_hooks.capture_out = &spec_cap; + hp = &spec_hooks; + } + const bool ok = deepseek4_step( + backend_, w_, cache_, embed.data(), n_tok, pos, logits, + moe_hybrid_.get(), tokens.data() + i, + moe_hybrid_ ? &stream_engine_ : nullptr, + timing ? &step_tel : nullptr, + routing_stats_.get(), hp); + if (ok && hp && !spec_cap.empty()) { + const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; + for (int t = 0; t < n_tok; ++t) { + spec_feat_window_.insert( + spec_feat_window_.end(), + spec_cap.begin() + (size_t) t * feat_row, + spec_cap.begin() + (size_t) (t + 1) * feat_row); + } + } + if (!ok) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); return -1; } @@ -721,6 +762,35 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, // Decode auto t1 = Clock::now(); + if (spec_enabled_ && spec_drafter_ && req.n_gen > 0) { + if (last_logits_.empty()) { result.error = "spec: no prefill logits"; return result; } + int seed = 0; + { float mv = last_logits_[0]; + for (int i = 1; i < w_.n_vocab; i++) if (last_logits_[i] > mv) { mv = last_logits_[i]; seed = i; } } + std::vector gen; + gen.push_back(seed); + io.emit(seed); + float accept_rate = 0.0f; + if (!deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { + const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; + const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; + std::vector spec_toks; + run_deepseek4_dspark_spec_decode(backend_, w_, cache_, *spec_drafter_, + committed, seed, req.n_gen - 1, + win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, + spec_toks, &accept_rate); + for (int t : spec_toks) io.emit(t); + gen.insert(gen.end(), spec_toks.begin(), spec_toks.end()); + } + result.ok = true; + result.tokens = std::move(gen); + result.decode_s = elapsed_s(t1); + std::fprintf(stderr, "[deepseek4] DSpark decode: %zu tok in %.3fs (%.1f tok/s) accept_rate=%.2f\n", + result.tokens.size(), result.decode_s, + result.decode_s > 0 ? result.tokens.size() / result.decode_s : 0.0, accept_rate); + maybe_save_routing_stats(); + return result; + } std::vector gen_tokens; gen_tokens.reserve(req.n_gen); @@ -777,7 +847,9 @@ bool DeepSeek4Backend::handle_compress(const std::string & line, } void DeepSeek4Backend::free_drafter() { - // No drafter in AR-only mode + spec_drafter_.reset(); + spec_enabled_ = false; + spec_feat_window_.clear(); } void DeepSeek4Backend::maybe_save_routing_stats() { diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 537d3ba44..64ce91464 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -13,6 +13,7 @@ #include "../common/moe_hybrid_storage.h" #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" +#include "deepseek4_dspark.h" #include "ggml.h" #include "ggml-backend.h" @@ -76,6 +77,11 @@ class DeepSeek4Backend : public ModelBackend { DeepSeek4Snapshot snapshots_[PREFIX_SLOTS]; std::vector last_logits_; + // DSpark speculative decode (opt-in: DFLASH_DS4_SPEC=1 + DFLASH_DS4_DRAFT=). + bool spec_enabled_ = false; + std::unique_ptr spec_drafter_; + std::vector spec_feat_window_; + // Prefill prompt tokens in chunks, return absolute committed position. int do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset = 0); diff --git a/server/src/deepseek4/deepseek4_dspark.cpp b/server/src/deepseek4/deepseek4_dspark.cpp new file mode 100644 index 000000000..8975244d1 --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark.cpp @@ -0,0 +1,300 @@ +// DeepSeek-V4-Flash "DSpark" drafter loader. See deepseek4_dspark.h. +// +// Self-contained (shares nothing with the target loader) so it cannot regress +// the target path. Loads a "deepseek4-dflash-draft" GGUF into a DSparkDrafter: +// the n_layer decoder blocks reuse DeepSeek4Weights leaf-name bindings, and the +// DSpark-specific tensors (dflash.fc / hidden_norm / dspark.*) bind separately. + +#include "deepseek4_dspark.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include +#include + +#include +#include + +namespace dflash::common { + +namespace { + +std::string g_dspark_err; +void set_err(const std::string & m) { g_dspark_err = m; std::fprintf(stderr, "[ds4-dspark] %s\n", m.c_str()); } + +const char * ARCH = "deepseek4-dflash-draft"; + +uint32_t kv_u32(gguf_context * g, const std::string & key, uint32_t def) { + const int64_t id = gguf_find_key(g, key.c_str()); + if (id < 0) return def; + if (gguf_get_kv_type(g, id) != GGUF_TYPE_UINT32) return def; + return gguf_get_val_u32(g, id); +} +float kv_f32(gguf_context * g, const std::string & key, float def) { + const int64_t id = gguf_find_key(g, key.c_str()); + if (id < 0) return def; + if (gguf_get_kv_type(g, id) != GGUF_TYPE_FLOAT32) return def; + return gguf_get_val_f32(g, id); +} + +// Read an INT32/UINT32 array KV into ints. Returns false if the key is absent. +bool kv_i32_array(gguf_context * g, const std::string & key, std::vector & out) { + const int64_t id = gguf_find_key(g, key.c_str()); + if (id < 0) return false; + if (gguf_get_kv_type(g, id) != GGUF_TYPE_ARRAY) return false; + const enum gguf_type at = gguf_get_arr_type(g, id); + const size_t n = gguf_get_arr_n(g, id); + out.resize(n); + const void * raw = gguf_get_arr_data(g, id); + if (at == GGUF_TYPE_INT32) { + const int32_t * v = static_cast(raw); + for (size_t i = 0; i < n; i++) out[i] = (int)v[i]; + } else if (at == GGUF_TYPE_UINT32) { + const uint32_t * v = static_cast(raw); + for (size_t i = 0; i < n; i++) out[i] = (int)v[i]; + } else { + out.clear(); + return false; + } + return true; +} + +// Suffix after "blk.." — returns "" if `name` isn't a block tensor. +bool block_suffix(const char * name, int & il, std::string & suffix) { + if (std::strncmp(name, "blk.", 4) != 0) return false; + const char * p = name + 4; + if (*p < '0' || *p > '9') return false; + char * end = nullptr; + long v = std::strtol(p, &end, 10); + if (!end || *end != '.' || v < 0) return false; + il = (int)v; + suffix = std::string(end + 1); + return true; +} + +} // namespace + +const char * deepseek4_dspark_last_error() { return g_dspark_err.c_str(); } + +bool load_deepseek4_dspark_drafter(const std::string & path, + ggml_backend_t backend, + DSparkDrafter & out) { + ggml_context * meta = nullptr; + gguf_init_params gip{}; + gip.no_alloc = true; + gip.ctx = &meta; + gguf_context * g = gguf_init_from_file(path.c_str(), gip); + if (!g) { set_err("gguf_init failed: " + path); return false; } + + // ── Arch check ────────────────────────────────────────────────────── + { + const int64_t aid = gguf_find_key(g, "general.architecture"); + if (aid < 0) { set_err("missing general.architecture"); gguf_free(g); if (meta) ggml_free(meta); return false; } + const char * arch = gguf_get_val_str(g, aid); + if (std::string(arch) != ARCH) { + set_err(std::string("unexpected arch: ") + arch + " (expected " + ARCH + ")"); + gguf_free(g); if (meta) ggml_free(meta); return false; + } + } + + const std::string P = std::string(ARCH) + "."; + DeepSeek4Weights & w = out.core; + + // ── Core hparams (mirror the target loader's defaults) ────────────── + const uint32_t n_layer = kv_u32(g, P + "block_count", 3); + w.n_layer = (int)n_layer; + w.n_embd = (int)kv_u32(g, P + "embedding_length", 4096); + w.n_vocab = (int)kv_u32(g, P + "vocab_size", 129280); + w.n_head = (int)kv_u32(g, P + "attention.head_count", 64); + w.n_head_kv = (int)kv_u32(g, P + "attention.head_count_kv", 1); + w.head_dim = (int)kv_u32(g, P + "attention.key_length", 512); + w.n_rot = (int)kv_u32(g, P + "rope.dimension_count", 64); + w.n_lora_q = (int)kv_u32(g, P + "attention.q_lora_rank", 1024); + w.n_lora_o = (int)kv_u32(g, P + "attention.output_lora_rank", 1024); + w.n_out_group = (int)kv_u32(g, P + "attention.output_group_count", 8); + w.n_expert = (int)kv_u32(g, P + "expert_count", 256); + w.n_expert_used = (int)kv_u32(g, P + "expert_used_count", 6); + w.n_expert_shared = (int)kv_u32(g, P + "expert_shared_count", 1); + w.n_ff_exp = (int)kv_u32(g, P + "expert_feed_forward_length", 2048); + w.n_hash_layer = (int)kv_u32(g, P + "hash_layer_count", 3); + w.n_swa = (int)kv_u32(g, P + "attention.sliding_window", 128); + w.n_indexer_head = (int)kv_u32(g, P + "attention.indexer.head_count", 64); + w.n_indexer_head_dim = (int)kv_u32(g, P + "attention.indexer.key_length", 128); + w.n_indexer_top_k = (int)kv_u32(g, P + "attention.indexer.top_k", 512); + w.n_hc = (int)kv_u32(g, P + "hyper_connection.count", 4); + w.n_hc_sinkhorn_iter = (int)kv_u32(g, P + "hyper_connection.sinkhorn_iterations", 20); + w.expert_weight_scale = kv_f32(g, P + "expert_weights_scale", 1.5f); + w.rope_freq_base = kv_f32(g, P + "rope.freq_base", 10000.0f); + w.rope_scale_factor = kv_f32(g, P + "rope.scaling.factor", 16.0f); + w.rope_yarn_beta_fast = kv_f32(g, P + "rope.scaling.yarn_beta_fast", 32.0f); + w.rope_yarn_beta_slow = kv_f32(g, P + "rope.scaling.yarn_beta_slow", 1.0f); + w.compress_rope_freq_base = kv_f32(g, P + "attention.compress_rope_freq_base", 160000.0f); + w.rms_eps = kv_f32(g, P + "attention.layer_norm_rms_epsilon", 1e-6f); + w.hc_eps = kv_f32(g, P + "hyper_connection.epsilon", 1e-6f); + w.swiglu_clamp_exp = kv_f32(g, P + "swiglu_clamp_exp", 10.0f); + w.eos_id = -1; // drafter carries no tokenizer; EOS comes from the target + w.eos_chat_id = -1; + + // DSpark drafter layers have NO KV compression (DSparkAttention asserts + // compress_ratio==0). Force all-zero so no compressor/indexer tensors are + // expected — do NOT run compute_compress_ratios (would give [0,0,4,...]). + w.compress_ratios.assign(n_layer, 0u); + + // The DSpark blocks are NOT hash-routing layers: in the checkpoint their + // real layer ids are n_layers+stage = 43/44/45, all >= num_hash_layers (3), + // and the drafter GGUF carries no ffn_gate_tid2eid. Force 0 so build_moe_ffn + // always takes the score-routed (sqrt-softplus + top-k) path. + w.n_hash_layer = 0; + + // ── DFlash / DSpark metadata ──────────────────────────────────────── + out.n_target_layers = (int)kv_u32(g, P + "dflash.n_target_layers", 3); + out.block_size = (int)kv_u32(g, P + "dflash.block_size", 5); + out.mask_token_id = (int)kv_u32(g, P + "dflash.mask_token_id", 128799); + out.head_hc_enabled = kv_u32(g, P + "dflash.head_hc_enabled", 0) != 0; + if (!kv_i32_array(g, P + "dflash.capture_layer_ids", out.capture_layer_ids)) { + out.capture_layer_ids.clear(); + } + out.dspark_enabled = kv_u32(g, P + "dflash.dspark.enabled", 0) != 0; + out.markov_rank = (int)kv_u32(g, P + "dflash.dspark.markov_rank", 256); + out.vocab_size = (int)kv_u32(g, P + "dflash.dspark.vocab_size", w.n_vocab); + out.confidence_dim = (int)kv_u32(g, P + "dflash.dspark.confidence_dim", 0); + + w.layers.resize(n_layer); + w.backend = backend; + + // ── Allocate all tensors from the meta ctx into one backend buffer ── + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(meta, backend); + if (!buf) { set_err("ggml_backend_alloc_ctx_tensors failed"); gguf_free(g); ggml_free(meta); return false; } + ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + // ── Stream tensor bytes from the file (pread per tensor) ──────────── + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { set_err("open failed: " + path); ggml_backend_buffer_free(buf); gguf_free(g); ggml_free(meta); return false; } + const size_t data_off = gguf_get_data_offset(g); + const int64_t n_tensors = gguf_get_n_tensors(g); + std::vector staging; + bool ok = true; + for (int64_t ti = 0; ti < n_tensors && ok; ti++) { + const char * tname = gguf_get_tensor_name(g, ti); + ggml_tensor * t = ggml_get_tensor(meta, tname); + if (!t) { set_err(std::string("meta tensor missing: ") + tname); ok = false; break; } + const size_t off = data_off + gguf_get_tensor_offset(g, ti); + const size_t size = gguf_get_tensor_size(g, ti); + staging.resize(size); + size_t done = 0; + while (done < size) { + const ssize_t r = ::pread(fd, staging.data() + done, size - done, (off_t)(off + done)); + if (r <= 0) { set_err(std::string("pread failed for ") + tname); ok = false; break; } + done += (size_t)r; + } + if (!ok) break; + ggml_backend_tensor_set(t, staging.data(), 0, size); + } + ::close(fd); + if (!ok) { ggml_backend_buffer_free(buf); gguf_free(g); ggml_free(meta); return false; } + + // ── Bind pointers by name ─────────────────────────────────────────── + for (int64_t ti = 0; ti < n_tensors; ti++) { + const char * name = gguf_get_tensor_name(g, ti); + ggml_tensor * t = ggml_get_tensor(meta, name); + + if (std::strcmp(name, "output_norm.weight") == 0) { w.out_norm = t; continue; } + if (std::strcmp(name, "output_hc_base.weight") == 0) { w.output_hc_base = t; continue; } + if (std::strcmp(name, "output_hc_fn.weight") == 0) { w.output_hc_fn = t; continue; } + if (std::strcmp(name, "output_hc_scale.weight") == 0) { w.output_hc_scale = t; continue; } + if (std::strcmp(name, "dflash.fc.weight") == 0) { out.main_proj = t; continue; } + if (std::strcmp(name, "dflash.hidden_norm.weight") == 0) { out.main_norm = t; continue; } + if (std::strcmp(name, "dflash.dspark.markov.w1") == 0) { out.markov_w1 = t; continue; } + if (std::strcmp(name, "dflash.dspark.markov.w2") == 0) { out.markov_w2 = t; continue; } + if (std::strcmp(name, "dflash.dspark.confidence.weight") == 0) { out.confidence_w = t; continue; } + if (std::strcmp(name, "dflash.dspark.confidence.bias") == 0) { out.confidence_b = t; continue; } + + int il = -1; std::string suffix; + if (!block_suffix(name, il, suffix) || il < 0 || il >= (int)n_layer) continue; + DeepSeek4Layer & L = w.layers[il]; + + if (suffix == "attn_norm.weight") L.attn_norm = t; + else if (suffix == "attn_q_a.weight") L.attn_q_a = t; + else if (suffix == "attn_q_a_norm.weight") L.attn_q_a_norm = t; + else if (suffix == "attn_q_b.weight") L.attn_q_b = t; + else if (suffix == "attn_kv.weight") L.attn_kv = t; + else if (suffix == "attn_kv_a_norm.weight") L.attn_kv_a_norm = t; + else if (suffix == "attn_sinks.weight") L.attn_sinks = t; + else if (suffix == "attn_output_a.weight") L.attn_output_a = t; + else if (suffix == "attn_output_b.weight") L.attn_output_b = t; + else if (suffix == "hc_attn_fn.weight") L.hc_attn_fn = t; + else if (suffix == "hc_attn_scale.weight") L.hc_attn_scale = t; + else if (suffix == "hc_attn_base.weight") L.hc_attn_base = t; + else if (suffix == "ffn_norm.weight") L.ffn_norm = t; + else if (suffix == "ffn_gate_inp.weight") L.ffn_gate_inp = t; + else if (suffix == "exp_probs_b.bias") L.ffn_exp_probs_b = t; + else if (suffix == "ffn_gate_tid2eid.weight") L.ffn_gate_tid2eid = t; + else if (suffix == "ffn_gate_exps.weight") L.ffn_gate_exps = t; + else if (suffix == "ffn_up_exps.weight") L.ffn_up_exps = t; + else if (suffix == "ffn_down_exps.weight") L.ffn_down_exps = t; + else if (suffix == "ffn_gate_shexp.weight") L.ffn_gate_shexp = t; + else if (suffix == "ffn_up_shexp.weight") L.ffn_up_shexp = t; + else if (suffix == "ffn_down_shexp.weight") L.ffn_down_shexp = t; + else if (suffix == "hc_ffn_fn.weight") L.hc_ffn_fn = t; + else if (suffix == "hc_ffn_scale.weight") L.hc_ffn_scale = t; + else if (suffix == "hc_ffn_base.weight") L.hc_ffn_base = t; + } + + w.ctx = meta; + w.buf = buf; + gguf_free(g); // meta_ctx now owned by w.ctx; do not free here + + std::fprintf(stderr, + "[ds4-dspark] loaded %s: n_layer=%d n_embd=%d vocab=%d block_size=%d " + "n_target_layers=%d markov_rank=%d confidence_dim=%d mask_tok=%d dspark=%d head_hc=%d\n", + path.c_str(), w.n_layer, w.n_embd, w.n_vocab, out.block_size, + out.n_target_layers, out.markov_rank, out.confidence_dim, out.mask_token_id, + (int)out.dspark_enabled, (int)out.head_hc_enabled); + return true; +} + +void free_deepseek4_dspark_drafter(DSparkDrafter & d) { + if (d.core.buf) { ggml_backend_buffer_free(d.core.buf); d.core.buf = nullptr; } + if (d.core.ctx) { ggml_free(d.core.ctx); d.core.ctx = nullptr; } + d = DSparkDrafter{}; +} + +// Validation dump used by the load smoke test. +void deepseek4_dspark_dump(const DSparkDrafter & d) { + const DeepSeek4Weights & w = d.core; + auto shp = [](ggml_tensor * t) -> std::string { + if (!t) return "NULL"; + char b[128]; + std::snprintf(b, sizeof(b), "[%lld,%lld,%lld,%lld] %s", + (long long)t->ne[0], (long long)t->ne[1], (long long)t->ne[2], (long long)t->ne[3], + ggml_type_name(t->type)); + return b; + }; + std::fprintf(stderr, "── DSpark drafter dump ──\n"); + std::fprintf(stderr, " main_proj %s\n", shp(d.main_proj).c_str()); + std::fprintf(stderr, " main_norm %s\n", shp(d.main_norm).c_str()); + std::fprintf(stderr, " out_norm %s\n", shp(w.out_norm).c_str()); + std::fprintf(stderr, " output_hc_fn %s\n", shp(w.output_hc_fn).c_str()); + std::fprintf(stderr, " markov_w1 %s\n", shp(d.markov_w1).c_str()); + std::fprintf(stderr, " markov_w2 %s\n", shp(d.markov_w2).c_str()); + std::fprintf(stderr, " conf_w %s\n", shp(d.confidence_w).c_str()); + std::fprintf(stderr, " conf_b %s\n", shp(d.confidence_b).c_str()); + for (int il = 0; il < w.n_layer; il++) { + const DeepSeek4Layer & L = w.layers[il]; + std::fprintf(stderr, " blk.%d: attn_norm=%s q_a=%s q_b=%s kv=%s o_a=%s o_b=%s sinks=%s\n", + il, shp(L.attn_norm).c_str(), shp(L.attn_q_a).c_str(), shp(L.attn_q_b).c_str(), + shp(L.attn_kv).c_str(), shp(L.attn_output_a).c_str(), shp(L.attn_output_b).c_str(), + shp(L.attn_sinks).c_str()); + std::fprintf(stderr, " ffn_gate_exps=%s up=%s down=%s shexp(g=%s) gate_inp=%s probs_b=%s hc_attn_fn=%s hc_ffn_fn=%s\n", + shp(L.ffn_gate_exps).c_str(), shp(L.ffn_up_exps).c_str(), shp(L.ffn_down_exps).c_str(), + shp(L.ffn_gate_shexp).c_str(), shp(L.ffn_gate_inp).c_str(), shp(L.ffn_exp_probs_b).c_str(), + shp(L.hc_attn_fn).c_str(), shp(L.hc_ffn_fn).c_str()); + } +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h new file mode 100644 index 000000000..9fe20da75 --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -0,0 +1,141 @@ +// DeepSeek-V4-Flash "DSpark" speculative-decode drafter. +// +// The DSpark drafter is a small (n_layer≈3) DeepSeek-V4 block stack stored under +// the checkpoint's mtp.* namespace and converted to a GGUF with arch +// "deepseek4-dflash-draft" (see scripts/convert_ds4_dspark_draft_to_gguf.py). +// +// Reference forward: deepseek-ai/DeepSeek-V4-Flash-DSpark inference/model.py +// (DSparkBlock / DSparkAttention / DSparkMarkovHead / DSparkConfidenceHead, +// Transformer.forward_spec). Key facts encoded here: +// - The target captures h.mean(dim=2) (mean over the hc_mult HC copies) after +// each of dspark_target_layer_ids (=[40,41,42]) -> main_hidden [.., 3*n_embd]. +// - forward_embed (stage 0): main_x = main_norm(main_proj(main_hidden)); the +// noise block = embed([seed]+[MASK]*(block_size-1)), HC-expanded. +// - DSparkAttention: compress_ratio==0 (no compressor/indexer). Each layer +// projects the shared main_x -> main_kv via its own wkv, and the block's +// block_size query positions attend BIDIRECTIONALLY over +// [sliding-window main-context KV] ++ [block KV]. +// - forward_head (last stage): hc_head collapse -> out_norm -> tied target +// lm_head -> per-position Markov correction + confidence gate. +// +// The drafter's token embedding and lm_head are TIED to the target, so the +// drafter GGUF carries neither token_embd nor output; embedding + projection +// go through the DeepSeek4DFlashTarget adapter. + +#pragma once + +#include "deepseek4_internal.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include + +namespace dflash::common { + +// The drafter weights. `core` reuses DeepSeek4Weights for the n_layer decoder +// blocks + per-layer tensors + metadata + out_norm + output_hc_* tail; its +// tok_embd/output stay null (tied to target). The DSpark-specific tensors below +// live in the same ggml context / backend buffer as `core`. +struct DSparkDrafter { + DeepSeek4Weights core; + + // Captured-feature fusion (stage 0 only in the checkpoint, but stored global). + ggml_tensor * main_proj = nullptr; // dflash.fc.weight [n_tgt*n_embd, n_embd] + ggml_tensor * main_norm = nullptr; // dflash.hidden_norm.weight [n_embd] + + // DSpark heads (last stage). + ggml_tensor * markov_w1 = nullptr; // dflash.dspark.markov.w1 [markov_rank, vocab] + ggml_tensor * markov_w2 = nullptr; // dflash.dspark.markov.w2 [markov_rank, vocab] + ggml_tensor * confidence_w = nullptr; // dflash.dspark.confidence.weight [confidence_dim, 1] + ggml_tensor * confidence_b = nullptr; // dflash.dspark.confidence.bias [1] + + int block_size = 5; + int n_target_layers = 3; + int markov_rank = 256; + int vocab_size = 129280; + int confidence_dim = 0; + int mask_token_id = 128799; + bool dspark_enabled = false; + bool head_hc_enabled = false; + std::vector capture_layer_ids; // [40,41,42] +}; + +// Load a "deepseek4-dflash-draft" GGUF into `out`. Returns false on error; +// deepseek4_dspark_last_error() has the message. +bool load_deepseek4_dspark_drafter(const std::string & path, + ggml_backend_t backend, + DSparkDrafter & out); + +void free_deepseek4_dspark_drafter(DSparkDrafter & d); + +const char * deepseek4_dspark_last_error(); + +// One drafter forward. Produces block_size normed hidden states (the input to +// the tied lm_head + Markov head), conditioned on a window of captured target +// features. All host-side f32 for a simple v1 (GPU feature-ring plumbing can +// come later). +// +// seed_tok : the committed anchor token (block position 0's embedding) +// noise_embed : [n_embd * block_size] embeds of [seed]+[MASK]*(block_size-1) +// ctx_features : [n_target_layers*n_embd * ctx_len] captured features, +// ordered oldest..newest, absolute positions +// [committed-ctx_len .. committed-1] +// ctx_len : number of context feature columns (<= n_swa) +// committed : absolute position of the seed (block position 0) +// out_hidden : filled with [n_embd * block_size] = out_norm(hc_head(block)) +// +// Defined in deepseek4_graph.cpp (needs the static DS4 sub-builders). +bool deepseek4_dspark_draft_forward(ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + const float * ctx_features, + int ctx_len, + int committed, + std::vector & out_hidden); + +// Batched target verify forward WITH feature capture (defined in +// deepseek4_graph.cpp so it can reuse the target sub-builders). Runs the DS4 +// target over `n_tokens` embeddings at absolute position `kv_start` in one +// forward and returns: +// argmax_out : per-position argmax token id (size n_tokens) +// logits_out : if non-null, full [n_tokens * n_vocab] f32 logits +// capture_out : [n_target_layers*n_embd * n_tokens] f32, per-position mean +// over the hc_mult HC copies after each capture layer +// (concatenated in capture_layer_ids order) — the drafter's +// main_hidden feed. +// Advances/updates the target cache exactly like a decode of these tokens. +bool deepseek4_dspark_verify_forward(ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & capture_layer_ids, + const float * embed, + const int32_t * token_ids, + int n_tokens, + int kv_start, + std::vector & argmax_out, + std::vector * logits_out, + std::vector & capture_out, + DeepSeek4StepTelemetry * telemetry = nullptr, + bool allow_graph_reuse = true); + +// Run DSpark speculative decode: draft block_size candidates with `drafter`, +// verify against the DS4 target in one batched forward, accept the matching +// prefix, and loop. Returns generated tokens via `io.emit`. Mirrors the laguna +// DSpark loop. accept_rate_out (optional) gets mean accepted / block. +struct GenerateRequest; // fwd (from common/…); the loop only needs n_gen + committed +bool run_deepseek4_dspark_spec_decode( + ggml_backend_t backend, + const DeepSeek4Weights & target_w, + DeepSeek4Cache & target_cache, + const DSparkDrafter & drafter, + int committed, + int last_tok, + int n_gen, + const float * prompt_feature_window, // [n_target_layers*n_embd * win_len] captured during prefill + int win_len, + std::vector & out_tokens, + float * accept_rate_out); + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp new file mode 100644 index 000000000..f3cec9b3d --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -0,0 +1,686 @@ +// DeepSeek-V4-Flash DSpark speculative decode: DFlashTarget adapter + spec loop. +// +// The drafter (DSparkDrafter, deepseek4_dspark.{h,cpp}) proposes block_size +// candidates conditioned on captured target features; the DS4 target verifies +// them in one batched forward. Feature capture + verify live in +// deepseek4_dspark_verify_forward (deepseek4_graph.cpp). The DSpark Markov head +// (common/dspark_head.cpp) is target-agnostic and reused verbatim. +// +// Fast spec loop (default): ONE batched verify per step, verify width capped +// at DS4_SPEC_Q=4 tokens (seed + 3 candidates). With q <= ratio(4) the verify +// crosses at most one compression boundary and never aliases rolling-state +// rows, so rejection rollback needs no KV snapshot at all: +// - raw ring rows are position-addressed (pos % n_swa) -> idempotent, +// - comp rows are index-addressed (pos / ratio) -> idempotent, +// - n_comp / n_index_comp are pure functions of commit position, +// - the ONLY non-idempotent state is the ratio-4 compressor prev-half +// (4 rows/state, flushed cur->prev at chunk boundaries), a few KB/layer, +// saved host-side before the verify and restored only when the flush +// happened at-or-past the rollback point. +// The legacy full-snapshot + double-verify path is kept behind +// DFLASH_DS4_FULL_SNAP=1 for A/B validation. + +#include "deepseek4_dspark.h" +#include "deepseek4_internal.h" +#include "internal.h" +#include "common/dspark_head.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +// ── DFlashTarget adapter over the DS4 target ──────────────────────────── +class DeepSeek4DFlashTarget : public DFlashTarget { +public: + DeepSeek4DFlashTarget(const DeepSeek4Weights & w, DeepSeek4Cache & cache, + ggml_backend_t backend, ggml_backend_t snap_backend, + std::vector capture_ids, int mask_tok) + : w_(w), cache_(cache), backend_(backend), snap_backend_(snap_backend), + capture_ids_(std::move(capture_ids)), mask_tok_(mask_tok) {} + + bool verify_batch(const std::vector & tokens, int base_pos, int & last_tok, + std::vector * all_argmax = nullptr, + bool capture_ssm_intermediates = false) override { + (void) capture_ssm_intermediates; + const int n = (int) tokens.size(); + embed_buf_.resize((size_t) n * w_.n_embd); + if (!w_.embedder.embed(tokens.data(), n, embed_buf_.data())) { + std::fprintf(stderr, "[ds4-verify] embed FAILED n=%d tok0=%d tok1=%d vocab=%d\n", + n, n > 0 ? tokens[0] : -1, n > 1 ? tokens[1] : -1, w_.n_vocab); + return false; + } + // Sequential exact verify (measurement mode): q single-token forwards + // through the legacy AR decode path. Causal by construction, compressor + // fed every token. Slow; used to measure the drafter's TRUE accept + // rate and produce exact greedy output. Enable: DFLASH_DS4_SEQ_VERIFY=1 + // (pair with DFLASH_DS4_FULL_SNAP=1 so rollback/replay stay exact). + static const bool seq_verify = [] { + const char * v = std::getenv("DFLASH_DS4_SEQ_VERIFY"); + return v && *v && *v != '0'; + }(); + if (seq_verify) { + std::vector am_all; + std::vector feat_all; + std::vector logits_all; + am_all.reserve(n); + for (int t = 0; t < n; t++) { + std::vector am1; + std::vector feat1; + std::vector logits1; + if (!deepseek4_dspark_verify_forward(backend_, w_, cache_, capture_ids_, + embed_buf_.data() + (size_t) t * w_.n_embd, + tokens.data() + t, 1, base_pos + t, am1, + keep_logits_ ? &logits1 : nullptr, + feat1, telemetry_, + /*allow_graph_reuse=*/false)) { + return false; + } + if (am1.empty()) return false; + am_all.push_back(am1[0]); + feat_all.insert(feat_all.end(), feat1.begin(), feat1.end()); + if (keep_logits_) logits_all.insert(logits_all.end(), logits1.begin(), logits1.end()); + } + verify_features_ = std::move(feat_all); + if (keep_logits_) verify_logits_ = std::move(logits_all); + last_tok = am_all.back(); + verify_n_ = n; + if (all_argmax) *all_argmax = std::move(am_all); + return true; + } + std::vector am; + // n==1 must take the dynamic (non-reuse) path: the reused decode graph + // skips the capture/all-logits hooks (backend HC), which this needs. + if (!deepseek4_dspark_verify_forward(backend_, w_, cache_, capture_ids_, + embed_buf_.data(), tokens.data(), n, base_pos, am, + keep_logits_ ? &verify_logits_ : nullptr, + verify_features_, telemetry_, + /*allow_graph_reuse=*/n > 1)) { + return false; + } + if (am.empty()) return false; + last_tok = am.back(); + verify_n_ = n; + if (all_argmax) *all_argmax = std::move(am); + return true; + } + + bool read_verify_logits(int n_tokens, std::vector & out) override { + if (!keep_logits_ || verify_logits_.empty()) return false; + const size_t need = (size_t) n_tokens * w_.n_vocab; + if (verify_logits_.size() < need) return false; + out.assign(verify_logits_.begin(), verify_logits_.begin() + need); + return true; + } + + bool snapshot_kv() override { return deepseek4_snapshot_save(cache_, snap_backend_, snap_); } + bool restore_kv() override { return deepseek4_snapshot_restore(snap_, cache_); } + + bool is_eos(int token) const override { return deepseek4_is_eos_tok(token, w_); } + + bool embed_tokens(const int32_t * tokens, int n, float * out) const override { + return w_.embedder.embed(tokens, n, out); + } + + bool project_hidden_to_tokens(const float * hidden, int n_tokens, + std::vector & tokens_out) override { + std::vector logits; + if (!project_hidden_to_logits(hidden, n_tokens, logits)) return false; + tokens_out.resize(n_tokens); + for (int t = 0; t < n_tokens; t++) { + const float * row = logits.data() + (size_t) t * w_.n_vocab; + int best = 0; float bv = row[0]; + for (int i = 1; i < w_.n_vocab; i++) if (row[i] > bv) { bv = row[i]; best = i; } + tokens_out[t] = best; + } + return true; + } + + // The drafter hidden is already out_norm'd (drafter tail); project with the + // tied target lm_head only (mul_mat, no norm), matching the reference head. + bool project_hidden_to_logits(const float * hidden, int n_tokens, + std::vector & logits_out) override { + if (n_tokens <= 0) return false; + ggml_init_params ip{}; + ip.mem_size = 32u * 1024 * 1024; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * h = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, w_.n_embd, n_tokens); + ggml_set_input(h); + ggml_tensor * logits = ggml_mul_mat(ctx, w_.output, h); // [n_vocab, n_tokens] + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, gf)) { + if (alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(h, hidden, 0, sizeof(float) * (size_t) w_.n_embd * n_tokens); + const ggml_status st = ggml_backend_graph_compute(backend_, gf); + if (st != GGML_STATUS_SUCCESS) { ggml_gallocr_free(alloc); ggml_free(ctx); return false; } + logits_out.resize((size_t) n_tokens * w_.n_vocab); + ggml_backend_tensor_get(logits, logits_out.data(), 0, sizeof(float) * logits_out.size()); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return true; + } + + ggml_tensor * lm_head_tensor() override { return w_.output; } + int hidden_size() const override { return w_.n_embd; } + int mask_token_id() const override { return mask_tok_; } + const std::vector & capture_layer_ids() const override { return capture_ids_; } + + void set_keep_logits(bool b) { keep_logits_ = b; } + void set_telemetry(DeepSeek4StepTelemetry * t) { telemetry_ = t; } + const std::vector & last_features() const { return verify_features_; } + int last_verify_n() const { return verify_n_; } + +private: + const DeepSeek4Weights & w_; + DeepSeek4Cache & cache_; + ggml_backend_t backend_; + ggml_backend_t snap_backend_; + std::vector capture_ids_; + int mask_tok_; + DeepSeek4Snapshot snap_{}; + DeepSeek4StepTelemetry * telemetry_ = nullptr; + bool keep_logits_ = false; + int verify_n_ = 0; + std::vector embed_buf_; + std::vector verify_logits_; + std::vector verify_features_; +}; + +namespace { + +// Build the DraftWeights shim the target-agnostic dspark head expects (only the +// DSpark head fields + n_embd are read). +DraftWeights make_dspark_shim(const DSparkDrafter & d) { + DraftWeights dw{}; + dw.n_embd = d.core.n_embd; + dw.dspark.enabled = d.dspark_enabled; + dw.dspark.markov_rank = d.markov_rank; + dw.dspark.vocab_size = d.vocab_size; + dw.dspark.confidence_dim = d.confidence_dim; + dw.dspark.markov_w1 = d.markov_w1; + dw.dspark.markov_w2 = d.markov_w2; + dw.dspark.confidence_w = d.confidence_w; + dw.dspark.confidence_b = d.confidence_b; + return dw; +} + +bool spec_env_flag(const char * name) { + const char * v = std::getenv(name); + return v && *v && *v != '0'; +} + +// ── Light rollback state ──────────────────────────────────────────────── +// Only the non-position-idempotent cache pieces: the prev-half rows of every +// ratio-4 rolling compressor state (attn + indexer) plus hc_state. Everything +// else (raw ring, comp rows, counters) is position/index-addressed and either +// overwritten idempotently or recomputable from the commit position. +struct SpecRollback { + int cur_pos = 0; + struct Layer { + std::vector attn_kv, attn_sc, idx_kv, idx_sc; + }; + std::vector layers; + std::vector hc_state; +}; + +// prev-half = first 4 rows of a [comp_width, 8] ratio-4 rolling state. +// ratio-128 states ([comp_width, 128]) are pure position rings -> skip. +void save_prev_half(ggml_tensor * t, std::vector & buf) { + if (!t || t->ne[1] != 8) { buf.clear(); return; } + const size_t bytes = (size_t) t->nb[1] * 4; + if (buf.size() != bytes) buf.resize(bytes); + ggml_backend_tensor_get(t, buf.data(), 0, bytes); +} + +void restore_prev_half(ggml_tensor * t, const std::vector & buf) { + if (!t || buf.empty()) return; + ggml_backend_tensor_set(t, buf.data(), 0, buf.size()); +} + +void spec_rollback_save(const DeepSeek4Cache & cache, SpecRollback & rb) { + rb.cur_pos = cache.cur_pos; + rb.layers.resize(cache.layers.size()); + for (size_t il = 0; il < cache.layers.size(); ++il) { + const DeepSeek4LayerCache & lc = cache.layers[il]; + SpecRollback::Layer & s = rb.layers[il]; + save_prev_half(lc.attn_compressor.state_kv, s.attn_kv); + save_prev_half(lc.attn_compressor.state_score, s.attn_sc); + save_prev_half(lc.indexer_compressor.state_kv, s.idx_kv); + save_prev_half(lc.indexer_compressor.state_score, s.idx_sc); + } + if (cache.hc_state) { + const size_t bytes = ggml_nbytes(cache.hc_state); + if (rb.hc_state.size() != bytes) rb.hc_state.resize(bytes); + ggml_backend_tensor_get(cache.hc_state, rb.hc_state.data(), 0, bytes); + } +} + +// Truncate the cache to commit_pos. restore_prev is set when the verify +// crossed a ratio-4 boundary at-or-past commit_pos: that flush filled the +// prev-half rows with a chunk containing rejected tokens, so put the +// pre-verify rows back. (A boundary strictly inside the committed range is a +// legitimate flush and must be kept.) +void spec_rollback_apply(const SpecRollback & rb, const DeepSeek4Weights & w, + DeepSeek4Cache & cache, int commit_pos, bool restore_prev) { + cache.cur_pos = commit_pos; + for (size_t il = 0; il < cache.layers.size(); ++il) { + DeepSeek4LayerCache & lc = cache.layers[il]; + const uint32_t ratio = il < w.compress_ratios.size() ? w.compress_ratios[il] : 0; + if (ratio > 0) lc.n_comp = commit_pos / (int) ratio; + if (ratio == 4) lc.n_index_comp = commit_pos / 4; + if (restore_prev && il < rb.layers.size()) { + const SpecRollback::Layer & s = rb.layers[il]; + restore_prev_half(lc.attn_compressor.state_kv, s.attn_kv); + restore_prev_half(lc.attn_compressor.state_score, s.attn_sc); + restore_prev_half(lc.indexer_compressor.state_kv, s.idx_kv); + restore_prev_half(lc.indexer_compressor.state_score, s.idx_sc); + } + } + if (restore_prev && cache.hc_state && !rb.hc_state.empty()) { + ggml_backend_tensor_set(cache.hc_state, rb.hc_state.data(), 0, rb.hc_state.size()); + } +} + +using SpecClock = std::chrono::steady_clock; + +double spec_ms_since(SpecClock::time_point t0) { + return std::chrono::duration_cast(SpecClock::now() - t0).count() / 1000.0; +} + +} // namespace + +// Batched target verify + capture: wraps the existing multi-token +// deepseek4_step_layer_range (dynamic attention + batched HC), which never +// touches the fused single-token 23 tok/s path, with the Ds4VerifyHooks that +// add per-layer mean-over-HC capture and full per-position logits. +bool deepseek4_dspark_verify_forward(ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & capture_layer_ids, + const float * embed, + const int32_t * token_ids, + int n_tokens, + int kv_start, + std::vector & argmax_out, + std::vector * logits_out, + std::vector & capture_out, + DeepSeek4StepTelemetry * telemetry, + bool allow_graph_reuse) { + std::vector hc_state; + std::vector all_logits; + std::vector last_logits; + Ds4VerifyHooks hooks; + hooks.capture_layer_ids = &capture_layer_ids; + hooks.capture_out = &capture_out; + hooks.all_logits_out = &all_logits; + if (!deepseek4_step_layer_range(backend, w, cache, hc_state, embed, n_tokens, kv_start, + 0, w.n_layer, &last_logits, token_ids, + telemetry, allow_graph_reuse, + &hooks)) { + std::fprintf(stderr, "[ds4-verify] step_layer_range returned false (n_tokens=%d kv_start=%d)\n", + n_tokens, kv_start); + return false; + } + if ((int) all_logits.size() < w.n_vocab * n_tokens) { + std::fprintf(stderr, "[ds4-verify] all_logits too small: got=%zu need=%d (cap=%zu)\n", + all_logits.size(), w.n_vocab * n_tokens, capture_out.size()); + return false; + } + argmax_out.resize(n_tokens); + for (int t = 0; t < n_tokens; t++) { + const float * row = all_logits.data() + (size_t) t * w.n_vocab; + int best = 0; float bv = row[0]; + for (int i = 1; i < w.n_vocab; i++) if (row[i] > bv) { bv = row[i]; best = i; } + argmax_out[t] = best; + } + if (logits_out) *logits_out = std::move(all_logits); + return true; +} + +bool run_deepseek4_dspark_spec_decode( + ggml_backend_t backend, + const DeepSeek4Weights & target_w, + DeepSeek4Cache & target_cache, + const DSparkDrafter & drafter, + int committed, + int last_tok, + int n_gen, + const float * prompt_feature_window, + int win_len, + std::vector & out_tokens, + float * accept_rate_out) { + const int n_embd = target_w.n_embd; + const int n_tgt = drafter.n_target_layers; + const int block = drafter.block_size; + const int n_swa = target_w.n_swa; + const int feat_row = n_tgt * n_embd; + + const bool debug = spec_env_flag("DFLASH_DS4_DSPARK_DEBUG"); + const bool timing = spec_env_flag("DFLASH_DS4_TIMING"); + const bool full_snap = spec_env_flag("DFLASH_DS4_FULL_SNAP"); + const bool seq_verify_mode = spec_env_flag("DFLASH_DS4_SEQ_VERIFY"); + // Laguna-style adaptive verify width: EWMA of accepted candidates, width = + // ewma + 2 (avg_commit << block means the wide tail is usually wasted). + // /tmp/ds4_awidth: 1 = on, 0 = off (default on). + bool adaptive_width = true; + if (std::FILE * f = std::fopen("/tmp/ds4_awidth", "r")) { + int v = 1; + if (std::fscanf(f, "%d", &v) == 1) adaptive_width = (v != 0); + std::fclose(f); + } + double ewma_accept = 1.5; + + // Fast path caps the verify at the compression ratio (4): one boundary max, + // no rolling-state row aliasing -> snapshot-free rollback stays exact. + // The full-snapshot A/B path may use the whole block. + int q_cap = full_snap ? block + 1 : 4; + if (const char * qs = std::getenv("DFLASH_DS4_SPEC_Q")) { + const int v = std::atoi(qs); + if (v >= 2 && v <= block + 1) q_cap = full_snap ? v : std::min(v, 4); + } + if (std::FILE * qf = std::fopen("/tmp/ds4_spec_q", "r")) { + // Per-request override for perf experiments (no server restart needed). + // q=1 disables drafting: pure AR pushed through the batched-verify path + // (diagnoses batched-vs-sequential target divergence). + int v = 0; + if (std::fscanf(qf, "%d", &v) == 1 && v >= 1 && v <= block + 1) { + q_cap = full_snap ? v : std::min(v, 4); + } + std::fclose(qf); + } + + // Snapshot backend for the legacy full-snapshot rollback path. + ggml_backend_t snap_backend = ggml_backend_cpu_init(); + if (!snap_backend) { std::fprintf(stderr, "[ds4-spec] no CPU snapshot backend\n"); return false; } + + DeepSeek4DFlashTarget target(target_w, target_cache, backend, snap_backend, + drafter.capture_layer_ids, drafter.mask_token_id); + DraftWeights dw = make_dspark_shim(drafter); + SpecRollback rollback; + DeepSeek4StepTelemetry tel{}; + if (timing) target.set_telemetry(&tel); + + // Host feature window ring [feat_row, n_swa] of absolute positions + // [committed-N .. committed-1]. Seed from the prefill window. + std::vector feat_win((size_t) feat_row * n_swa, 0.0f); + int win_have = win_len > n_swa ? n_swa : win_len; + if (prompt_feature_window && win_have > 0) { + // copy the last win_have columns of the prefill window + const int src_off = (win_len - win_have); + std::memcpy(feat_win.data(), + prompt_feature_window + (size_t) src_off * feat_row, + sizeof(float) * (size_t) feat_row * win_have); + } + int feat_count = win_have; // number of valid feature columns ending at committed-1 + + auto push_feature = [&](const float * col) { + // Shift-append one feature column (keep last n_swa). + if (feat_count >= n_swa) { + std::memmove(feat_win.data(), feat_win.data() + feat_row, + sizeof(float) * (size_t) feat_row * (n_swa - 1)); + std::memcpy(feat_win.data() + (size_t) feat_row * (n_swa - 1), col, + sizeof(float) * feat_row); + } else { + std::memcpy(feat_win.data() + (size_t) feat_row * feat_count, col, + sizeof(float) * feat_row); + feat_count++; + } + }; + + int lt = last_tok; + int pos = committed; // absolute position of the seed (block slot 0) + int n_generated = 0; + long accept_sum = 0, steps = 0; + + std::vector noise_embed((size_t) n_embd * block); + std::vector noise_ids(block); + std::vector local_hidden, padded_hidden((size_t) n_embd * (block + 1), 0.0f); + std::vector draft_tok, tgt_am; + + // Cumulative phase timings (ms). + double tm_draft = 0, tm_head = 0, tm_save = 0, tm_verify = 0, tm_apply = 0, tm_feat = 0; + const SpecClock::time_point run_t0 = SpecClock::now(); + + while (n_generated < n_gen) { + const int ctx_len = feat_count < n_swa ? feat_count : n_swa; + + // Noise block = [seed] + [MASK]*(block-1). + SpecClock::time_point t0 = SpecClock::now(); + if (q_cap >= 2) { + noise_ids[0] = lt; + for (int i = 1; i < block; i++) noise_ids[i] = drafter.mask_token_id; + if (!target.embed_tokens(noise_ids.data(), block, noise_embed.data())) break; + + // Drafter forward -> block normed hidden states. + if (!deepseek4_dspark_draft_forward(backend, drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, pos, local_hidden)) { + std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); + break; + } + } + tm_draft += spec_ms_since(t0); + + if (debug) { + size_t lh_nan = 0; double lh_ss = 0; + for (float v : local_hidden) { if (!std::isfinite(v)) lh_nan++; else lh_ss += (double) v * v; } + std::fprintf(stderr, "[ds4-spec] hidden nnan=%zu/%zu rms=%.4f ctx_len=%d\n", + lh_nan, local_hidden.size(), + lh_nan < local_hidden.size() ? std::sqrt(lh_ss / (double) local_hidden.size()) : 0.0, + ctx_len); + } + + // DSpark Markov chain over the first q_cap-1 candidates. Reference + // predicts token i+1 from block slot i, so prepend a dummy row 0 and + // let the (row-0-skipping) chain use slots 1..q-1. + t0 = SpecClock::now(); + draft_tok.clear(); + bool ds_ok = false; + // Batched-verify exactness: the batch must not cross a ratio-4 + // boundary except at its last token (state rows stay distinct and the + // comp emission matches AR). Boundaries sit at p % 4 == 3. The + // sequential verify has no such limit. + static const bool fused_verify_mode = [] { + const char * v = std::getenv("DFLASH_DS4_FUSED_VERIFY"); + return v && *v && *v != '0'; + }(); + int q_step_cap = (seq_verify_mode || fused_verify_mode) + ? std::min(q_cap, 4) + : std::min(q_cap, 4 - (pos & 3)); + if (adaptive_width && !seq_verify_mode) { + const int w_cap = (int) ewma_accept + 2; + if (w_cap < q_step_cap) q_step_cap = w_cap; + } + if (q_step_cap >= 2) { + std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), + sizeof(float) * (size_t) n_embd * block); + ds_ok = dspark_markov_correct_greedy_chain_fused(dw, backend, target.lm_head_tensor(), + padded_hidden.data(), q_step_cap, lt, draft_tok); + if (!ds_ok) { + ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, + padded_hidden.data(), q_step_cap, lt, 0.0f, draft_tok); + } + if (!ds_ok || (int) draft_tok.size() < 2) { + // Fallback: plain projection of the block hiddens. + std::vector pj; + if (!target.project_hidden_to_tokens(local_hidden.data(), q_step_cap - 1, pj)) break; + draft_tok.clear(); + draft_tok.push_back(lt); + for (int i = 0; i < q_step_cap - 1; i++) draft_tok.push_back(pj[i]); + } + } else { + draft_tok.push_back(lt); // q=1: seed only, no speculation + } + if ((int) draft_tok.size() > q_step_cap) draft_tok.resize(q_step_cap); + const int q = (int) draft_tok.size(); // seed + candidates + tm_head += spec_ms_since(t0); + + if (debug) { + std::fprintf(stderr, "[ds4-spec] dbg ds_ok=%d q=%d lt=%d draft=[%d %d %d %d]\n", + (int) ds_ok, q, lt, + q > 0 ? draft_tok[0] : -1, q > 1 ? draft_tok[1] : -1, + q > 2 ? draft_tok[2] : -1, q > 3 ? draft_tok[3] : -1); + } + + // ── Rollback state save (cheap) or legacy full snapshot ── + t0 = SpecClock::now(); + if (full_snap) { + if (!target.snapshot_kv()) { std::fprintf(stderr, "[ds4-spec] snapshot failed\n"); break; } + } else { + spec_rollback_save(target_cache, rollback); + } + tm_save += spec_ms_since(t0); + + // First ratio-4 boundary position touched by this verify (p % 4 == 3). + const int first_boundary = pos + (3 - (pos & 3)); + const bool boundary_crossed = first_boundary <= pos + q - 1; + + // ── ONE batched verify (writes cache + captures features for all q) ── + t0 = SpecClock::now(); + int verify_last = -1; + if (!target.verify_batch(draft_tok, pos, verify_last, &tgt_am)) { + if (full_snap) { + target.restore_kv(); + } else { + spec_rollback_apply(rollback, target_w, target_cache, pos, boundary_crossed); + } + std::fprintf(stderr, "[ds4-spec] verify failed\n"); + break; + } + tm_verify += spec_ms_since(t0); + + // Accept the longest matching prefix. accept counts the seed (slot 0) + // plus each candidate the target agrees with. + int accept = 1; + for (int i = 0; i < q - 1; i++) { + if (draft_tok[i + 1] == tgt_am[i]) accept++; + else break; + } + const int matched = accept - 1; // accepted candidates + const int bonus = tgt_am[accept - 1]; // target's token at the accept point + const int commit_pos = pos + accept; // seed + accepted candidates in KV + + if (timing && steps < 8 && q >= 2) { + // Alignment probe: draft candidate i should match tgt_am[i-1]. A + // consistent draft[i]==tgt_am[i] pattern instead = off-by-one. + std::fprintf(stderr, "[ds4-spec-cmp] step=%ld pos=%d draft=[%d %d %d] tgt=[%d %d %d %d] acc=%d\n", + steps, pos, + q > 1 ? draft_tok[1] : -1, q > 2 ? draft_tok[2] : -1, q > 3 ? draft_tok[3] : -1, + tgt_am[0], q > 1 ? tgt_am[1] : -1, q > 2 ? tgt_am[2] : -1, q > 3 ? tgt_am[3] : -1, + accept); + } + + // ── Rollback: truncate to the committed prefix ── + // The bonus token is DEFERRED: it becomes the next step's seed, whose + // KV is written then. + t0 = SpecClock::now(); + if (full_snap) { + // Legacy: full restore + replay the committed tokens through the + // target so ring/compressor/n_comp advance exactly. + std::vector kv_toks; + kv_toks.push_back(lt); + for (int i = 1; i < accept; i++) kv_toks.push_back(draft_tok[i]); + target.restore_kv(); + int replay_last = -1; + std::vector replay_am; + if (!target.verify_batch(kv_toks, pos, replay_last, &replay_am)) { + std::fprintf(stderr, "[ds4-spec] replay verify failed\n"); + break; + } + } else if (accept < q) { + // The prev-half flush is bad only if the boundary sits at-or-past + // the commit point (its chunk then contains rejected tokens). + const bool restore_prev = boundary_crossed && first_boundary >= commit_pos; + spec_rollback_apply(rollback, target_w, target_cache, commit_pos, restore_prev); + } + // accept == q on the fast path: cur_pos/n_comp already exact, keep. + tm_apply += spec_ms_since(t0); + + // Push the committed positions' features (slots 0..accept-1 = positions + // pos..pos+accept-1) into the drafter's context window. + t0 = SpecClock::now(); + const std::vector & feats = target.last_features(); + const int fN = full_snap ? target.last_verify_n() : accept; + for (int i = 0; i < fN; i++) push_feature(feats.data() + (size_t) i * feat_row); + tm_feat += spec_ms_since(t0); + + // Output tokens this step = accepted candidates + bonus. + bool hit_eos = false; + for (int i = 1; i <= accept; i++) { + const int t = (i < accept) ? draft_tok[i] : bonus; + out_tokens.push_back(t); + n_generated++; + if (target.is_eos(t)) { hit_eos = true; break; } + if (n_generated >= n_gen) break; + } + pos = commit_pos; // seed + accepted candidates now in KV + lt = bonus; // deferred bonus becomes next seed + accept_sum += matched; + ewma_accept = 0.7 * ewma_accept + 0.3 * (double) matched; + steps++; + if (timing && (steps <= 4 || (steps & 31) == 0)) { + std::fprintf(stderr, + "[ds4-spec-t] step=%ld q=%d acc=%d | draft=%.1f head=%.1f save=%.1f " + "verify=%.1f apply=%.1f feat=%.1f ms (cum means)\n", + steps, q, accept, + tm_draft / steps, tm_head / steps, tm_save / steps, + tm_verify / steps, tm_apply / steps, tm_feat / steps); + } + if (hit_eos) break; + } + + const double total_ms = spec_ms_since(run_t0); + if (accept_rate_out && steps > 0) { + const int denom = q_cap > 1 ? q_cap - 1 : 1; + *accept_rate_out = (float) accept_sum / (float) (steps * denom); + } + std::fprintf(stderr, "[ds4-spec] gen=%d steps=%ld mean_accept=%.2f/%d q_cap=%d full_snap=%d\n", + n_generated, steps, steps ? (double) accept_sum / steps : 0.0, q_cap - 1, q_cap, + (int) full_snap); + if (steps > 0) { + std::fprintf(stderr, + "[ds4-spec-t] TOTAL %.1f ms, %ld steps (%.1f ms/step), %d tok (%.1f tok/s) | " + "means: draft=%.1f head=%.1f save=%.1f verify=%.1f apply=%.1f feat=%.1f ms\n", + total_ms, steps, total_ms / steps, n_generated, + total_ms > 0 ? n_generated * 1000.0 / total_ms : 0.0, + tm_draft / steps, tm_head / steps, tm_save / steps, + tm_verify / steps, tm_apply / steps, tm_feat / steps); + } + if (timing && steps > 0) { + const double s = 1000.0 * steps; // us -> ms per-step means + std::fprintf(stderr, + "[ds4-spec-t] verify tel/step: hc_pre_a=%.1f attn_b=%.1f attn_c=%.1f attn_r=%.1f " + "hc_post_a=%.1f hc_pre_f=%.1f route(b/c/r/s)=%.1f/%.1f/%.1f/%.1f " + "ffn(b/c/r)=%.1f/%.1f/%.1f eval=%.1f hot=%.1f cold=%.1f comb=%.1f part=%.1f " + "ghits=%llu gbuilds=%llu ms\n", + tel.hc_pre_attn_us / s, tel.attn_build_us / s, tel.attn_compute_us / s, + tel.attn_read_us / s, tel.hc_post_attn_us / s, tel.hc_pre_ffn_us / s, + tel.route_build_us / s, tel.route_compute_us / s, tel.route_read_us / s, + tel.route_select_us / s, + tel.ffn_build_us / s, tel.ffn_compute_us / s, tel.ffn_read_us / s, + tel.ffn_eval_us / s, tel.ffn_hot_us / s, tel.ffn_cold_us / s, + tel.ffn_combine_us / s, tel.ffn_partition_us / s, + (unsigned long long) tel.ffn_hot_graph_hits, + (unsigned long long) tel.ffn_hot_graph_builds); + } + ggml_backend_free(snap_backend); + return true; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc new file mode 100644 index 000000000..824d609e0 --- /dev/null +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -0,0 +1,517 @@ +// ─── Fused batched VERIFY graph (n_tokens = q in [2,4]) ───────────────── +// One whole-model graph per (q, flush, padded-comp) shape: batched attention +// (causal mask input) + batched MoE, per-token fused HC chains, in-graph +// drafter-feature capture and q-wide logits. The spec loop guarantees a +// ratio-4 boundary can only sit at the batch's LAST token (dynamic q), so the +// compressor pools at most once per step, exactly like an AR step at the last +// position. Enable with DFLASH_DS4_FUSED_VERIFY=1. + +// tensor_set that tolerates inputs the graph never consumed (gallocr leaves +// them unallocated, e.g. comp-emission bundles in non-flush graphs). +static void ds4_fv_set(ggml_tensor * t, const void * data, size_t nbytes) { + if (t && t->buffer) ggml_backend_tensor_set(t, data, 0, nbytes); +} +static bool ds4_fused_verify_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_FUSED_VERIFY") ? 1 : 0; + } + return enabled == 1; +} + +struct Ds4FusedVerifyCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + bool disabled = false; + uint64_t counter = 0; + std::array slots; + // verify-specific inputs/outputs per slot (parallel arrays, same index) + struct Extra { + ggml_tensor * pos_q = nullptr; // i32 [q] + ggml_tensor * neg_q = nullptr; // i32 [q] + ggml_tensor * rawrows = nullptr; // i64 [1,q] + ggml_tensor * ape4 = nullptr; // i32 [q] + ggml_tensor * ape128 = nullptr; // i32 [q] + ggml_tensor * st4 = nullptr; // i64 [1,q] + ggml_tensor * st128 = nullptr; // i64 [1,q] + ggml_tensor * capture = nullptr; // f32 [n_embd*ncap*q], order [ci][t] + int q = 0; + void reset() { *this = Extra{}; } + }; + std::array extra; + void destroy() { + for (auto & s : slots) s.destroy(); + for (auto & e : extra) e.reset(); + } +}; + +static Ds4FusedVerifyCache fused_verify_graph_cache; + +static bool ds4_build_fused_verify_graph( + DeepSeek4FusedDecodeCache & mc, // fn mirrors (shared with AR fused) + DeepSeek4FusedDecodeGraph & fg, + Ds4FusedVerifyCache::Extra & ex, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + int kv_start, + int q, + bool have_token_ids, + const std::vector & capture_ids, + std::vector && shape_key) { + step_graph_free(fg.sg); + fg.reset_nodes(); + ex.reset(); + fg.hash_ids.assign((size_t) w.n_layer, nullptr); + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int token_pos = kv_start + q - 1; // last batch position + + const size_t arena_size = 256u * 1024 * 1024; + if (fg.sg.meta_arena.size() < arena_size) fg.sg.meta_arena.resize(arena_size); + ggml_init_params params{}; + params.mem_size = fg.sg.meta_arena.size(); + params.mem_buffer = fg.sg.meta_arena.data(); + params.no_alloc = true; + fg.sg.ctx = ggml_init(params); + if (!fg.sg.ctx) return false; + ggml_context * ctx = fg.sg.ctx; + fg.sg.gf = ggml_new_graph_custom(ctx, 65536, false); + ggml_cgraph * gf = fg.sg.gf; + + fg.inp_embed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, q); + ggml_set_input(fg.inp_embed); + ex.q = q; + ex.pos_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.pos_q); + ex.neg_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.neg_q); + ex.rawrows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.rawrows); + ex.ape4 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.ape4); + ex.ape128 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.ape128); + ex.st4 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.st4); + ex.st128 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.st128); + // comp emission scalars per layer: i32 {ape_row(unused batched), comp_pos}, + // i64 {comp_row}; reuse the decode-style bundles (2 i32 + 1 i64 per layer). + fg.i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2 * (int64_t) w.n_layer); + ggml_set_input(fg.i32_bundle); + fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1 * (int64_t) w.n_layer); + ggml_set_input(fg.i64_bundle); + + // mask bundle: per layer [n_swa + padded + q] rows × q tokens + int64_t mask_total = 0; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + int padded = 0; + if (ratio > 0 && cache.layers[(size_t) il].comp_kv) { + const int n_comp = ds4_comp_rows_used(cache.layers[(size_t) il].comp_kv, + cache.layers[(size_t) il].n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) cache.layers[(size_t) il].comp_kv->ne[1]); + } + mask_total += (int64_t) (w.n_swa + padded + q) * q; + } + fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); + ggml_set_input(fg.mask_bundle); + int64_t mask_off = 0; + + // per-token HC streams + std::vector hc_cur(q); + for (int t = 0; t < q; ++t) { + ggml_tensor * col = ggml_view_2d(ctx, fg.inp_embed, n_embd, 1, + fg.inp_embed->nb[1], (size_t) t * fg.inp_embed->nb[1]); + hc_cur[(size_t) t] = ggml_repeat_4d(ctx, col, n_embd, n_hc, 1, 1); + } + + std::vector capture_pieces; // order [ci][t] + capture_pieces.reserve(capture_ids.size() * (size_t) q); + + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const int ratio = (int) w.compress_ratios[il]; + + // ── HC pre (attention), per token ── + std::vector hc_flat(q), split_attn(q); + ggml_tensor * attn_in = nullptr; + for (int t = 0; t < q; ++t) { + hc_flat[(size_t) t] = ggml_reshape_1d(ctx, hc_cur[(size_t) t], (int64_t) n_embd * n_hc); + ggml_tensor * working = ds4_build_fused_hc_pre(ctx, w, hc_flat[(size_t) t], + mc.fn_attn_f16[(size_t) il], L.hc_attn_base, + hlw.attn, &split_attn[(size_t) t]); + if (!working) return false; + ggml_tensor * w2 = ggml_reshape_2d(ctx, working, n_embd, 1); + attn_in = attn_in ? ggml_concat(ctx, attn_in, w2, 1) : w2; + } + + // ── Batched attention ── + DeepSeek4AttentionGraphInputs ain{}; + ain.rope_pos = ex.pos_q; + ain.neg_pos = ex.neg_q; + ain.raw_kv_rows = ex.rawrows; + if (ratio > 0) { + ain.attn_ape_row = (ratio == 4) ? ex.ape4 : ex.ape128; + ain.attn_state_rows = (ratio == 4) ? ex.st4 : ex.st128; + ain.attn_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 2 + 1) * sizeof(int32_t)); + ain.attn_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + (size_t) il * sizeof(int64_t)); + } + if (ratio == 4) { + ain.index_ape_row = ex.ape4; + ain.index_state_rows = ex.st4; + ain.index_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 2 + 1) * sizeof(int32_t)); + ain.index_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + (size_t) il * sizeof(int64_t)); + } + int padded = 0; + if (ratio > 0 && lc.comp_kv) { + const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + const int64_t n_attn = (int64_t) w.n_swa + padded + q; + ain.attn_row_mask = ggml_view_2d(ctx, fg.mask_bundle, n_attn, q, + n_attn * sizeof(float), (size_t) mask_off * sizeof(float)); + ain.padded_comp = padded; + mask_off += n_attn * q; + + std::vector i32b; + std::vector i32ab; + std::vector i64ab; + ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); + ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, + kv_start, q, &ain, i32b, i32ab, i64ab); + if (!attn_out) return false; + if (!i32b.empty() || !i32ab.empty() || !i64ab.empty()) { + std::fprintf(stderr, "[ds4-fused-verify] layer %d dynamic bindings; cannot fuse\n", il); + return false; + } + + // ── HC post (attention) + HC pre (FFN), per token ── + ggml_tensor * ffn_in = nullptr; + std::vector split_ffn(q); + for (int t = 0; t < q; ++t) { + ggml_tensor * ao = ggml_view_2d(ctx, attn_out, n_embd, 1, + attn_out->nb[1], (size_t) t * attn_out->nb[1]); + ggml_tensor * ao_flat = ggml_reshape_1d(ctx, ggml_cont(ctx, ao), n_embd); + ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_flat[(size_t) t], ao_flat, + split_attn[(size_t) t], n_hc); + hc_cur[(size_t) t] = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + hc_flat[(size_t) t] = ggml_reshape_1d(ctx, hc_cur[(size_t) t], (int64_t) n_embd * n_hc); + ggml_tensor * fworking = ds4_build_fused_hc_pre(ctx, w, hc_flat[(size_t) t], + mc.fn_ffn_f16[(size_t) il], L.hc_ffn_base, + hlw.ffn, &split_ffn[(size_t) t]); + if (!fworking) return false; + ggml_tensor * f2 = ggml_reshape_2d(ctx, fworking, n_embd, 1); + ffn_in = ffn_in ? ggml_concat(ctx, ffn_in, f2, 1) : f2; + } + + // ── Batched FFN ── + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = nullptr; + const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && + have_token_ids && hash_tables[(size_t) il].loaded; + if (hash_routed) { + // ds4_build_hash_routed_ffn is single-token; run per token, concat. + ggml_tensor * hids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, w.n_expert_used, q); + ggml_set_input(hids); + fg.hash_ids[(size_t) il] = hids; + for (int t = 0; t < q; ++t) { + ggml_tensor * fcol = ggml_view_2d(ctx, ffn_normed, n_embd, 1, + ffn_normed->nb[1], (size_t) t * ffn_normed->nb[1]); + ggml_tensor * hcol = ggml_view_2d(ctx, hids, w.n_expert_used, 1, + hids->nb[1], (size_t) t * hids->nb[1]); + ggml_tensor * fo = ds4_build_hash_routed_ffn(ctx, w, L, ggml_cont(ctx, fcol), hcol); + if (!fo) { ffn_out = nullptr; break; } + ffn_out = ffn_out ? ggml_concat(ctx, ffn_out, fo, 1) : fo; + } + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, q); + } + if (!ffn_out) return false; + + // ── HC post (FFN), per token; capture at drafter layers ── + for (int t = 0; t < q; ++t) { + ggml_tensor * fo = ggml_view_2d(ctx, ffn_out, n_embd, 1, + ffn_out->nb[1], (size_t) t * ffn_out->nb[1]); + ggml_tensor * fo_flat = ggml_reshape_1d(ctx, ggml_cont(ctx, fo), n_embd); + ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_flat[(size_t) t], fo_flat, + split_ffn[(size_t) t], n_hc); + hc_cur[(size_t) t] = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + } + for (size_t ci = 0; ci < capture_ids.size(); ++ci) { + if (capture_ids[ci] != il) continue; + for (int t = 0; t < q; ++t) { + ggml_tensor * hs = hc_cur[(size_t) t]; // [n_embd, n_hc] + ggml_tensor * hsT = ggml_cont(ctx, ggml_transpose(ctx, hs)); // [n_hc, n_embd] + ggml_tensor * summed = ggml_sum_rows(ctx, hsT); // [1, n_embd] + ggml_tensor * mean = ggml_scale(ctx, summed, 1.0f / (float) n_hc); + capture_pieces.push_back(ggml_reshape_1d(ctx, ggml_cont(ctx, mean), n_embd)); + } + } + } + + // ── Output: per-token HC merge → batched out_norm + lm_head ── + ggml_tensor * final_all = nullptr; + for (int t = 0; t < q; ++t) { + ggml_tensor * hc_flat = ggml_reshape_1d(ctx, hc_cur[(size_t) t], (int64_t) n_embd * n_hc); + ggml_tensor * onorm = ggml_rms_norm(ctx, hc_flat, w.hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, mc.fn_out_f16, onorm); + omix = ggml_reshape_1d(ctx, omix, ggml_nelements(omix)); + ggml_tensor * obase = ds4_fused_hc_base_f32(ctx, w.output_hc_base); + if (!obase || hc_out_weights.scale_data.empty()) return false; + ggml_tensor * fe = ggml_ds4_hc_out(ctx, omix, obase, hc_flat, n_hc, + hc_out_weights.scale_data[0]); + ggml_tensor * fe2 = ggml_reshape_2d(ctx, fe, n_embd, 1); + final_all = final_all ? ggml_concat(ctx, final_all, fe2, 1) : fe2; + } + ggml_tensor * out_normed = build_rms_norm(ctx, final_all, w.out_norm, w.rms_eps); + fg.logits = ggml_mul_mat(ctx, w.output, out_normed); // [n_vocab, q] + ggml_set_output(fg.logits); + ggml_build_forward_expand(gf, fg.logits); + + if (!capture_pieces.empty()) { + ggml_tensor * cap = capture_pieces[0]; + for (size_t i = 1; i < capture_pieces.size(); ++i) { + cap = ggml_concat(ctx, cap, capture_pieces[i], 0); + } + ex.capture = cap; + ggml_set_output(ex.capture); + ggml_build_forward_expand(gf, ex.capture); + } + + if (!fg.sg.alloc) { + fg.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + if (!fg.sg.alloc || !ggml_gallocr_alloc_graph(fg.sg.alloc, fg.sg.gf)) { + std::fprintf(stderr, "[ds4-fused-verify] graph alloc failed\n"); + return false; + } + fg.shape_key = std::move(shape_key); + return true; +} + +static int ds4_try_fused_verify_step( + Ds4FusedVerifyCache & vc, + DeepSeek4FusedDecodeCache & mc, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + std::vector & hash_scratch, + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + const int32_t * token_ids, + Ds4VerifyHooks * hooks, + DeepSeek4StepTelemetry * telemetry) { + if (vc.disabled || mc.disabled) return 0; + if (!hooks || !hooks->capture_layer_ids) return 0; + if (!hc_out_weights.loaded || hc_out_weights.scale_data.empty() || + !w.output_hc_fn || !w.output_hc_base) { vc.disabled = true; return 0; } + for (int il = 0; il < w.n_layer; ++il) { + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const DeepSeek4Layer & L = w.layers[(size_t) il]; + if (!hlw.attn.loaded || hlw.attn.scale_data.size() < 3 || + !hlw.ffn.loaded || hlw.ffn.scale_data.size() < 3 || + !L.hc_attn_fn || !L.hc_ffn_fn || !L.hc_attn_base || !L.hc_ffn_base) { + vc.disabled = true; + return 0; + } + } + if (vc.owner_ctx != w.ctx || vc.backend != backend) { + vc.destroy(); + vc.owner_ctx = w.ctx; + vc.backend = backend; + } + if (mc.owner_ctx != w.ctx || mc.backend != backend) { + mc.destroy(); + mc.owner_ctx = w.ctx; + mc.backend = backend; + } + if (!ds4_fused_ensure_fn_mirrors(mc, backend, w, hc_weights, hc_out_weights)) { + vc.disabled = true; + return 0; + } + + const int q = n_tokens; + const int token_pos = kv_start + q - 1; + std::vector key; + key.reserve((size_t) w.n_layer + 3); + key.push_back(w.n_swa); + key.push_back(q); + key.push_back(token_ids ? 1 : 0); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + int padded = 0; + if (ratio > 0 && lc.comp_kv) { + const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + int b_idx = -1; + if (ratio > 0) { + for (int t = 0; t < q; ++t) { + if (((kv_start + t + 1) % ratio) == 0) { b_idx = t; break; } + } + } + key.push_back(((int64_t) padded << 4) | (int64_t) (b_idx + 1)); + } + + vc.counter++; + DeepSeek4FusedDecodeGraph * fg = nullptr; + Ds4FusedVerifyCache::Extra * ex = nullptr; + for (size_t i = 0; i < vc.slots.size(); ++i) { + if (vc.slots[i].built() && vc.slots[i].shape_key == key) { + fg = &vc.slots[i]; ex = &vc.extra[i]; break; + } + } + if (!fg) { + size_t pick = 0; + for (size_t i = 0; i < vc.slots.size(); ++i) { + if (!vc.slots[i].built()) { pick = i; break; } + if (vc.slots[i].last_use < vc.slots[pick].last_use) pick = i; + } + fg = &vc.slots[pick]; ex = &vc.extra[pick]; + const auto build_t0 = Ds4TimingClock::now(); + if (!ds4_build_fused_verify_graph(mc, *fg, *ex, backend, w, cache, + hc_weights, hc_out_weights, hash_tables, + kv_start, q, token_ids != nullptr, + *hooks->capture_layer_ids, std::move(key))) { + step_graph_free(fg->sg); + fg->reset_nodes(); + ex->reset(); + vc.disabled = true; + return 0; + } + if (telemetry) telemetry->full_graph_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + } + fg->last_use = vc.counter; + + // ── Fill inputs ── + ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); + std::vector iv(q); + std::vector lv(q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = kv_start + i; + ds4_fv_set(ex->pos_q, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = -(kv_start + i); + ds4_fv_set(ex->neg_q, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) lv[(size_t) i] = (kv_start + i) % w.n_swa; + ds4_fv_set(ex->rawrows, lv.data(), sizeof(int64_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % 4; + ds4_fv_set(ex->ape4, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) lv[(size_t) i] = 4 + (kv_start + i) % 4; + ds4_fv_set(ex->st4, lv.data(), sizeof(int64_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % 128; + ds4_fv_set(ex->ape128, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) lv[(size_t) i] = (kv_start + i) % 128; + ds4_fv_set(ex->st128, lv.data(), sizeof(int64_t) * q); + + std::vector i32v((size_t) w.n_layer * 2, 0); + std::vector i64v((size_t) w.n_layer * 1, 0); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + if (ratio > 0) { + int pos_b = -1; // boundary position inside the batch (if any) + for (int t = 0; t < q; ++t) { + if (((kv_start + t + 1) % ratio) == 0) { pos_b = kv_start + t; break; } + } + i32v[(size_t) il * 2 + 0] = token_pos % ratio; + i32v[(size_t) il * 2 + 1] = (pos_b >= 0 ? pos_b : token_pos) + 1 - ratio; + i64v[(size_t) il] = (pos_b >= 0 ? pos_b : token_pos) / ratio; + } + } + ds4_fv_set(fg->i32_bundle, i32v.data(), sizeof(int32_t) * i32v.size()); + ds4_fv_set(fg->i64_bundle, i64v.data(), sizeof(int64_t) * i64v.size()); + + // causal mask values + { + std::vector maskv((size_t) ggml_nelements(fg->mask_bundle), 0.0f); + size_t off = 0; + const int e = kv_start + q; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + int padded = 0; + if (ratio > 0 && lc.comp_kv) { + const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + const size_t n_attn = (size_t) w.n_swa + padded + q; + for (int i = 0; i < q; ++i) { + float * col = maskv.data() + off + (size_t) i * n_attn; + const int pos_i = kv_start + i; + for (int r = 0; r < w.n_swa; ++r) { + // position held by slot r AFTER this batch's ring writes + const int p_r = (e <= w.n_swa) ? (r < e ? r : -1) + : (e - 1) - ((e - 1 - r) % w.n_swa); + if (p_r < 0 || p_r > pos_i) col[r] = -1.0e30f; + } + const int vis = (ratio > 0 && lc.comp_kv) + ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i) : 0; + for (int c = 0; c < padded; ++c) { + if (c >= vis) col[(size_t) w.n_swa + c] = -1.0e30f; + } + for (int j = 0; j < q; ++j) { + const bool visible = (j > i) && (kv_start + j >= w.n_swa); + if (!visible) col[(size_t) w.n_swa + padded + j] = -1.0e30f; + } + } + off += n_attn * q; + } + ds4_fv_set(fg->mask_bundle, maskv.data(), sizeof(float) * maskv.size()); + } + + if (token_ids) { + for (int il = 0; il < w.n_layer; ++il) { + ggml_tensor * hids = fg->hash_ids[(size_t) il]; + if (!hids) continue; + const auto & ht = hash_tables[(size_t) il].ids; + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used * q); + for (int i = 0; i < q; ++i) { + std::memcpy(hash_scratch.data() + (size_t) i * n_used, + ht.data() + (size_t) token_ids[i] * (size_t) n_used, + (size_t) n_used * sizeof(int32_t)); + } + ggml_backend_tensor_set(hids, hash_scratch.data(), 0, + sizeof(int32_t) * (size_t) n_used * q); + } + } + + // ── Compute + read back ── + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, fg->sg.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[ds4-fused-verify] compute failed\n"); + return -1; + } + if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + + const int ncap = (int) hooks->capture_layer_ids->size(); + if (hooks->all_logits_out) { + hooks->all_logits_out->resize((size_t) w.n_vocab * q); + ggml_backend_tensor_get(fg->logits, hooks->all_logits_out->data(), 0, + sizeof(float) * (size_t) w.n_vocab * q); + } + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get(fg->logits, out_logits.data(), + (size_t) (q - 1) * (size_t) w.n_vocab * sizeof(float), + sizeof(float) * (size_t) w.n_vocab); + if (hooks->capture_out && ex->capture && ncap > 0) { + std::vector flat((size_t) w.n_embd * ncap * q); + ggml_backend_tensor_get(ex->capture, flat.data(), 0, sizeof(float) * flat.size()); + hooks->capture_out->assign((size_t) ncap * w.n_embd * q, 0.0f); + for (int ci = 0; ci < ncap; ++ci) { + for (int t = 0; t < q; ++t) { + const float * src = flat.data() + ((size_t) ci * q + t) * w.n_embd; + float * dst = hooks->capture_out->data() + + (size_t) t * ncap * w.n_embd + (size_t) ci * w.n_embd; + std::memcpy(dst, src, sizeof(float) * (size_t) w.n_embd); + } + } + } + return 1; +} diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 9879fee36..ec2f79cea 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -174,6 +174,11 @@ struct DeepSeek4I64ArrayBinding { std::vector values; }; +struct DeepSeek4F32ArrayBinding { + ggml_tensor * tensor = nullptr; + std::vector values; +}; + static ggml_tensor * build_rms_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * weight, float eps); static ggml_tensor * build_clamped_swiglu(ggml_context * ctx, @@ -603,7 +608,11 @@ static void build_compressor_step( ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, std::vector & i32_array_inputs, - ggml_tensor ** comp_cache_source_out = nullptr) { + ggml_tensor ** comp_cache_source_out = nullptr, + ggml_tensor * flush_rows_inp = nullptr, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -622,18 +631,89 @@ static void build_compressor_step( ggml_tensor * state_score_source = state.state_score; ggml_tensor * comp_cache_source = comp_cache; - ggml_tensor * ape_col = nullptr; - if (ape_row_inp) { - ape_col = ggml_get_rows(ctx, ape, ape_row_inp); - ape_col = ggml_reshape_2d(ctx, ape_col, comp_width, 1); - } else { - ape_col = ggml_view_2d( - ctx, ape, comp_width, 1, ape->nb[1], (size_t)pos_mod * ape->nb[1]); - ape_col = ggml_cast(ctx, ape_col, GGML_TYPE_F32); + // Causal-batch verify: every token's contribution lands in its + // position-addressed state row (the rows are distinct because the batch + // never crosses a ratio window; the boundary may only be the last token). + const bool batched_state = (cur_all != nullptr && n_tokens_all > 1 && + !state_rows_inp && kv_start_all >= 0); + if (batched_state) { + ggml_tensor * kv_all = ggml_mul_mat(ctx, kv_proj, cur_all); + ggml_tensor * sc_all = ggml_mul_mat(ctx, gate_proj, cur_all); + for (int ti = 0; ti < n_tokens_all; ti++) { + const int pm_ti = (kv_start_all + ti) % ratio; + const int row_ti = (ratio == 4) ? (ratio + pm_ti) : pm_ti; + ggml_tensor * kv_ti = ggml_view_2d(ctx, kv_all, comp_width, 1, kv_all->nb[1], + (size_t) ti * kv_all->nb[1]); + ggml_tensor * sc_ti = ggml_view_2d(ctx, sc_all, comp_width, 1, sc_all->nb[1], + (size_t) ti * sc_all->nb[1]); + ggml_tensor * ape_ti = ggml_view_2d(ctx, ape, comp_width, 1, ape->nb[1], + (size_t) pm_ti * ape->nb[1]); + sc_ti = ggml_add(ctx, sc_ti, ggml_cast(ctx, ape_ti, GGML_TYPE_F32)); + ggml_tensor * kv_slot_ti = ggml_view_2d(ctx, state.state_kv, comp_width, 1, + state.state_kv->nb[1], + (size_t) row_ti * state.state_kv->nb[1]); + ggml_tensor * sc_slot_ti = ggml_view_2d(ctx, state.state_score, comp_width, 1, + state.state_score->nb[1], + (size_t) row_ti * state.state_score->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_cast(ctx, kv_ti, state.state_kv->type), kv_slot_ti)); + ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_ti, sc_slot_ti)); + } } - sc_cur = ggml_add(ctx, sc_cur, ape_col); - if (state_rows_inp) { + const bool batched_rows = (state_rows_inp && cur_all != nullptr && n_tokens_all > 1); + int batched_b = -1; // boundary index within the batch (batched_rows) + int batched_nB = 0; // tokens after the boundary + int batched_span_off = 0; + ggml_tensor * batched_kv_all = nullptr; + ggml_tensor * batched_sc_all = nullptr; + ggml_tensor * ape_col = nullptr; + if (!batched_rows) { + if (ape_row_inp) { + ape_col = ggml_get_rows(ctx, ape, ape_row_inp); + ape_col = ggml_reshape_2d(ctx, ape_col, comp_width, 1); + } else { + ape_col = ggml_view_2d( + ctx, ape, comp_width, 1, ape->nb[1], (size_t)pos_mod * ape->nb[1]); + ape_col = ggml_cast(ctx, ape_col, GGML_TYPE_F32); + } + sc_cur = ggml_add(ctx, sc_cur, ape_col); + } + + if (batched_state) { + // state rows already written above (batched) + } else if (batched_rows) { + // Fused verify: batched state writes with ONE boundary allowed at ANY + // batch index b (q <= ratio keeps every pos_mod distinct). Graph order: + // writes[0..b] -> pool(boundary, reads through span A) -> rotate + // cur->prev (ratio-4) -> writes[b+1..]. The pooling and rotation code + // below read state_*_source, which span A set. + ggml_tensor * kv_all = ggml_mul_mat(ctx, kv_proj, cur_all); + ggml_tensor * sc_all = ggml_mul_mat(ctx, gate_proj, cur_all); + ggml_tensor * ape_cols = ggml_get_rows(ctx, ape, ape_row_inp); // [comp_width, q] + sc_all = ggml_add(ctx, sc_all, ape_cols); + for (int ti = 0; ti < n_tokens_all; ++ti) { + if (((kv_start_all + ti + 1) % ratio) == 0) { batched_b = ti; break; } + } + const int nA = (batched_b >= 0) ? (batched_b + 1) : n_tokens_all; + batched_nB = n_tokens_all - nA; + auto write_span = [&](int off, int count, ggml_tensor ** kv_src, ggml_tensor ** sc_src) { + if (count <= 0) return; + ggml_tensor * kv_v = ggml_cont(ctx, ggml_view_2d(ctx, kv_all, comp_width, count, + kv_all->nb[1], (size_t) off * kv_all->nb[1])); + ggml_tensor * sc_v = ggml_cont(ctx, ggml_view_2d(ctx, sc_all, comp_width, count, + sc_all->nb[1], (size_t) off * sc_all->nb[1])); + ggml_tensor * rows_v = ggml_view_1d(ctx, state_rows_inp, count, + (size_t) off * state_rows_inp->nb[0]); + *kv_src = ggml_set_rows(ctx, state.state_kv, kv_v, rows_v); + *sc_src = ggml_set_rows(ctx, state.state_score, sc_v, rows_v); + ggml_build_forward_expand(gf, *kv_src); + ggml_build_forward_expand(gf, *sc_src); + }; + write_span(0, nA, &state_kv_source, &state_score_source); + batched_kv_all = kv_all; + batched_sc_all = sc_all; + batched_span_off = nA; + } else if (state_rows_inp) { state_kv_source = ggml_set_rows(ctx, state.state_kv, kv_cur, state_rows_inp); state_score_source = ggml_set_rows(ctx, state.state_score, sc_cur, state_rows_inp); ggml_build_forward_expand(gf, state_kv_source); @@ -649,7 +729,14 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_cur, sc_slot)); } - if (((token_pos + 1) % ratio) != 0) { + if (batched_rows && batched_b < 0) { + // no boundary inside the batch: state rows written, nothing to emit + return; + } + if (!batched_rows && !flush_rows_inp && ((token_pos + 1) % ratio) != 0) { + // Legacy per-layer graphs only pool at flush boundaries. The fused + // stable-topology graph (flush_rows_inp set) pools every step; the + // partial result lands on the masked running comp row. return; } @@ -738,7 +825,53 @@ static void build_compressor_step( *comp_cache_source_out = comp_cache_source; } - if (ratio == 4) { + if (batched_rows) { + if (ratio == 4) { + // completed window: rotate current half -> prev half, reading + // through the span-A writes so ordering is explicit. + for (int r = 0; r < ratio; ++r) { + ggml_tensor * src_kv = ggml_view_2d(ctx, state_kv_source, comp_width, 1, + state_kv_source->nb[1], + (size_t)(ratio + r) * state_kv_source->nb[1]); + ggml_tensor * dst_kv = ggml_view_2d(ctx, state.state_kv, comp_width, 1, + state.state_kv->nb[1], + (size_t) r * state.state_kv->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx, src_kv, dst_kv)); + ggml_tensor * src_sc = ggml_view_2d(ctx, state_score_source, comp_width, 1, + state_score_source->nb[1], + (size_t)(ratio + r) * state_score_source->nb[1]); + ggml_tensor * dst_sc = ggml_view_2d(ctx, state.state_score, comp_width, 1, + state.state_score->nb[1], + (size_t) r * state.state_score->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx, src_sc, dst_sc)); + } + } + if (batched_nB > 0) { + ggml_tensor * kv_v = ggml_cont(ctx, ggml_view_2d(ctx, batched_kv_all, comp_width, batched_nB, + batched_kv_all->nb[1], (size_t) batched_span_off * batched_kv_all->nb[1])); + ggml_tensor * sc_v = ggml_cont(ctx, ggml_view_2d(ctx, batched_sc_all, comp_width, batched_nB, + batched_sc_all->nb[1], (size_t) batched_span_off * batched_sc_all->nb[1])); + ggml_tensor * rows_v = ggml_view_1d(ctx, state_rows_inp, batched_nB, + (size_t) batched_span_off * state_rows_inp->nb[0]); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_kv, kv_v, rows_v)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_score, sc_v, rows_v)); + } + return; + } + if (ratio == 4 && flush_rows_inp) { + // Stable-topology flush: copy the cur half onto rows given by the + // input (prev half [0..3] at flush, cur half itself [4..7] = no-op + // otherwise). Values are read through the set_rows source so this + // step's state write is ordered first. + ggml_tensor * cur_kv_vals = ggml_cont(ctx, ggml_view_2d( + ctx, state_kv_source, comp_width, ratio, state_kv_source->nb[1], + (size_t) ratio * state_kv_source->nb[1])); + ggml_tensor * cur_sc_vals = ggml_cont(ctx, ggml_view_2d( + ctx, state_score_source, comp_width, ratio, state_score_source->nb[1], + (size_t) ratio * state_score_source->nb[1])); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_kv, cur_kv_vals, flush_rows_inp)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_score, cur_sc_vals, flush_rows_inp)); + } else if (ratio == 4) { for (int r = 0; r < ratio; ++r) { ggml_tensor * src_kv = ggml_view_2d(ctx, state.state_kv, comp_width, 1, state.state_kv->nb[1], @@ -779,7 +912,11 @@ static void build_indexer_compressor_step( ggml_tensor * comp_rows_inp, ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, - std::vector & i32_array_inputs) { + std::vector & i32_array_inputs, + ggml_tensor * flush_rows_inp = nullptr, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, @@ -803,7 +940,11 @@ static void build_indexer_compressor_step( comp_pos_inp, i64_array_inputs, i32_array_inputs, - nullptr); + nullptr, + flush_rows_inp, + cur_all, + n_tokens_all, + kv_start_all); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { @@ -925,7 +1066,8 @@ static ggml_tensor * build_mla_attention( const DeepSeek4AttentionGraphInputs * cached_inputs, std::vector & i32_inputs, std::vector & i32_array_inputs, - std::vector & i64_array_inputs) { + std::vector & i64_array_inputs, + std::vector * f32_array_inputs = nullptr) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; @@ -985,6 +1127,51 @@ static ggml_tensor * build_mla_attention( rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + // ── Causal batched step (exact multi-token target semantics) ─── + // The target model is causal: token i must not attend to batch tokens + // j > i, must see the compressed-row count as of its own position, and — + // once the ring has wrapped — must still see the OLD contents of ring + // slots that later batch tokens overwrite. Default ON for multi-token + // steps on this path; DFLASH_DS4_NO_CAUSAL_VERIFY=1 restores the legacy + // (bidirectional) behavior for A/B comparison. + const bool causal_batch = (n_tokens > 1) && !cached_inputs && f32_array_inputs && + !ds4_env_flag("DFLASH_DS4_NO_CAUSAL_VERIFY"); + ggml_tensor * old_rows_scratch = nullptr; + int n_old_rows = 0; + const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; + if (fused_causal) { + // Fused verify: ALWAYS q preserved rows so the topology is stable; + // unwrapped/garbage rows are masked by the host-filled mask values. + for (int ti = 0; ti < n_tokens; ti++) { + ggml_tensor * slot = ggml_view_2d( + ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], + (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + ggml_tensor * saved = ggml_cont(ctx, slot); + ggml_build_forward_expand(gf, saved); + old_rows_scratch = old_rows_scratch + ? ggml_concat(ctx, old_rows_scratch, saved, 1) : saved; + n_old_rows++; + } + old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); + } else if (causal_batch) { + // Copy the to-be-overwritten rows FIRST; same-stream build order runs + // these before the ring writes below. + for (int ti = 0; ti < n_tokens; ti++) { + if (kv_start + ti < w.n_swa) continue; // slot never held an older pos + ggml_tensor * slot = ggml_view_2d( + ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], + (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + ggml_tensor * saved = ggml_cont(ctx, slot); + ggml_build_forward_expand(gf, saved); + old_rows_scratch = old_rows_scratch + ? ggml_concat(ctx, old_rows_scratch, saved, 1) : saved; + n_old_rows++; + } + if (old_rows_scratch) { + old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); + } + } + // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. ggml_tensor * raw_kv_source = lc.raw_kv; @@ -1035,7 +1222,11 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->attn_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, - &comp_kv_source); + &comp_kv_source, + cached_inputs ? cached_inputs->flush_rows : nullptr, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start); } if (ratio == 4 && L.indexer_compressor_kv) { @@ -1045,7 +1236,11 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->index_comp_rows : nullptr, cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + cached_inputs ? cached_inputs->flush_rows : nullptr, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start); (void)build_indexer_score(ctx, qr_last, cur_last, w, L, lc, token_pos, i32_inputs); } @@ -1054,14 +1249,14 @@ static ggml_tensor * build_mla_attention( // raw_kv: [head_dim, n_swa] F16 persistent ring buffer (single KV head, shared) // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) - const bool masked_kv = (n_tokens == 1) && cached_inputs && cached_inputs->attn_row_mask; + const bool masked_kv = cached_inputs && cached_inputs->attn_row_mask; const int n_comp_live = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; // Stable path reads the full physical ring (masking not-yet-written slots) // and a padded compressed-row span; the plain path reads only valid rows. const int n_raw = masked_kv ? w.n_swa : std::min(kv_start + n_tokens, w.n_swa); const int n_comp_attn = masked_kv ? cached_inputs->padded_comp : n_comp_live; const int n_valid_raw = std::min(kv_start + n_tokens, w.n_swa); - const int n_attn = n_raw + n_comp_attn; + const int n_attn = n_raw + n_comp_attn + n_old_rows; const float kq_scale = 1.0f / sqrtf((float)head_dim); // Get valid KV rows. For single-token decode, include the current in-graph @@ -1107,6 +1302,9 @@ static ggml_tensor * build_mla_attention( comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } + if (old_rows_scratch) { + kv_attn = ggml_concat(ctx, kv_attn, old_rows_scratch, 1); + } // kv_attn: [head_dim, n_attn] // Flatten q to [head_dim, n_head*n_tokens] for batched matmul @@ -1116,9 +1314,47 @@ static ggml_tensor * build_mla_attention( // → [n_attn, n_head*n_tokens] ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); scores = ggml_scale(ctx, scores, kq_scale); - if (masked_kv) { + if (masked_kv && n_tokens > 1) { + // Per-token causal mask [n_attn, n_tokens] from the host-filled bundle. + ggml_tensor * m3 = ggml_reshape_3d(ctx, cached_inputs->attn_row_mask, n_attn, 1, n_tokens); + ggml_tensor * scores3d = ggml_reshape_3d(ctx, scores, n_attn, n_head, n_tokens); + scores3d = ggml_add(ctx, scores3d, m3); + scores = ggml_reshape_2d(ctx, scores3d, n_attn, n_head * n_tokens); + } else if (masked_kv) { // Broadcast-add the [n_attn,1] additive mask across all query columns. scores = ggml_add(ctx, scores, cached_inputs->attn_row_mask); + } else if (causal_batch) { + // Per-token causal mask over [ring rows | comp rows | old rows]. + ggml_tensor * cmask = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); + ggml_set_input(cmask); + std::vector mvals((size_t) n_attn * n_tokens, 0.0f); + const int e = kv_start + n_tokens; // exclusive end position + for (int i = 0; i < n_tokens; i++) { + const int pos_i = kv_start + i; + float * col = mvals.data() + (size_t) i * n_attn; + for (int r = 0; r < n_raw; r++) { + // position held by ring slot r AFTER this batch's writes + const int p_r = (e <= w.n_swa) ? r + : (e - 1) - ((e - 1 - r) % w.n_swa); + if (p_r > pos_i) col[r] = -1e30f; + } + if (n_comp_attn > 0) { + const int vis = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i); + for (int c = vis; c < n_comp_attn; c++) col[n_raw + c] = -1e30f; + } + // old contents of slot overwritten by batch token j are visible + // exactly to tokens i < j (still inside their SWA window) + int oi = 0; + for (int tj = 0; tj < n_tokens; tj++) { + if (kv_start + tj < w.n_swa) continue; + if (tj <= i) col[n_raw + n_comp_attn + oi] = -1e30f; + oi++; + } + } + f32_array_inputs->push_back({cmask, std::move(mvals)}); + ggml_tensor * scores3d = ggml_reshape_3d(ctx, scores, n_attn, n_head, n_tokens); + scores3d = ggml_add(ctx, scores3d, cmask); + scores = ggml_reshape_2d(ctx, scores3d, n_attn, n_head * n_tokens); } (void) n_valid_raw; @@ -2355,6 +2591,17 @@ static void cpu_matvec_f16_pooled(float * out, const uint16_t * mat, const float ds4_hc_matvec_pool().run(mat, x, out, rows, cols); } +// Token-level persistent-pool parallel-for: same splitting semantics as +// ds4_parallel_for_tokens but without per-call thread spawns (a multi-token +// step issues ~86 batched-HC calls, so spawn cost dominates at small n). +// Inner work must stay serial (serial_fn=true paths) - the pool is not +// reentrant. +static void ds4_pool_for_tokens(int n_tokens, const std::function & fn) { + if (n_tokens <= 1) { fn(0, n_tokens); return; } + static Ds4HcMatvecPool token_pool; + token_pool.run_custom(n_tokens, [&fn](int t) { fn(t, t + 1); }); +} + static void cpu_hc_sinkhorn(float * out, const float * mix, const float * scale, const float * base, int n_hc, int iters, float eps) { const float pre_scale = scale[0]; @@ -2589,7 +2836,7 @@ static void hc_pre_batch(std::vector & working, post.resize((size_t)n_tokens * (size_t)n_hc); comb.resize((size_t)n_tokens * (size_t)n_hc * (size_t)n_hc); - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { std::vector flat(hc_dim); float mix[24]; for (int t = t0; t < t1; ++t) { @@ -2651,7 +2898,7 @@ static void hc_post_batch(std::vector & out_hc, }); return; } - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { for (int t = t0; t < t1; ++t) { cpu_hc_post(out_hc.data() + (size_t)t * hc_dim, block_out + (size_t)t * n_embd, @@ -2673,7 +2920,7 @@ static void hc_output_batch(std::vector & final_embd, float hc_eps) { const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; final_embd.resize((size_t)n_tokens * (size_t)n_embd); - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { std::vector flat(hc_dim); std::vector pre((size_t)n_hc); std::vector hc_weights((size_t)n_hc); @@ -2801,7 +3048,8 @@ static bool deepseek4_step_hybrid( const int32_t * token_ids, MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, - MoeHybridRoutingStats * routing_stats) { + MoeHybridRoutingStats * routing_stats, + Ds4VerifyHooks * verify_hooks) { const auto step_t0 = Ds4TimingClock::now(); const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -3285,7 +3533,8 @@ bool deepseek4_step( std::vector hc_state; return deepseek4_step_layer_range( backend, w, cache, hc_state, embed, n_tokens, kv_start, - 0, w.n_layer, &out_logits, token_ids, telemetry); + 0, w.n_layer, &out_logits, token_ids, telemetry, + /*allow_decode_graph_reuse=*/verify_hooks == nullptr, verify_hooks); } // ─── Fused single-graph decode (n_tokens == 1) ────────────────────────── @@ -3681,6 +3930,8 @@ static bool ds4_build_fused_decode_graph( return true; } +#include "deepseek4_fused_verify.inc" + // Returns 1 on success (out_logits filled), 0 to fall back to the per-layer // path, -1 on a hard failure after cache state may have been touched. static int ds4_try_fused_decode_step( @@ -3887,7 +4138,9 @@ bool deepseek4_step_layer_range( int layer_end, std::vector * out_logits, const int32_t * token_ids, - DeepSeek4StepTelemetry * telemetry) { + DeepSeek4StepTelemetry * telemetry, + bool allow_decode_graph_reuse, + Ds4VerifyHooks * verify_hooks) { const auto step_t0 = Ds4TimingClock::now(); // ── Partial layer-range forward with HC ───────────────────────────── @@ -3982,7 +4235,30 @@ bool deepseek4_step_layer_range( const int n_expert_used = ds4_effective_expert_count(w); scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, n_expert_used); - if (n_tokens == 1 && layer_begin == 0 && is_last_shard && + if (n_tokens >= 2 && n_tokens <= 4 && verify_hooks && layer_begin == 0 && is_last_shard && + out_logits && ds4_backend_is_gpu(backend) && w.fused_decode && + ds4_fused_verify_enabled()) { + const int vrc = ds4_try_fused_verify_step( + fused_verify_graph_cache, fused_decode_graph_cache, backend, w, cache, + hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, + scratch.hash_expert_ids, embed, n_tokens, kv_start, *out_logits, token_ids, + verify_hooks, telemetry); + if (vrc < 0) return false; + if (vrc > 0) { + const int np = kv_start + n_tokens; + for (int il = layer_begin; il < layer_end; ++il) { + const uint32_t vratio = w.compress_ratios[il]; + if (vratio <= 0) continue; + cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, np / (int) vratio); + if (vratio == 4) cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, np / (int) vratio); + } + cache.cur_pos = np; + if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + return true; + } + } + std::vector fused_debug_logits; + if (n_tokens == 1 && allow_decode_graph_reuse && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend) && w.fused_decode) { const int rc = ds4_try_fused_decode_step( fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, @@ -4240,6 +4516,7 @@ bool deepseek4_step_layer_range( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f32_array_inputs; const size_t graph_size = ds4_attn_step_graph_size(n_tokens); gf = ggml_new_graph_custom(ctx, graph_size, false); @@ -4247,7 +4524,7 @@ bool deepseek4_step_layer_range( attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs); + i64_array_inputs, &f32_array_inputs); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -4275,6 +4552,8 @@ bool deepseek4_step_layer_range( ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(int32_t) * b.values.size()); for (const auto & b : i64_array_inputs) ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(int64_t) * b.values.size()); + for (const auto & b : f32_array_inputs) + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(float) * b.values.size()); } const auto attn_compute_t0 = Ds4TimingClock::now(); @@ -4490,6 +4769,25 @@ bool deepseek4_step_layer_range( n_hc); std::memcpy(hc_state.data(), next_hc.data(), next_hc.size() * sizeof(float)); if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); + if (verify_hooks && verify_hooks->capture_layer_ids && verify_hooks->capture_out) { + const std::vector & _ids = *verify_hooks->capture_layer_ids; + for (size_t _ci = 0; _ci < _ids.size(); ++_ci) { + if (_ids[_ci] != il) continue; + const int _ncap = (int) _ids.size(); + std::vector & _cap = *verify_hooks->capture_out; + if ((int) _cap.size() != _ncap * n_embd * n_tokens) + _cap.assign((size_t) _ncap * n_embd * n_tokens, 0.0f); + for (int _t = 0; _t < n_tokens; ++_t) { + float * _dst = _cap.data() + (size_t) _t * _ncap * n_embd + (size_t) _ci * n_embd; + const float * _hs = hc_state.data() + (size_t) _t * hc_dim; + for (int _d = 0; _d < n_embd; ++_d) { + float _acc = 0.0f; + for (int _h = 0; _h < n_hc; ++_h) _acc += _hs[(size_t) _h * n_embd + _d]; + _dst[_d] = _acc / (float) n_hc; + } + } + } + } } } } @@ -4568,6 +4866,11 @@ bool deepseek4_step_layer_range( const size_t logits_offset = (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); ggml_backend_tensor_get(logits, out_logits->data(), logits_offset, sizeof(float) * (size_t)w.n_vocab); + if (verify_hooks && verify_hooks->all_logits_out) { + verify_hooks->all_logits_out->resize((size_t) w.n_vocab * n_tokens); + ggml_backend_tensor_get(logits, verify_hooks->all_logits_out->data(), 0, + sizeof(float) * (size_t) w.n_vocab * n_tokens); + } ggml_free(ctx); } if (telemetry) telemetry->output_us += ds4_elapsed_us(output_t0, Ds4TimingClock::now()); @@ -4890,3 +5193,385 @@ void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { } } // namespace dflash::common + +// ══════════════════════════════════════════════════════════════════════ +// DSpark drafter forward graph (appended to deepseek4_graph.cpp so it can +// reuse the file-static DS4 sub-builders: build_rms_norm, build_tail_rope_*, +// build_moe_ffn, build_shared_ffn). See deepseek4_dspark.h for the contract. +// +// Mirrors deepseek-ai/DeepSeek-V4-Flash-DSpark inference/model.py: +// forward_embed -> main_x = main_norm(main_proj(cat[h40,h41,h42])) +// per layer (DSparkBlock): HC-pre (per block position) -> attn_norm -> +// DSparkAttention (bidirectional over [ctx main-KV ++ block-KV]) -> +// HC-post ; HC-pre -> ffn_norm -> MoE -> HC-post +// tail: hc_head collapse -> out_norm (input to the tied lm_head + Markov) +// +// The ggml_ds4_hc_* ops are single-token, so HC-pre/HC-post run per block +// position; attention batches all block positions together (bidirectional). +// ══════════════════════════════════════════════════════════════════════ + +#include "deepseek4_dspark.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Fresh MLA attention for the drafter: no KV cache, no compression. The 5 +// block queries attend over an explicit [ctx main-context KV ++ block KV] +// tensor with full (bidirectional) visibility, plus the learned per-head sink. +static ggml_tensor * build_dspark_attention( + ggml_context * ctx, + ggml_tensor * cur, // [n_embd, block] (post attn_norm) + ggml_tensor * main_x, // [n_embd, ctx_len] (post main_norm, shared) + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int ctx_len, + ggml_tensor * pos_block, // I32[block] absolute positions committed..committed+block-1 + ggml_tensor * neg_block, // I32[block] -(block positions) + ggml_tensor * pos_ctx) { // I32[ctx_len] absolute positions committed-ctx_len..committed-1 + const int n_embd = w.n_embd; + const int head_dim = w.head_dim; + const int n_head = w.n_head; + const int n_rot = w.n_rot; + const int n_lora_o = w.n_lora_o; + const int n_out_group = w.n_out_group; + const int block = (int) cur->ne[1]; + const float eps = w.rms_eps; + // DSparkAttention has compress_ratio==0 -> base RoPE, YaRN disabled. + const float rope_freq = w.rope_freq_base; + const float rope_scale = 1.0f, rope_ext = 0.0f, rope_attn = 1.0f; + const int rope_orig = (int) w.rope_orig_ctx; + + // ── Q path (block queries) ────────────────────────────────────── + ggml_tensor * qr = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_q_a, cur), L.attn_q_a_norm, eps); + ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); // [n_head*head_dim, block] + q = ggml_reshape_3d(ctx, q, head_dim, n_head, block); + q = ggml_rms_norm(ctx, q, eps); // per-head unweighted + q = build_tail_rope_3d(ctx, q, pos_block, n_rot, head_dim, n_head, block, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + + // ── KV: block positions ───────────────────────────────────────── + ggml_tensor * kv_b = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_kv, cur), L.attn_kv_a_norm, eps); + kv_b = build_tail_rope_2d(ctx, kv_b, pos_block, n_rot, head_dim, block, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + + // ── KV: context positions (from shared main_x) ────────────────── + ggml_tensor * kv_attn = kv_b; + int n_attn = block; + if (ctx_len > 0) { + ggml_tensor * kv_c = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_kv, main_x), L.attn_kv_a_norm, eps); + kv_c = build_tail_rope_2d(ctx, kv_c, pos_ctx, n_rot, head_dim, ctx_len, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + kv_attn = ggml_concat(ctx, kv_c, kv_b, 1); // [head_dim, ctx_len+block] + n_attn = ctx_len + block; + } + + // ── Scores + sink softmax (full visibility, no causal mask) ───── + ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, n_head * block); + ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); // [n_attn, n_head*block] + scores = ggml_scale(ctx, scores, 1.0f / sqrtf((float) head_dim)); + ggml_tensor * probs = nullptr; + if (L.attn_sinks) { + ggml_tensor * sink = ggml_reshape_2d(ctx, L.attn_sinks, 1, n_head); + ggml_tensor * sink_shape = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n_head * block); + sink = ggml_repeat(ctx, sink, sink_shape); + ggml_tensor * sws = ggml_concat(ctx, scores, sink, 0); // [n_attn+1, n_head*block] + ggml_tensor * pws = ggml_soft_max(ctx, sws); + probs = ggml_view_2d(ctx, pws, n_attn, n_head * block, pws->nb[1], 0); + } else { + probs = ggml_soft_max(ctx, scores); + } + + // ── Context, inverse RoPE, grouped low-rank output ────────────── + ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); // [n_attn, head_dim] + ggml_tensor * context = ggml_mul_mat(ctx, kv_T, probs); // [head_dim, n_head*block] + context = ggml_reshape_3d(ctx, context, head_dim, n_head, block); + context = build_tail_rope_3d(ctx, context, neg_block, n_rot, head_dim, n_head, block, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, block); + const int group_dim = head_dim * (n_head / n_out_group); + attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, block); + attn_out = ggml_cont(ctx, ggml_permute(ctx, attn_out, 0, 2, 1, 3)); // [group_dim, block, n_out_group] + ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); + ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); // [n_lora_o, block, n_out_group] + attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); // [n_lora_o, n_out_group, block] + attn_low = ggml_reshape_2d(ctx, attn_low, n_lora_o * n_out_group, block); + return ggml_mul_mat(ctx, L.attn_output_b, attn_low); // [n_embd, block] +} + +// Read a small F32 GPU tensor (HC scale, [k]) into host floats. +static void ds4_read_f32(ggml_tensor * t, float * dst, int k) { + if (t) ggml_backend_tensor_get(t, dst, 0, sizeof(float) * (size_t) k); + else for (int i = 0; i < k; i++) dst[i] = 0.0f; +} + +} // namespace + +// ── Cached drafter graph ──────────────────────────────────────────────── +// The drafter forward runs every spec step with identical topology (ctx_len +// is constant once the feature window fills at n_swa). Rebuilding the +// multi-thousand-node graph, zero-initializing a fresh 256 MB arena and +// re-planning gallocr each call used to cost more than the 3-layer compute +// itself (~63 ms/step). Cache the built graph keyed by (ctx_len, block, +// drafter instance) and re-set only the inputs per call. +namespace { + +struct DsparkDraftCache { + int ctx_len = -1; + int block = -1; + const void * drafter = nullptr; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_gallocr_t alloc = nullptr; + ggml_cgraph * gf = nullptr; + ggml_tensor * inp_noise = nullptr; + ggml_tensor * inp_ctx = nullptr; + ggml_tensor * pos_block = nullptr; + ggml_tensor * neg_block = nullptr; + ggml_tensor * pos_ctx = nullptr; + ggml_tensor * out = nullptr; + std::vector> dbg_taps; + // HC scales are immutable weights: read from the backend once. + std::vector> s_attn, s_ffn; + float s_out = 0.0f; +}; + +thread_local DsparkDraftCache g_dspark_draft_cache; + +} // namespace + +bool deepseek4_dspark_draft_forward(ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + const float * ctx_features, + int ctx_len, + int committed, + std::vector & out_hidden) { + const DeepSeek4Weights & w = d.core; + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int block = d.block_size; + const int fc_in = d.n_target_layers * n_embd; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + const float hc_eps = w.hc_eps; + if (ctx_len < 0) ctx_len = 0; + + DsparkDraftCache & C = g_dspark_draft_cache; + const bool DS4_DBG = std::getenv("DFLASH_DS4_DSPARK_DEBUG") != nullptr; + + if (C.drafter != (const void *) &d) { + // HC scales (host) per layer + output — immutable, read once per drafter. + C.s_attn.assign((size_t) w.n_layer, {0.0f, 0.0f, 0.0f}); + C.s_ffn.assign((size_t) w.n_layer, {0.0f, 0.0f, 0.0f}); + for (int il = 0; il < w.n_layer; il++) { + ds4_read_f32(w.layers[il].hc_attn_scale, C.s_attn[il].data(), 3); + ds4_read_f32(w.layers[il].hc_ffn_scale, C.s_ffn[il].data(), 3); + } + float so[1] = {0.0f}; + ds4_read_f32(w.output_hc_scale, so, 1); + C.s_out = so[0]; + } + + if (!C.ctx || C.ctx_len != ctx_len || C.block != block || C.drafter != (const void *) &d) { + // ── (Re)build the graph ───────────────────────────────────────── + if (C.ctx) { ggml_free(C.ctx); C.ctx = nullptr; } + C.gf = nullptr; + C.dbg_taps.clear(); + if (C.arena.empty()) C.arena.resize(256u * 1024 * 1024); + ggml_init_params ip{}; + ip.mem_size = C.arena.size(); + ip.mem_buffer = C.arena.data(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + C.ctx = ctx; + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); + C.gf = gf; + auto dbg_tap = [&](const std::string & nm, ggml_tensor * t) { + if (!DS4_DBG || !t) return; + ggml_set_output(t); + ggml_build_forward_expand(gf, t); + C.dbg_taps.push_back({nm, t}); + }; + + // Inputs. + C.inp_noise = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, block); + ggml_set_input(C.inp_noise); + C.inp_ctx = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, fc_in, ctx_len > 0 ? ctx_len : 1); + ggml_set_input(C.inp_ctx); + C.pos_block = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, block); + ggml_set_input(C.pos_block); + C.neg_block = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, block); + ggml_set_input(C.neg_block); + C.pos_ctx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ctx_len > 0 ? ctx_len : 1); + ggml_set_input(C.pos_ctx); + + // main_x = main_norm(main_proj(ctx_features)). Shared across layers. + ggml_tensor * main_x = nullptr; + if (ctx_len > 0) { + // Captured target features have large magnitude (rms ~1e3 — HC streams + // accumulate over 40+ layers). main_proj is rocmfp4-quantized and its + // activation quantization overflows on inputs that big -> NaN. Since + // main_norm (RMSNorm) normalizes main_proj's output and RMSNorm is + // scale-invariant, pre-normalizing the features to unit RMS gives a + // mathematically identical main_x while keeping the rocmfp4 activation + // in a safe range: main_norm(main_proj(f/rms(f))) == main_norm(main_proj(f)). + ggml_tensor * fc_in_normed = ggml_rms_norm(ctx, C.inp_ctx, w.rms_eps); + ggml_tensor * fc_out = ggml_mul_mat(ctx, d.main_proj, fc_in_normed); // [n_embd, ctx_len] + dbg_tap("fc_out", fc_out); + main_x = build_rms_norm(ctx, fc_out, d.main_norm, w.rms_eps); + dbg_tap("main_x", main_x); + } + + // HC state: [n_embd, n_hc, block], init = block embeds replicated over streams. + ggml_tensor * noise3 = ggml_reshape_3d(ctx, C.inp_noise, n_embd, 1, block); + ggml_tensor * hc_cur = ggml_repeat_4d(ctx, noise3, n_embd, n_hc, block, 1); + + auto hc_col = [&](ggml_tensor * hc, int p) -> ggml_tensor * { + // Contiguous [n_embd*n_hc] slab for block position p. + return ggml_view_1d(ctx, hc, (int64_t) n_embd * n_hc, + (size_t) p * hc->nb[2]); + }; + + for (int il = 0; il < w.n_layer; il++) { + const DeepSeek4Layer & L = w.layers[il]; + + // ── HC pre (attention), per block position ────────────────── + std::vector split_attn(block), work_cols(block); + for (int p = 0; p < block; p++) { + ggml_tensor * hcf = hc_col(hc_cur, p); + ggml_tensor * normed = ggml_rms_norm(ctx, hcf, hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, L.hc_attn_fn, normed); + mix = ggml_reshape_1d(ctx, mix, mix_dim); + ggml_tensor * base = ggml_reshape_1d(ctx, L.hc_attn_base, mix_dim); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hcf, n_hc, + w.n_hc_sinkhorn_iter, + C.s_attn[il][0], C.s_attn[il][1], C.s_attn[il][2]); + work_cols[p] = ggml_reshape_2d(ctx, ggml_view_1d(ctx, pre, n_embd, 0), n_embd, 1); + split_attn[p] = ggml_view_1d(ctx, pre, mix_dim, (size_t) n_embd * sizeof(float)); + } + ggml_tensor * attn_in = work_cols[0]; + for (int p = 1; p < block; p++) attn_in = ggml_concat(ctx, attn_in, work_cols[p], 1); + ggml_tensor * attn_normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); + ggml_tensor * attn_out = build_dspark_attention(ctx, attn_normed, main_x, w, L, + ctx_len, C.pos_block, C.neg_block, C.pos_ctx); + dbg_tap(std::string("attn_L") + std::to_string(il), attn_out); + // ── HC post (attention), per block position ───────────────── + ggml_tensor * hc_next = nullptr; + for (int p = 0; p < block; p++) { + ggml_tensor * bo = ggml_view_1d(ctx, attn_out, n_embd, (size_t) p * attn_out->nb[1]); + ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_col(hc_cur, p), bo, split_attn[p], n_hc); + hp = ggml_reshape_3d(ctx, hp, n_embd, n_hc, 1); + hc_next = hc_next ? ggml_concat(ctx, hc_next, hp, 2) : hp; + } + hc_cur = ggml_cont(ctx, hc_next); + + // ── HC pre (FFN), per block position ──────────────────────── + std::vector split_ffn(block), fwork(block); + for (int p = 0; p < block; p++) { + ggml_tensor * hcf = hc_col(hc_cur, p); + ggml_tensor * normed = ggml_rms_norm(ctx, hcf, hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, L.hc_ffn_fn, normed); + mix = ggml_reshape_1d(ctx, mix, mix_dim); + ggml_tensor * base = ggml_reshape_1d(ctx, L.hc_ffn_base, mix_dim); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hcf, n_hc, + w.n_hc_sinkhorn_iter, + C.s_ffn[il][0], C.s_ffn[il][1], C.s_ffn[il][2]); + fwork[p] = ggml_reshape_2d(ctx, ggml_view_1d(ctx, pre, n_embd, 0), n_embd, 1); + split_ffn[p] = ggml_view_1d(ctx, pre, mix_dim, (size_t) n_embd * sizeof(float)); + } + ggml_tensor * ffn_in = fwork[0]; + for (int p = 1; p < block; p++) ffn_in = ggml_concat(ctx, ffn_in, fwork[p], 1); + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, block); + if (!ffn_out) { ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; return false; } + dbg_tap(std::string("ffn_L") + std::to_string(il), ffn_out); + // ── HC post (FFN) ─────────────────────────────────────────── + hc_next = nullptr; + for (int p = 0; p < block; p++) { + ggml_tensor * bo = ggml_view_1d(ctx, ffn_out, n_embd, (size_t) p * ffn_out->nb[1]); + ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_col(hc_cur, p), bo, split_ffn[p], n_hc); + hp = ggml_reshape_3d(ctx, hp, n_embd, n_hc, 1); + hc_next = hc_next ? ggml_concat(ctx, hc_next, hp, 2) : hp; + } + hc_cur = ggml_cont(ctx, hc_next); + dbg_tap(std::string("hcL") + std::to_string(il), hc_cur); + } + + // ── Tail: hc_head collapse -> out_norm, per block position ────── + ggml_tensor * out = nullptr; + for (int p = 0; p < block; p++) { + ggml_tensor * hcf = hc_col(hc_cur, p); + ggml_tensor * onorm = ggml_rms_norm(ctx, hcf, hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, w.output_hc_fn, onorm); + omix = ggml_reshape_1d(ctx, omix, n_hc); + ggml_tensor * obase = ggml_reshape_1d(ctx, w.output_hc_base, n_hc); + ggml_tensor * final_embd = ggml_ds4_hc_out(ctx, omix, obase, hcf, n_hc, C.s_out); + ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, 1); + ggml_tensor * hidden_p = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); + out = out ? ggml_concat(ctx, out, hidden_p, 1) : hidden_p; + } + ggml_set_output(out); + ggml_build_forward_expand(gf, out); + C.out = out; + + if (!C.alloc) C.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!C.alloc || !ggml_gallocr_alloc_graph(C.alloc, gf)) { + ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; + return false; + } + C.ctx_len = ctx_len; + C.block = block; + C.drafter = (const void *) &d; + } + + // ── Set inputs + compute (cached graph) ───────────────────────────── + ggml_backend_tensor_set(C.inp_noise, noise_embed, 0, sizeof(float) * (size_t) n_embd * block); + if (ctx_len > 0) { + ggml_backend_tensor_set(C.inp_ctx, ctx_features, 0, sizeof(float) * (size_t) fc_in * ctx_len); + std::vector pc(ctx_len); + for (int i = 0; i < ctx_len; i++) pc[i] = committed - ctx_len + i; + ggml_backend_tensor_set(C.pos_ctx, pc.data(), 0, sizeof(int32_t) * ctx_len); + } + std::vector pb(block), nb(block); + for (int i = 0; i < block; i++) { pb[i] = committed + i; nb[i] = -(committed + i); } + ggml_backend_tensor_set(C.pos_block, pb.data(), 0, sizeof(int32_t) * block); + ggml_backend_tensor_set(C.neg_block, nb.data(), 0, sizeof(int32_t) * block); + + const ggml_status st = ggml_backend_graph_compute(backend, C.gf); + if (st != GGML_STATUS_SUCCESS) { + // Invalidate: a failed compute leaves no reusable state guarantees. + ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; C.ctx_len = -1; + return false; + } + out_hidden.resize((size_t) n_embd * block); + ggml_backend_tensor_get(C.out, out_hidden.data(), 0, sizeof(float) * out_hidden.size()); + + if (DS4_DBG) { + for (auto & tp : C.dbg_taps) { + const size_t ne = ggml_nelements(tp.second); + std::vector buf(ne); + ggml_backend_tensor_get(tp.second, buf.data(), 0, sizeof(float) * ne); + double ss = 0.0; size_t nnan = 0; float mn = 1e30f, mx = -1e30f; + for (float v : buf) { + if (!std::isfinite(v)) { nnan++; } + else { ss += (double) v * v; if (v < mn) mn = v; if (v > mx) mx = v; } + } + std::fprintf(stderr, "[ds4-dspark-dbg] %-10s ne=%zu nnan=%zu rms=%.4f min=%.3f max=%.3f\n", + tp.first.c_str(), ne, nnan, std::sqrt(ss / (double) ne), mn, mx); + } + } + + return true; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 27fb5a1f1..b52ca002b 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -325,6 +325,8 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, // embed: [n_embd, n_tokens] input embeddings (post-embedding lookup). // hc_state: [n_hc * n_embd] persistent HC residual (updated in-place). // Returns logits for last token. +struct Ds4VerifyHooks; + bool deepseek4_step( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -337,7 +339,18 @@ bool deepseek4_step( const int32_t * token_ids = nullptr, MoeHybridStreamEngine * stream_engine = nullptr, DeepSeek4StepTelemetry * telemetry = nullptr, - MoeHybridRoutingStats * routing_stats = nullptr); + MoeHybridRoutingStats * routing_stats = nullptr, + Ds4VerifyHooks * verify_hooks = nullptr); + +// Optional hooks for the DSpark spec-decode batched verify (deepseek4_dspark). +// When set on a multi-token deepseek4_step_layer_range call they add: per-layer +// mean-over-HC feature capture and full per-position logits. Null on the normal +// (23 tok/s) decode path so it is completely unaffected. +struct Ds4VerifyHooks { + const std::vector * capture_layer_ids = nullptr; // e.g. {40,41,42} + std::vector * capture_out = nullptr; // [n_cap*n_embd * n_tokens] + std::vector * all_logits_out = nullptr; // [n_vocab * n_tokens] +}; bool deepseek4_step_layer_range( ggml_backend_t backend, @@ -351,7 +364,9 @@ bool deepseek4_step_layer_range( int layer_end, std::vector * out_logits, const int32_t * token_ids = nullptr, - DeepSeek4StepTelemetry * telemetry = nullptr); + DeepSeek4StepTelemetry * telemetry = nullptr, + bool allow_decode_graph_reuse = true, + Ds4VerifyHooks * verify_hooks = nullptr); bool build_deepseek4_moe_hybrid_storage_from_file( const std::string & path, diff --git a/server/test/test_ds4_dspark_load.cpp b/server/test/test_ds4_dspark_load.cpp new file mode 100644 index 000000000..5340e4bb4 --- /dev/null +++ b/server/test/test_ds4_dspark_load.cpp @@ -0,0 +1,153 @@ +// Smoke test: load the DeepSeek-V4-Flash DSpark drafter GGUF and dump bindings. +// +// test_ds4_dspark_load +// +// Verifies the "deepseek4-dflash-draft" GGUF <-> DSparkDrafter loader contract: +// all block tensors + DSpark heads bind, dspark_enabled, capture ids present. + +#include "deepseek4/deepseek4_dspark.h" + +#include "ggml.h" +#include "ggml-backend.h" +#if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) +#include "ggml-cuda.h" +#define DS4_HAVE_GPU 1 +#endif + +#include +#include +#include +#include +#include + +namespace dflash::common { void deepseek4_dspark_dump(const DSparkDrafter &); } + +int main(int argc, char ** argv) { + using namespace dflash::common; + const char * path = argc > 1 ? argv[1] + : "/home/lucebox2/models/DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4.gguf"; + + ggml_backend_t backend = nullptr; +#ifdef DS4_HAVE_GPU + backend = ggml_backend_cuda_init(0); +#endif + if (!backend) { std::fprintf(stderr, "FAIL: no GPU backend\n"); return 2; } + + DSparkDrafter d; + if (!load_deepseek4_dspark_drafter(path, backend, d)) { + std::fprintf(stderr, "FAIL: load: %s\n", deepseek4_dspark_last_error()); + return 1; + } + deepseek4_dspark_dump(d); + + // Contract checks. + int rc = 0; + auto need = [&](bool cond, const char * what) { + if (!cond) { std::fprintf(stderr, "FAIL: missing %s\n", what); rc = 1; } + }; + need(d.core.n_layer == 3, "n_layer==3"); + need(d.main_proj && d.main_norm, "main_proj/main_norm (dflash.fc/hidden_norm)"); + need(d.markov_w1 && d.markov_w2, "markov heads"); + need(d.dspark_enabled, "dspark_enabled"); + need((int)d.capture_layer_ids.size() == d.n_target_layers, "capture_layer_ids count"); + need(d.core.output == nullptr && d.core.tok_embd == nullptr, "tied embed/lm_head (no output tensors)"); + for (int il = 0; il < d.core.n_layer; il++) { + const auto & L = d.core.layers[il]; + need(L.attn_q_a && L.attn_q_b && L.attn_kv && L.attn_output_a && L.attn_output_b, + "MLA weights"); + need(L.ffn_gate_exps && L.ffn_up_exps && L.ffn_down_exps, "routed experts"); + need(L.ffn_gate_shexp && L.ffn_up_shexp && L.ffn_down_shexp, "shared expert"); + need(L.hc_attn_fn && L.hc_ffn_fn, "HC weights"); + need(d.core.compress_ratios[il] == 0, "compress_ratio==0"); + need(L.attn_compressor_kv == nullptr, "no compressor tensors"); + } + // ── Weight sanity: dequantize main_proj directly (no matmul) ──── + if (rc == 0 && d.main_proj) { + ggml_init_params ip{}; ip.mem_size = 64u*1024*1024; ip.no_alloc = true; + ggml_context * c = ggml_init(ip); + ggml_cgraph * g = ggml_new_graph(c); + ggml_tensor * f32 = ggml_cast(c, d.main_proj, GGML_TYPE_F32); + ggml_set_output(f32); ggml_build_forward_expand(g, f32); + ggml_gallocr_t al = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (al && ggml_gallocr_alloc_graph(al, g) && ggml_backend_graph_compute(backend, g) == GGML_STATUS_SUCCESS) { + size_t ne = ggml_nelements(f32); + std::vector buf(ne); + ggml_backend_tensor_get(f32, buf.data(), 0, sizeof(float)*ne); + size_t nnan=0; double ss=0; float mn=1e30f, mx=-1e30f; + for (float v : buf) { if(!std::isfinite(v)) nnan++; else { ss+=(double)v*v; if(vmx)mx=v; } } + std::fprintf(stderr, "[weight-check] main_proj dequant: ne=%zu nnan=%zu rms=%.5f min=%.4f max=%.4f\n", + ne, nnan, std::sqrt(ss/(double)ne), mn, mx); + } else { std::fprintf(stderr, "[weight-check] main_proj cast failed (type %s not castable?)\n", ggml_type_name(d.main_proj->type)); } + if (al) ggml_gallocr_free(al); ggml_free(c); + } + + // ── Isolate: mul_mat(weight, ones) for main_proj and attn_q_a ─── + if (rc == 0) { + auto mm_check = [&](const char * nm, ggml_tensor * W) { + if (!W) return; + ggml_init_params ip{}; ip.mem_size = 64u*1024*1024; ip.no_alloc = true; + ggml_context * c = ggml_init(ip); + ggml_cgraph * g = ggml_new_graph(c); + ggml_tensor * x = ggml_new_tensor_2d(c, GGML_TYPE_F32, W->ne[0], 1); + ggml_set_input(x); + ggml_tensor * y = ggml_mul_mat(c, W, x); + ggml_set_output(y); ggml_build_forward_expand(g, y); + ggml_gallocr_t al = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (al && ggml_gallocr_alloc_graph(al, g)) { + std::vector ones((size_t)W->ne[0], 1.0f); + ggml_backend_tensor_set(x, ones.data(), 0, sizeof(float)*ones.size()); + if (ggml_backend_graph_compute(backend, g) == GGML_STATUS_SUCCESS) { + size_t ne = ggml_nelements(y); std::vector buf(ne); + ggml_backend_tensor_get(y, buf.data(), 0, sizeof(float)*ne); + size_t nn=0; for(float v:buf) if(!std::isfinite(v)) nn++; + std::fprintf(stderr, "[mm-check] %-12s mul_mat(W[%lld,%lld], ones): nnan=%zu/%zu first=%.3f\n", + nm,(long long)W->ne[0],(long long)W->ne[1],nn,ne,buf.empty()?0:buf[0]); + } + } + if (al) ggml_gallocr_free(al); ggml_free(c); + }; + mm_check("main_proj", d.main_proj); + mm_check("blk0.attn_q_a", d.core.layers[0].attn_q_a); + mm_check("blk0.ffn_gate_shexp", d.core.layers[0].ffn_gate_shexp); + } + + // ── Exercise the drafter forward with dummy inputs ────────────── + // Validates the whole graph runs on the GPU (HC ops, MoE, rocmfp + // matmuls, bidirectional attention, tail) and returns finite output. + if (rc == 0) { + const int n_embd = d.core.n_embd; + const int block = d.block_size; + const int fc_in = d.n_target_layers * n_embd; + const char * cle = std::getenv("DS4_CTX_LEN"); + const int ctx_len = cle ? atoi(cle) : 8; + std::vector noise((size_t) n_embd * block); + std::vector feats((size_t) fc_in * ctx_len); + for (size_t i = 0; i < noise.size(); i++) noise[i] = 0.06f * std::sin(0.31f * (float) i + 1.3f); + for (size_t i = 0; i < feats.size(); i++) feats[i] = 0.05f * std::cos(0.17f * (float) i + 0.4f); + std::vector hidden; + const int committed = 64; // arbitrary >= ctx_len + std::fprintf(stderr, "\n── drafter forward (ctx_len=%d block=%d) ──\n", ctx_len, block); + if (!deepseek4_dspark_draft_forward(backend, d, noise.data(), feats.data(), + ctx_len, committed, hidden)) { + std::fprintf(stderr, "FAIL: drafter forward returned false\n"); + rc = 1; + } else { + bool finite = true; double sumsq = 0.0; + for (float v : hidden) { if (!std::isfinite(v)) finite = false; sumsq += (double) v * v; } + std::fprintf(stderr, "forward out: %zu floats, finite=%d, rms=%.4f, first=[%.4f %.4f %.4f]\n", + hidden.size(), (int) finite, + std::sqrt(sumsq / (double) hidden.size()), + hidden.size() > 0 ? hidden[0] : 0.0f, + hidden.size() > 1 ? hidden[1] : 0.0f, + hidden.size() > 2 ? hidden[2] : 0.0f); + need(finite, "finite forward output"); + need(hidden.size() == (size_t) n_embd * block, "forward output size"); + } + } + + std::fprintf(stderr, rc == 0 ? "\nPASS: DSpark drafter load + forward OK\n" : "\nFAIL\n"); + + free_deepseek4_dspark_drafter(d); + ggml_backend_free(backend); + return rc; +} From edd1b89a57559216ec0e23b24af0eb724187c282 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:45:44 +0200 Subject: [PATCH 03/12] perf(deepseek4): use DSpark confidence for adaptive verify width Return confidence scores from the fused Markov graph in the existing token-id synchronization, then select q=2/q=3/q=4 from cumulative prefix confidence. Compatible DSpark artifacts enable the policy automatically; artifacts without a confidence head retain the existing EWMA controller. No new deployment configuration is introduced.\n\nValidated on gfx1151: GSM+Math 10/10 at 29.25 tok/s weighted, within 0.8% of fixed q=4, with adaptive behavior retained on low-acceptance prompts. --- server/docs/DS4.md | 40 ++++++++++++---- server/src/common/dspark_head.cpp | 46 +++++++++++++++++-- server/src/common/dspark_head.h | 8 ++-- .../src/deepseek4/deepseek4_dspark_spec.cpp | 42 +++++++++++++++-- 4 files changed, 117 insertions(+), 19 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index c1d8ea264..ee2b5da69 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -138,13 +138,13 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. DeepSeek4 no longer uses the old expert-split environment variables or expert-worker tuning knobs. Those retired knobs were removed from the codebase rather than left behind as unsupported debug switches. -## DSpark Aux Draft Heads +## DSpark Speculative Decode -The current DSpark runtime lives in the shared DFlash speculative stack and is -used by the Laguna backend when the draft GGUF carries `dflash.dspark.*` -tensors. DeepSeek4 DSpark work is still a draft bridge: the DeepSeek4/MTP -artifact stores compatible aux heads under the `mtp.2.*` namespace, and the -converter can now map those names into the existing GGUF contract. +DeepSeek4 uses the shared DFlash DSpark head implementation together with a +DeepSeek4-specific three-layer drafter and fused target verification. The draft +GGUF carries its auxiliary projections under the existing `dflash.dspark.*` +tensor contract. DeepSeek4/MTP checkpoints store compatible heads under the +`mtp.2.*` namespace, which the converter maps as follows. Supported DeepSeek4/MTP input tensors: @@ -169,10 +169,30 @@ python server/scripts/convert_dflash_to_gguf.py \ --aux-heads /path/to/hf-ds4-flash-dspark/model-00048-of-00048.safetensors ``` -This does not by itself make the production DeepSeek4 layer-split backend use -DSpark. It makes the DeepSeek4 DSpark aux artifact consumable by the existing -`dflash.dspark.*` GGUF loader/runtime so the follow-up runtime PR can wire it -into the DS4 decode path with a clean tensor contract. +Run the converted drafter against a DeepSeek4 target with: + +```bash +export DFLASH_DS4_SPEC=1 +export DFLASH_DS4_FUSED_VERIFY=1 +export DFLASH_DS4_DRAFT=/path/to/dflash-draft.gguf +export DFLASH_DS4_SPEC_Q=4 + +./server/build-hip/dflash_server /path/to/deepseek4-target.gguf +``` + +Adaptive width is automatic. When the draft artifact has a compatible +confidence projection, the runtime selects q=2, q=3, or q=4 from the cumulative +confidence of the proposed prefix. It adds the projection to the same fused +Markov graph and reads its scores in the existing token-id synchronization; no +additional host round trip is introduced. Artifacts without a compatible +confidence head transparently retain the existing acceptance-EWMA policy. + +On the gfx1151 validation host, confidence-adaptive width retained 10/10 +GSM+Math accuracy and measured 29.25 tok/s weighted, within 0.8% of fixed q=4 +at 29.49 tok/s. On the low-acceptance stress prompt it measured 21.9/21.8 +tok/s warm, effectively tied with EWMA while avoiding fixed q=4's wasted wide +verification. These numbers are workload-specific; the confidence policy is +opt-in. ## Example: CUDA + Halo Layer Split diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index 4aea171b7..2e1cca95d 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -182,6 +182,7 @@ struct MarkovChainGraph { ggml_tensor * base = nullptr; // [vocab, n_positions] std::vector toks; // corrected argmax per depth std::vector corrected; // corrected logits per depth + std::vector confidence; // optional sigmoid score per depth }; // Guards shared by the fused Markov paths: head present, usable inputs, and @@ -218,6 +219,7 @@ bool build_markov_chain_graph(const DraftWeights & dw, ggml_tensor * lm_head, int n_positions, int first_corrected, bool corrected_are_outputs, + bool confidence_are_outputs, std::vector & arena, MarkovChainGraph & out) { const int hdim = dw.n_embd; @@ -252,6 +254,12 @@ bool build_markov_chain_graph(const DraftWeights & dw, ggml_tensor * prev_ids = out.inp_seed; out.toks.assign((size_t)n_corr, nullptr); out.corrected.assign((size_t)n_corr, nullptr); + out.confidence.assign((size_t)n_corr, nullptr); + const bool have_confidence = confidence_are_outputs && + dw.dspark.confidence_w != nullptr && + dw.dspark.confidence_b != nullptr && + (dw.dspark.confidence_dim == hdim || + dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); for (int i = 0; i < n_corr; ++i) { const int row = first_corrected + i; ggml_tensor * prev_emb = ggml_get_rows(out.ctx, dw.dspark.markov_w1, prev_ids); @@ -268,6 +276,23 @@ bool build_markov_chain_graph(const DraftWeights & dw, ggml_build_forward_expand(out.gf, tok); out.corrected[(size_t)i] = corrected; out.toks[(size_t)i] = tok; + if (have_confidence) { + ggml_tensor * hidden_i = ggml_view_2d( + out.ctx, out.inp_hidden, hdim, 1, out.inp_hidden->nb[1], + (size_t)row * out.inp_hidden->nb[1]); + ggml_tensor * conf_in = hidden_i; + if (dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank) { + conf_in = ggml_concat(out.ctx, hidden_i, prev_emb, 0); + } + ggml_tensor * conf = ggml_mul_mat(out.ctx, dw.dspark.confidence_w, conf_in); + conf = ggml_add( + out.ctx, conf, + ggml_reshape_2d(out.ctx, dw.dspark.confidence_b, 1, 1)); + conf = ggml_sigmoid(out.ctx, conf); + ggml_set_output(conf); + ggml_build_forward_expand(out.gf, conf); + out.confidence[(size_t)i] = conf; + } prev_ids = tok; } return true; @@ -281,7 +306,8 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, const float * local_hidden, int q_len, int32_t last_tok, - std::vector & draft_tok) { + std::vector & draft_tok, + std::vector * confidence_out) { if (q_len <= 1) return false; if (!dspark_fused_usable(dw, backend, lm_head, local_hidden, "dspark_fused")) return false; const int hdim = dw.n_embd; @@ -289,8 +315,12 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, static thread_local std::vector g_arena_chain; MarkovChainGraph g; + const bool want_confidence = confidence_out != nullptr; + if (confidence_out) confidence_out->clear(); if (!build_markov_chain_graph(dw, lm_head, n_cand, /*first_corrected=*/0, - /*corrected_are_outputs=*/false, g_arena_chain, g)) { + /*corrected_are_outputs=*/false, + /*confidence_are_outputs=*/want_confidence, + g_arena_chain, g)) { return false; } @@ -319,14 +349,22 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, draft_tok[0] = last_tok; // One synchronize instead of n_cand blocking readbacks. int32_t t_out[16]; + float c_out[16] = {}; const int n_get = n_cand < 16 ? n_cand : 16; for (int i = 0; i < n_get; ++i) { ggml_backend_tensor_get_async(backend, g.toks[(size_t)i], &t_out[i], 0, sizeof(int32_t)); + if (want_confidence && g.confidence[(size_t)i]) { + ggml_backend_tensor_get_async( + backend, g.confidence[(size_t)i], &c_out[i], 0, sizeof(float)); + } } ggml_backend_synchronize(backend); for (int i = 0; i < n_get; ++i) { draft_tok[(size_t)i + 1] = t_out[i]; } + if (want_confidence && !g.confidence.empty() && g.confidence[0]) { + confidence_out->assign(c_out, c_out + n_get); + } ggml_free(g.ctx); return true; } @@ -347,7 +385,9 @@ bool dspark_markov_project_topk(const DraftWeights & dw, static thread_local std::vector g_arena_topk; MarkovChainGraph g; if (!build_markov_chain_graph(dw, lm_head, n_tokens, /*first_corrected=*/1, - /*corrected_are_outputs=*/true, g_arena_topk, g)) { + /*corrected_are_outputs=*/true, + /*confidence_are_outputs=*/false, + g_arena_topk, g)) { return false; } diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index b76815ce5..ad778e897 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -20,15 +20,17 @@ bool dspark_markov_correct_greedy_chain(const DraftWeights & dw, // Fused variant: base logits (one lm_head matmul over all candidates) + // unrolled Markov correction chain + in-graph argmax feeding the next // step's get_rows, all in ONE graph on the draft backend. No host logits -// round-trip. Does not implement the confidence gate; callers wanting -// confidence-prefix truncation must use the unfused path. +// round-trip. When confidence_out is non-null and the checkpoint has a +// compatible confidence head, returns one score per candidate from the same +// graph and host synchronization as the token ids. bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, const float * local_hidden, int q_len, int32_t last_tok, - std::vector & draft_tok); + std::vector & draft_tok, + std::vector * confidence_out = nullptr); // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index f3cec9b3d..48527e391 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -225,6 +225,12 @@ bool spec_env_flag(const char * name) { return v && *v && *v != '0'; } +// Calibrated cumulative-confidence thresholds for widening the DS4 verify. +// They are part of the policy, not deployment knobs: artifacts without a +// compatible confidence head transparently retain the existing EWMA policy. +constexpr float kConfidenceQ3Threshold = 0.30f; +constexpr float kConfidenceQ4Threshold = 0.25f; + // ── Light rollback state ──────────────────────────────────────────────── // Only the non-position-idempotent cache pieces: the prev-half rows of every // ratio-4 rolling compressor state (attn + indexer) plus hc_state. Everything @@ -384,6 +390,13 @@ bool run_deepseek4_dspark_spec_decode( if (std::fscanf(f, "%d", &v) == 1) adaptive_width = (v != 0); std::fclose(f); } + const bool use_confidence_width = adaptive_width && !seq_verify_mode && + drafter.confidence_w != nullptr && drafter.confidence_b != nullptr && + (drafter.confidence_dim == n_embd || + drafter.confidence_dim == n_embd + drafter.markov_rank); + if (timing && use_confidence_width) { + std::fprintf(stderr, "[ds4-spec] adaptive width policy=confidence\n"); + } double ewma_accept = 1.5; // Fast path caps the verify at the compression ratio (4): one boundary max, @@ -452,6 +465,7 @@ bool run_deepseek4_dspark_spec_decode( std::vector noise_ids(block); std::vector local_hidden, padded_hidden((size_t) n_embd * (block + 1), 0.0f); std::vector draft_tok, tgt_am; + std::vector draft_confidence; // Cumulative phase timings (ms). double tm_draft = 0, tm_head = 0, tm_save = 0, tm_verify = 0, tm_apply = 0, tm_feat = 0; @@ -491,6 +505,7 @@ bool run_deepseek4_dspark_spec_decode( // let the (row-0-skipping) chain use slots 1..q-1. t0 = SpecClock::now(); draft_tok.clear(); + draft_confidence.clear(); bool ds_ok = false; // Batched-verify exactness: the batch must not cross a ratio-4 // boundary except at its last token (state rows stay distinct and the @@ -503,15 +518,17 @@ bool run_deepseek4_dspark_spec_decode( int q_step_cap = (seq_verify_mode || fused_verify_mode) ? std::min(q_cap, 4) : std::min(q_cap, 4 - (pos & 3)); - if (adaptive_width && !seq_verify_mode) { + if (adaptive_width && !use_confidence_width && !seq_verify_mode) { const int w_cap = (int) ewma_accept + 2; if (w_cap < q_step_cap) q_step_cap = w_cap; } if (q_step_cap >= 2) { std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), sizeof(float) * (size_t) n_embd * block); - ds_ok = dspark_markov_correct_greedy_chain_fused(dw, backend, target.lm_head_tensor(), - padded_hidden.data(), q_step_cap, lt, draft_tok); + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw, backend, target.lm_head_tensor(), padded_hidden.data(), + q_step_cap, lt, draft_tok, + use_confidence_width ? &draft_confidence : nullptr); if (!ds_ok) { ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, padded_hidden.data(), q_step_cap, lt, 0.0f, draft_tok); @@ -527,6 +544,25 @@ bool run_deepseek4_dspark_spec_decode( } else { draft_tok.push_back(lt); // q=1: seed only, no speculation } + // Confidence estimates are per candidate. Their cumulative product is + // the estimated probability that the target accepts the whole prefix + // unlocked by a wider verify. The defaults were calibrated on q=4 + // traces and keep q=4 for high-confidence prefixes while avoiding its + // extra verify cost on low-acceptance prompts. + if (use_confidence_width && draft_confidence.size() >= 2 && draft_tok.size() >= 3) { + const float p2 = draft_confidence[0] * draft_confidence[1]; + int selected_q = p2 >= kConfidenceQ3Threshold ? 3 : 2; + if (selected_q == 3 && draft_confidence.size() >= 3 && draft_tok.size() >= 4) { + const float p3 = p2 * draft_confidence[2]; + if (p3 >= kConfidenceQ4Threshold) selected_q = 4; + } + if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); + } else if (use_confidence_width && !seq_verify_mode) { + // The fused head should always return confidence for a compatible + // artifact. Preserve the old policy if a backend cannot do so. + const int selected_q = (int) ewma_accept + 2; + if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); + } if ((int) draft_tok.size() > q_step_cap) draft_tok.resize(q_step_cap); const int q = (int) draft_tok.size(); // seed + candidates tm_head += spec_ms_since(t0); From d60fb936fbca932d7e66cde2bf001561830fd1a8 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 11:42:32 -0700 Subject: [PATCH 04/12] fix(ds4): integrate DSpark with current backend APIs --- server/src/deepseek4/deepseek4_backend.cpp | 37 ++++++++++++++++------ server/src/deepseek4/deepseek4_graph.cpp | 15 +++++++-- 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index b3c1a14f6..c9f3e12e5 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -762,8 +762,13 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, // Decode auto t1 = Clock::now(); - if (spec_enabled_ && spec_drafter_ && req.n_gen > 0) { - if (last_logits_.empty()) { result.error = "spec: no prefill logits"; return result; } + const bool budget_requires_ar = !req.budget_hook.close_token_ids.empty(); + if (spec_enabled_ && spec_drafter_ && req.n_gen > 0 && + !req.force_ar_decode && !budget_requires_ar) { + if (last_logits_.empty()) { + result.fail(GenerateErrorCode::DecodeFailed, "spec: no prefill logits"); + return result; + } int seed = 0; { float mv = last_logits_[0]; for (int i = 1; i < w_.n_vocab; i++) if (last_logits_[i] > mv) { mv = last_logits_[i]; seed = i; } } @@ -771,20 +776,32 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, gen.push_back(seed); io.emit(seed); float accept_rate = 0.0f; - if (!deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { + bool spec_ran = false; + if (!io.cancelled && !deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; std::vector spec_toks; - run_deepseek4_dspark_spec_decode(backend_, w_, cache_, *spec_drafter_, - committed, seed, req.n_gen - 1, - win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, - spec_toks, &accept_rate); - for (int t : spec_toks) io.emit(t); - gen.insert(gen.end(), spec_toks.begin(), spec_toks.end()); + spec_ran = true; + if (!run_deepseek4_dspark_spec_decode( + backend_, w_, cache_, *spec_drafter_, committed, seed, + req.n_gen - 1, + win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, + spec_toks, &accept_rate)) { + result.fail(GenerateErrorCode::DecodeFailed, + "DSpark speculative decode failed"); + return result; + } + for (int32_t tok : spec_toks) { + if (io.cancelled) break; + gen.push_back(tok); + io.emit(tok); + } } - result.ok = true; + result.succeed(); result.tokens = std::move(gen); result.decode_s = elapsed_s(t1); + result.accept_rate = accept_rate; + result.spec_decode_ran = spec_ran; std::fprintf(stderr, "[deepseek4] DSpark decode: %zu tok in %.3fs (%.1f tok/s) accept_rate=%.2f\n", result.tokens.size(), result.decode_s, result.decode_s > 0 ? result.tokens.size() / result.decode_s : 0.0, accept_rate); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index ec2f79cea..b9c54137b 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -47,6 +48,11 @@ static uint64_t ds4_elapsed_us(Ds4TimingClock::time_point start, return (uint64_t)std::chrono::duration_cast(end - start).count(); } +static bool ds4_env_flag(const char * name) { + const char * value = std::getenv(name); + return value && value[0] && std::strcmp(value, "0") != 0; +} + static int ds4_effective_expert_count(const DeepSeek4Weights & w) { const int requested = w.routed_expert_top_k; if (requested > 0 && requested < w.n_expert_used) { @@ -259,6 +265,9 @@ struct DeepSeek4AttentionGraphInputs { // [n_swa raw rows ++ padded comp rows]; 0 for valid, -1e30 for padding. ggml_tensor * attn_row_mask = nullptr; int padded_comp = 0; // padded compressed-row count (>= n_comp) + // Optional stable-topology compressor rows. DSpark's fused verifier leaves + // this null and uses its batched state-row inputs instead. + ggml_tensor * flush_rows = nullptr; }; struct DeepSeek4CachedDecodeAttnGraph { @@ -3048,8 +3057,7 @@ static bool deepseek4_step_hybrid( const int32_t * token_ids, MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, - MoeHybridRoutingStats * routing_stats, - Ds4VerifyHooks * verify_hooks) { + MoeHybridRoutingStats * routing_stats) { const auto step_t0 = Ds4TimingClock::now(); const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -3523,7 +3531,8 @@ bool deepseek4_step( const int32_t * token_ids, MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, - MoeHybridRoutingStats * routing_stats) { + MoeHybridRoutingStats * routing_stats, + Ds4VerifyHooks * verify_hooks) { if (w.moe_hybrid && moe_hybrid != nullptr) { return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, embed, n_tokens, kv_start, out_logits, From b7c184ced2ec54f9317a2924b294b3fb52e2552e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 13 Jul 2026 23:18:12 +0200 Subject: [PATCH 05/12] perf(ds4): default gfx1151 verifier q4 to MMVQ --- server/docs/DS4.md | 7 +++++ server/src/deepseek4/deepseek4_backend.cpp | 32 ++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index ee2b5da69..cf5a53945 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -180,6 +180,13 @@ export DFLASH_DS4_SPEC_Q=4 ./server/build-hip/dflash_server /path/to/deepseek4-target.gguf ``` +On HIP `gfx1151`, enabling DSpark defaults `LUCE_MMVQ_MAX_NCOLS` to `4` when +the variable is unset. This keeps the four-row verifier on MMVQ; the validated +ROCm 7.2.4 high-clock GSM+Math run reached 30.23 tok/s weighted with 10/10 +correctness. Set `LUCE_MMVQ_MAX_NCOLS` explicitly to override the platform +default. AR, NVIDIA, and other HIP architectures retain the shared dispatch +default. + Adaptive width is automatic. When the draft artifact has a compatible confidence projection, the runtime selects q=2, q=3, or q=4 from the cumulative confidence of the proposed prefix. It adds the projection to the same fused diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index c9f3e12e5..5d9fb2fe7 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -4,6 +4,10 @@ #include "deepseek4_internal.h" #include "common/sampler.h" +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) +#include "common/gpu_runtime_compat.h" +#endif + #include "ggml.h" #include "ggml-backend.h" #include "ggml-cuda.h" @@ -33,6 +37,29 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +static void configure_gfx1151_dspark_mmvq_default(int gpu) { +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + if (!env_flag_enabled("DFLASH_DS4_SPEC") || + std::getenv("LUCE_MMVQ_MAX_NCOLS") != nullptr) { + return; + } + + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, gpu) != cudaSuccess || + std::strncmp(prop.gcnArchName, "gfx1151", 7) != 0) { + return; + } + + if (::setenv("LUCE_MMVQ_MAX_NCOLS", "4", 0) == 0) { + std::fprintf(stderr, + "[deepseek4] gfx1151 DSpark: defaulting " + "LUCE_MMVQ_MAX_NCOLS=4\n"); + } +#else + (void) gpu; +#endif +} + static double gib(uint64_t bytes) { return (double) bytes / 1024.0 / 1024.0 / 1024.0; } @@ -371,6 +398,11 @@ bool DeepSeek4Backend::load_model() { } bool DeepSeek4Backend::init() { + // The shared MMVQ/MMQ crossover defaults to q=3 for NVIDIA. On gfx1151, + // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged, + // and preserve LUCE_MMVQ_MAX_NCOLS as an explicit override. + configure_gfx1151_dspark_mmvq_default(cfg_.device.gpu); + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!backend_) { std::fprintf(stderr, "[deepseek4] failed to create CUDA backend (gpu=%d)\n", From 2c7378aadfef466575283c7cba1f9dafc0310eca Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 12:10:43 -0700 Subject: [PATCH 06/12] fix(ds4): harden speculative runtime lifecycle --- server/docs/DS4.md | 12 ++- server/src/deepseek4/deepseek4_backend.cpp | 42 +++++++-- server/src/deepseek4/deepseek4_dspark.cpp | 90 +++++++++++++++++++ server/src/deepseek4/deepseek4_dspark.h | 8 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 32 +++++-- server/src/deepseek4/deepseek4_graph.cpp | 13 +++ 6 files changed, 180 insertions(+), 17 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index cf5a53945..a76261634 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -177,9 +177,16 @@ export DFLASH_DS4_FUSED_VERIFY=1 export DFLASH_DS4_DRAFT=/path/to/dflash-draft.gguf export DFLASH_DS4_SPEC_Q=4 -./server/build-hip/dflash_server /path/to/deepseek4-target.gguf +./server/build-hip/dflash_server /path/to/deepseek4-target.gguf \ + --target-device hip:0 \ + --ds4-fused-decode ``` +DSpark currently requires monolithic target placement. On HIP, +`--ds4-fused-decode` selects that placement; if the target falls back to hybrid +expert placement, the server logs that DSpark is disabled and continues with +the normal autoregressive path. + On HIP `gfx1151`, enabling DSpark defaults `LUCE_MMVQ_MAX_NCOLS` to `4` when the variable is unset. This keeps the four-row verifier on MMVQ; the validated ROCm 7.2.4 high-clock GSM+Math run reached 30.23 tok/s weighted with 10/10 @@ -199,7 +206,8 @@ GSM+Math accuracy and measured 29.25 tok/s weighted, within 0.8% of fixed q=4 at 29.49 tok/s. On the low-acceptance stress prompt it measured 21.9/21.8 tok/s warm, effectively tied with EWMA while avoiding fixed q=4's wasted wide verification. These numbers are workload-specific; the confidence policy is -opt-in. +enabled only when DSpark is explicitly enabled and the draft artifact contains +a compatible confidence head. ## Example: CUDA + Halo Layer Split diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 5d9fb2fe7..9aaef45f1 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -445,8 +445,29 @@ bool DeepSeek4Backend::init() { if (dp && *dp) { spec_drafter_ = std::make_unique(); if (load_deepseek4_dspark_drafter(dp, backend_, *spec_drafter_)) { - spec_enabled_ = true; - std::fprintf(stderr, "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", dp); + const DSparkDrafter & d = *spec_drafter_; + bool compatible = d.core.n_embd == w_.n_embd && + d.core.n_vocab == w_.n_vocab && + d.vocab_size == w_.n_vocab && + d.mask_token_id >= 0 && d.mask_token_id < w_.n_vocab && + (int) d.capture_layer_ids.size() == d.n_target_layers; + for (int layer : d.capture_layer_ids) { + compatible = compatible && layer >= 0 && layer < w_.n_layer; + } + if (!compatible) { + std::fprintf(stderr, + "[deepseek4] DSpark drafter is incompatible with target " + "(target embd/vocab/layers=%d/%d/%d, draft=%d/%d)\n", + w_.n_embd, w_.n_vocab, w_.n_layer, + d.core.n_embd, d.vocab_size); + free_deepseek4_dspark_drafter(*spec_drafter_); + spec_drafter_.reset(); + } else { + spec_enabled_ = true; + std::fprintf(stderr, + "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", + dp); + } } else { std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", deepseek4_dspark_last_error()); @@ -818,16 +839,17 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, backend_, w_, cache_, *spec_drafter_, committed, seed, req.n_gen - 1, win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, - spec_toks, &accept_rate)) { + spec_toks, &accept_rate, + [&io](int32_t tok) { + if (io.cancelled) return false; + io.emit(tok); + return !io.cancelled; + })) { result.fail(GenerateErrorCode::DecodeFailed, "DSpark speculative decode failed"); return result; } - for (int32_t tok : spec_toks) { - if (io.cancelled) break; - gen.push_back(tok); - io.emit(tok); - } + gen.insert(gen.end(), spec_toks.begin(), spec_toks.end()); } result.succeed(); result.tokens = std::move(gen); @@ -896,6 +918,9 @@ bool DeepSeek4Backend::handle_compress(const std::string & line, } void DeepSeek4Backend::free_drafter() { + if (spec_drafter_) { + free_deepseek4_dspark_drafter(*spec_drafter_); + } spec_drafter_.reset(); spec_enabled_ = false; spec_feat_window_.clear(); @@ -912,6 +937,7 @@ void DeepSeek4Backend::maybe_save_routing_stats() { void DeepSeek4Backend::shutdown() { maybe_save_routing_stats(); + free_drafter(); for (int i = 0; i < PREFIX_SLOTS; i++) { free_deepseek4_snapshot(snapshots_[i]); } diff --git a/server/src/deepseek4/deepseek4_dspark.cpp b/server/src/deepseek4/deepseek4_dspark.cpp index 8975244d1..ac36d120c 100644 --- a/server/src/deepseek4/deepseek4_dspark.cpp +++ b/server/src/deepseek4/deepseek4_dspark.cpp @@ -245,6 +245,95 @@ bool load_deepseek4_dspark_drafter(const std::string & path, else if (suffix == "hc_ffn_base.weight") L.hc_ffn_base = t; } + // Reject incomplete artifacts here, before production code can build a + // graph that dereferences a missing tensor. The original smoke test + // checked these bindings only after load_deepseek4_dspark_drafter had + // already returned success. + std::vector missing; + auto need = [&](const void * ptr, const std::string & name) { + if (!ptr) missing.push_back(name); + }; + need(out.main_proj, "dflash.fc.weight"); + need(out.main_norm, "dflash.hidden_norm.weight"); + need(out.markov_w1, "dflash.dspark.markov.w1"); + need(out.markov_w2, "dflash.dspark.markov.w2"); + need(w.out_norm, "output_norm.weight"); + need(w.output_hc_fn, "output_hc_fn.weight"); + need(w.output_hc_scale, "output_hc_scale.weight"); + need(w.output_hc_base, "output_hc_base.weight"); + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + const std::string p = "blk." + std::to_string(il) + "."; + need(L.attn_norm, p + "attn_norm.weight"); + need(L.attn_q_a, p + "attn_q_a.weight"); + need(L.attn_q_a_norm, p + "attn_q_a_norm.weight"); + need(L.attn_q_b, p + "attn_q_b.weight"); + need(L.attn_kv, p + "attn_kv.weight"); + need(L.attn_kv_a_norm, p + "attn_kv_a_norm.weight"); + need(L.attn_output_a, p + "attn_output_a.weight"); + need(L.attn_output_b, p + "attn_output_b.weight"); + need(L.hc_attn_fn, p + "hc_attn_fn.weight"); + need(L.hc_attn_scale, p + "hc_attn_scale.weight"); + need(L.hc_attn_base, p + "hc_attn_base.weight"); + need(L.ffn_norm, p + "ffn_norm.weight"); + need(L.ffn_gate_inp, p + "ffn_gate_inp.weight"); + need(L.ffn_gate_exps, p + "ffn_gate_exps.weight"); + need(L.ffn_up_exps, p + "ffn_up_exps.weight"); + need(L.ffn_down_exps, p + "ffn_down_exps.weight"); + need(L.ffn_gate_shexp, p + "ffn_gate_shexp.weight"); + need(L.ffn_up_shexp, p + "ffn_up_shexp.weight"); + need(L.ffn_down_shexp, p + "ffn_down_shexp.weight"); + need(L.hc_ffn_fn, p + "hc_ffn_fn.weight"); + need(L.hc_ffn_scale, p + "hc_ffn_scale.weight"); + need(L.hc_ffn_base, p + "hc_ffn_base.weight"); + } + + std::string contract_error; + if (!out.dspark_enabled) { + contract_error = "dflash.dspark.enabled is false"; + } else if (w.n_layer <= 0 || w.n_embd <= 0 || w.n_vocab <= 0 || + w.n_hc <= 0 || out.n_target_layers <= 0 || + out.block_size < 2 || out.block_size > 16) { + contract_error = "invalid DSpark architecture metadata"; + } else if ((int) out.capture_layer_ids.size() != out.n_target_layers) { + contract_error = "capture_layer_ids count does not match dflash.n_target_layers"; + } else if (out.main_proj && + (out.main_proj->ne[0] != (int64_t) out.n_target_layers * w.n_embd || + out.main_proj->ne[1] != w.n_embd)) { + contract_error = "dflash.fc.weight shape does not match target-layer metadata"; + } else if (out.markov_w1 && out.markov_w2 && + (out.markov_w1->ne[0] != out.markov_rank || + out.markov_w1->ne[1] != out.vocab_size || + out.markov_w2->ne[0] != out.markov_rank || + out.markov_w2->ne[1] != out.vocab_size)) { + contract_error = "DSpark Markov tensor shape does not match metadata"; + } else if ((out.confidence_w == nullptr) != (out.confidence_b == nullptr)) { + contract_error = "incomplete DSpark confidence tensor pair"; + } else if (out.confidence_w && + (out.confidence_dim <= 0 || + out.confidence_w->ne[0] != out.confidence_dim || + out.confidence_w->ne[1] != 1 || + ggml_nelements(out.confidence_b) != 1)) { + contract_error = "DSpark confidence tensor shape does not match metadata"; + } + if (!missing.empty() || !contract_error.empty()) { + std::string message = "invalid DSpark drafter GGUF"; + if (!contract_error.empty()) message += ": " + contract_error; + if (!missing.empty()) { + message += "; missing "; + for (size_t i = 0; i < missing.size(); ++i) { + if (i != 0) message += ", "; + message += missing[i]; + } + } + set_err(message); + ggml_backend_buffer_free(buf); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } + w.ctx = meta; w.buf = buf; gguf_free(g); // meta_ctx now owned by w.ctx; do not free here @@ -259,6 +348,7 @@ bool load_deepseek4_dspark_drafter(const std::string & path, } void free_deepseek4_dspark_drafter(DSparkDrafter & d) { + reset_deepseek4_dspark_runtime_cache(); if (d.core.buf) { ggml_backend_buffer_free(d.core.buf); d.core.buf = nullptr; } if (d.core.ctx) { ggml_free(d.core.ctx); d.core.ctx = nullptr; } d = DSparkDrafter{}; diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 9fe20da75..9cddf3beb 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -29,6 +29,7 @@ #include "ggml.h" #include "ggml-backend.h" +#include #include #include @@ -70,6 +71,10 @@ bool load_deepseek4_dspark_drafter(const std::string & path, void free_deepseek4_dspark_drafter(DSparkDrafter & d); +// Drop thread-local graph/allocation state that retains drafter tensor and +// backend references. Called before releasing the drafter weights. +void reset_deepseek4_dspark_runtime_cache(); + const char * deepseek4_dspark_last_error(); // One drafter forward. Produces block_size normed hidden states (the input to @@ -136,6 +141,7 @@ bool run_deepseek4_dspark_spec_decode( const float * prompt_feature_window, // [n_target_layers*n_embd * win_len] captured during prefill int win_len, std::vector & out_tokens, - float * accept_rate_out); + float * accept_rate_out, + const std::function & on_token = {}); } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 48527e391..327217c82 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -370,7 +370,8 @@ bool run_deepseek4_dspark_spec_decode( const float * prompt_feature_window, int win_len, std::vector & out_tokens, - float * accept_rate_out) { + float * accept_rate_out, + const std::function & on_token) { const int n_embd = target_w.n_embd; const int n_tgt = drafter.n_target_layers; const int block = drafter.block_size; @@ -460,6 +461,8 @@ bool run_deepseek4_dspark_spec_decode( int pos = committed; // absolute position of the seed (block slot 0) int n_generated = 0; long accept_sum = 0, steps = 0; + bool ok = true; + bool stop_requested = false; std::vector noise_embed((size_t) n_embd * block); std::vector noise_ids(block); @@ -479,13 +482,17 @@ bool run_deepseek4_dspark_spec_decode( if (q_cap >= 2) { noise_ids[0] = lt; for (int i = 1; i < block; i++) noise_ids[i] = drafter.mask_token_id; - if (!target.embed_tokens(noise_ids.data(), block, noise_embed.data())) break; + if (!target.embed_tokens(noise_ids.data(), block, noise_embed.data())) { + ok = false; + break; + } // Drafter forward -> block normed hidden states. if (!deepseek4_dspark_draft_forward(backend, drafter, noise_embed.data(), ctx_len > 0 ? feat_win.data() : nullptr, ctx_len, pos, local_hidden)) { std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); + ok = false; break; } } @@ -536,7 +543,10 @@ bool run_deepseek4_dspark_spec_decode( if (!ds_ok || (int) draft_tok.size() < 2) { // Fallback: plain projection of the block hiddens. std::vector pj; - if (!target.project_hidden_to_tokens(local_hidden.data(), q_step_cap - 1, pj)) break; + if (!target.project_hidden_to_tokens(local_hidden.data(), q_step_cap - 1, pj)) { + ok = false; + break; + } draft_tok.clear(); draft_tok.push_back(lt); for (int i = 0; i < q_step_cap - 1; i++) draft_tok.push_back(pj[i]); @@ -577,7 +587,11 @@ bool run_deepseek4_dspark_spec_decode( // ── Rollback state save (cheap) or legacy full snapshot ── t0 = SpecClock::now(); if (full_snap) { - if (!target.snapshot_kv()) { std::fprintf(stderr, "[ds4-spec] snapshot failed\n"); break; } + if (!target.snapshot_kv()) { + std::fprintf(stderr, "[ds4-spec] snapshot failed\n"); + ok = false; + break; + } } else { spec_rollback_save(target_cache, rollback); } @@ -597,6 +611,7 @@ bool run_deepseek4_dspark_spec_decode( spec_rollback_apply(rollback, target_w, target_cache, pos, boundary_crossed); } std::fprintf(stderr, "[ds4-spec] verify failed\n"); + ok = false; break; } tm_verify += spec_ms_since(t0); @@ -637,6 +652,7 @@ bool run_deepseek4_dspark_spec_decode( std::vector replay_am; if (!target.verify_batch(kv_toks, pos, replay_last, &replay_am)) { std::fprintf(stderr, "[ds4-spec] replay verify failed\n"); + ok = false; break; } } else if (accept < q) { @@ -662,6 +678,10 @@ bool run_deepseek4_dspark_spec_decode( const int t = (i < accept) ? draft_tok[i] : bonus; out_tokens.push_back(t); n_generated++; + if (on_token && !on_token(t)) { + stop_requested = true; + break; + } if (target.is_eos(t)) { hit_eos = true; break; } if (n_generated >= n_gen) break; } @@ -678,7 +698,7 @@ bool run_deepseek4_dspark_spec_decode( tm_draft / steps, tm_head / steps, tm_save / steps, tm_verify / steps, tm_apply / steps, tm_feat / steps); } - if (hit_eos) break; + if (hit_eos || stop_requested) break; } const double total_ms = spec_ms_since(run_t0); @@ -716,7 +736,7 @@ bool run_deepseek4_dspark_spec_decode( (unsigned long long) tel.ffn_hot_graph_builds); } ggml_backend_free(snap_backend); - return true; + return ok; } } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index b9c54137b..971a7da6e 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -5358,6 +5358,19 @@ thread_local DsparkDraftCache g_dspark_draft_cache; } // namespace +void reset_deepseek4_dspark_runtime_cache() { + DsparkDraftCache & cache = g_dspark_draft_cache; + if (cache.alloc) { + ggml_gallocr_free(cache.alloc); + cache.alloc = nullptr; + } + if (cache.ctx) { + ggml_free(cache.ctx); + cache.ctx = nullptr; + } + cache = DsparkDraftCache{}; +} + bool deepseek4_dspark_draft_forward(ggml_backend_t backend, const DSparkDrafter & d, const float * noise_embed, From 415432f119e87d4c929894695e5271275c821814 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:01:57 -0700 Subject: [PATCH 07/12] fix(ds4): stabilize verifier snapshot lifecycle --- server/docs/DS4.md | 28 +++++++++++-------- .../src/deepseek4/deepseek4_dspark_spec.cpp | 16 ++++++++--- server/src/deepseek4/deepseek4_graph.cpp | 7 +++-- 3 files changed, 34 insertions(+), 17 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index a76261634..fcd230f8f 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -179,20 +179,28 @@ export DFLASH_DS4_SPEC_Q=4 ./server/build-hip/dflash_server /path/to/deepseek4-target.gguf \ --target-device hip:0 \ - --ds4-fused-decode + --ds4-fused-decode \ + --ds4-expert-top-k 4 ``` DSpark currently requires monolithic target placement. On HIP, `--ds4-fused-decode` selects that placement; if the target falls back to hybrid expert placement, the server logs that DSpark is disabled and continues with -the normal autoregressive path. +the normal autoregressive path. `--ds4-expert-top-k 4` is a separate, +approximate inference policy used by the validated Strix Halo profile; omit it +to retain the model's default six routed experts. On HIP `gfx1151`, enabling DSpark defaults `LUCE_MMVQ_MAX_NCOLS` to `4` when -the variable is unset. This keeps the four-row verifier on MMVQ; the validated -ROCm 7.2.4 high-clock GSM+Math run reached 30.23 tok/s weighted with 10/10 -correctness. Set `LUCE_MMVQ_MAX_NCOLS` explicitly to override the platform -default. AR, NVIDIA, and other HIP architectures retain the shared dispatch -default. +the variable is unset. This keeps the four-row verifier on MMVQ. On a 128 GiB +Strix Halo Radeon 8060S using ROCm 7.2.4, the rebased candidate measured 32.12 +tok/s weighted at fixed q=4 and 31.94 tok/s with confidence-adaptive width, +versus 25.31 tok/s autoregressive. All three configurations scored 10/10 on the +same five GSM and five Math prompts. The run used `--ds4-expert-top-k 4`, the +platform `performance` profile, and the GPU `high` performance level; fixed +q=4 with the model-default six routed experts measured 28.26 tok/s. Enabling +DSpark alone therefore does not guarantee 30 tok/s. Set +`LUCE_MMVQ_MAX_NCOLS` explicitly to override the platform default. AR, NVIDIA, +and other HIP architectures retain the shared dispatch default. Adaptive width is automatic. When the draft artifact has a compatible confidence projection, the runtime selects q=2, q=3, or q=4 from the cumulative @@ -202,10 +210,8 @@ additional host round trip is introduced. Artifacts without a compatible confidence head transparently retain the existing acceptance-EWMA policy. On the gfx1151 validation host, confidence-adaptive width retained 10/10 -GSM+Math accuracy and measured 29.25 tok/s weighted, within 0.8% of fixed q=4 -at 29.49 tok/s. On the low-acceptance stress prompt it measured 21.9/21.8 -tok/s warm, effectively tied with EWMA while avoiding fixed q=4's wasted wide -verification. These numbers are workload-specific; the confidence policy is +GSM+Math accuracy and measured 31.94 tok/s weighted, within 0.6% of fixed q=4 +at 32.12 tok/s. These numbers are workload-specific; the confidence policy is enabled only when DSpark is explicitly enabled and the draft artifact contains a compatible confidence head. diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 327217c82..cc5fa4fa0 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -47,6 +47,8 @@ class DeepSeek4DFlashTarget : public DFlashTarget { : w_(w), cache_(cache), backend_(backend), snap_backend_(snap_backend), capture_ids_(std::move(capture_ids)), mask_tok_(mask_tok) {} + ~DeepSeek4DFlashTarget() override { clear_snapshot(); } + bool verify_batch(const std::vector & tokens, int base_pos, int & last_tok, std::vector * all_argmax = nullptr, bool capture_ssm_intermediates = false) override { @@ -58,10 +60,11 @@ class DeepSeek4DFlashTarget : public DFlashTarget { n, n > 0 ? tokens[0] : -1, n > 1 ? tokens[1] : -1, w_.n_vocab); return false; } - // Sequential exact verify (measurement mode): q single-token forwards - // through the legacy AR decode path. Causal by construction, compressor - // fed every token. Slow; used to measure the drafter's TRUE accept - // rate and produce exact greedy output. Enable: DFLASH_DS4_SEQ_VERIFY=1 + // Sequential verify (measurement mode): q single-token forwards through + // the legacy AR decode path. Causal by construction, compressor fed every + // token. Slow; used to measure the drafter's token-at-a-time accept rate. + // It is not a bit-exact oracle: graph shape can change floating-point + // reduction order around near-tied logits. Enable: DFLASH_DS4_SEQ_VERIFY=1 // (pair with DFLASH_DS4_FULL_SNAP=1 so rollback/replay stay exact). static const bool seq_verify = [] { const char * v = std::getenv("DFLASH_DS4_SEQ_VERIFY"); @@ -185,6 +188,7 @@ class DeepSeek4DFlashTarget : public DFlashTarget { void set_telemetry(DeepSeek4StepTelemetry * t) { telemetry_ = t; } const std::vector & last_features() const { return verify_features_; } int last_verify_n() const { return verify_n_; } + void clear_snapshot() { free_deepseek4_snapshot(snap_); } private: const DeepSeek4Weights & w_; @@ -735,6 +739,10 @@ bool run_deepseek4_dspark_spec_decode( (unsigned long long) tel.ffn_hot_graph_hits, (unsigned long long) tel.ffn_hot_graph_builds); } + // Snapshot buffers must be released while their backend is still alive. + // clear_snapshot() is idempotent, so the target destructor remains a + // safety net for future exits that are added above this point. + target.clear_snapshot(); ggml_backend_free(snap_backend); return ok; } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 971a7da6e..e34b4bc61 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -4295,7 +4295,10 @@ bool deepseek4_step_layer_range( std::vector & attn_out_host = scratch.attn_out_host; std::vector & ffn_out_host = scratch.ffn_out_host; std::vector & hash_expert_ids_host = scratch.hash_expert_ids; - const bool reuse_decode_graphs = n_tokens == 1; + // Verification hooks need per-position captures and logits. The cached + // single-token decode graphs only expose the final decode outputs, so + // honor the caller's request to take the dynamic path for verification. + const bool reuse_decode_graphs = n_tokens == 1 && allow_decode_graph_reuse; Ds4DecodeSharedInputs * shared_inputs = nullptr; if (reuse_decode_graphs && ds4_backend_is_gpu(backend) && decode_shared_inputs.ensure(w, backend)) { @@ -5194,8 +5197,8 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { - if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; } if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } + if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; } s.layers.clear(); s.cur_pos = 0; s.hc_state_snap = nullptr; From 07b8ad783d41749de3de10f8499cfe3cd783efd9 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:27:09 -0700 Subject: [PATCH 08/12] fix(ds4): address speculative review findings --- server/scripts/convert_dflash_to_gguf.py | 4 - server/src/deepseek4/deepseek4_backend.cpp | 36 ++- server/src/deepseek4/deepseek4_dspark.cpp | 253 +++++++++++++++--- server/src/deepseek4/deepseek4_dspark.h | 2 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 16 +- .../src/deepseek4/deepseek4_fused_verify.inc | 36 ++- server/src/deepseek4/deepseek4_graph.cpp | 11 +- server/tests/test_deepseek4_unit.cpp | 168 ++++++++++++ 8 files changed, 469 insertions(+), 57 deletions(-) diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index 6c1c46291..b904d7ea5 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -411,10 +411,6 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): ) print("[warn] DSpark confidence head has no bias tensor; writing a zero bias") - if conf_missing and bias_names in conf_missing and len(conf_missing) > 1: - print("[warn] incomplete DSpark confidence head; Markov head will still load") - return - conf_w = conf_resolved[weight_names][1] conf_b = conf_resolved[bias_names][1] confidence_dim = int(conf_w.shape[1]) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 9aaef45f1..9ba2f12bf 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -652,7 +652,30 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, reset_deepseek4_cache(cache_); } last_logits_.clear(); - if (spec_enabled_ && kv_offset == 0) spec_feat_window_.clear(); + int spec_capture_from = n_total; + if (spec_enabled_ && spec_drafter_) { + const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; + if (kv_offset == 0 || n_total >= w_.n_swa || feat_row <= 0 || + spec_feat_window_.size() % (size_t) feat_row != 0) { + spec_feat_window_.clear(); + spec_capture_from = std::max(0, n_total - w_.n_swa); + } else { + // Keep enough prior rows for the new prompt suffix, then append all + // new rows. This bounds host capture storage at n_swa without + // shifting a multi-megabyte feature window after every token. + const size_t old_rows = spec_feat_window_.size() / (size_t) feat_row; + const size_t keep_rows = (size_t) std::max(0, w_.n_swa - n_total); + if (old_rows > keep_rows) { + const size_t drop_floats = (old_rows - keep_rows) * (size_t) feat_row; + const size_t keep_floats = keep_rows * (size_t) feat_row; + std::memmove(spec_feat_window_.data(), + spec_feat_window_.data() + drop_floats, + keep_floats * sizeof(float)); + spec_feat_window_.resize(keep_floats); + } + spec_capture_from = 0; + } + } const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); DeepSeek4StepTelemetry tel_acc; @@ -674,7 +697,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, Ds4VerifyHooks spec_hooks; std::vector spec_cap; Ds4VerifyHooks * hp = nullptr; - if (spec_enabled_ && spec_drafter_) { + if (spec_enabled_ && spec_drafter_ && i + n_tok > spec_capture_from) { spec_hooks.capture_layer_ids = &spec_drafter_->capture_layer_ids; spec_hooks.capture_out = &spec_cap; hp = &spec_hooks; @@ -687,7 +710,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, routing_stats_.get(), hp); if (ok && hp && !spec_cap.empty()) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; - for (int t = 0; t < n_tok; ++t) { + const int first_capture = std::max(0, spec_capture_from - i); + for (int t = first_capture; t < n_tok; ++t) { spec_feat_window_.insert( spec_feat_window_.end(), spec_cap.begin() + (size_t) t * feat_row, @@ -807,6 +831,12 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, } result.prefill_s = elapsed_s(t0); + if (io.cancelled) { + result.succeed(); + maybe_save_routing_stats(); + return result; + } + if (req.n_gen <= 0) { result.succeed(); maybe_save_routing_stats(); diff --git a/server/src/deepseek4/deepseek4_dspark.cpp b/server/src/deepseek4/deepseek4_dspark.cpp index ac36d120c..81300e48f 100644 --- a/server/src/deepseek4/deepseek4_dspark.cpp +++ b/server/src/deepseek4/deepseek4_dspark.cpp @@ -7,6 +7,8 @@ #include "deepseek4_dspark.h" +#include "common/gguf_bounds.h" + #include "ggml.h" #include "ggml-backend.h" #include "gguf.h" @@ -14,10 +16,12 @@ #include #include #include +#include #include #include #include +#include #include namespace dflash::common { @@ -77,6 +81,53 @@ bool block_suffix(const char * name, int & il, std::string & suffix) { return true; } +bool recognized_block_suffix(const std::string & suffix) { + static const char * const names[] = { + "attn_norm.weight", "attn_q_a.weight", "attn_q_a_norm.weight", + "attn_q_b.weight", "attn_kv.weight", "attn_kv_a_norm.weight", + "attn_sinks.weight", "attn_output_a.weight", "attn_output_b.weight", + "hc_attn_fn.weight", "hc_attn_scale.weight", "hc_attn_base.weight", + "ffn_norm.weight", "ffn_gate_inp.weight", "exp_probs_b.bias", + "ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight", + "ffn_gate_shexp.weight", "ffn_up_shexp.weight", "ffn_down_shexp.weight", + "hc_ffn_fn.weight", "hc_ffn_scale.weight", "hc_ffn_base.weight", + }; + for (const char * name : names) { + if (suffix == name) return true; + } + return false; +} + +bool recognized_dspark_tensor(const char * name, int n_layer) { + static const char * const names[] = { + "output_norm.weight", "output_hc_base.weight", "output_hc_fn.weight", + "output_hc_scale.weight", "dflash.fc.weight", + "dflash.hidden_norm.weight", "dflash.dspark.markov.w1", + "dflash.dspark.markov.w2", "dflash.dspark.confidence.weight", + "dflash.dspark.confidence.bias", + }; + for (const char * expected : names) { + if (std::strcmp(name, expected) == 0) return true; + } + int il = -1; + std::string suffix; + return block_suffix(name, il, suffix) && il >= 0 && il < n_layer && + recognized_block_suffix(suffix); +} + +bool tensor_shape_is(const ggml_tensor * tensor, + std::initializer_list dims) { + if (!tensor || dims.size() > GGML_MAX_DIMS) return false; + size_t i = 0; + for (int64_t dim : dims) { + if (tensor->ne[i++] != dim) return false; + } + for (; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] != 1) return false; + } + return true; +} + } // namespace const char * deepseek4_dspark_last_error() { return g_dspark_err.c_str(); } @@ -84,6 +135,7 @@ const char * deepseek4_dspark_last_error() { return g_dspark_err.c_str(); } bool load_deepseek4_dspark_drafter(const std::string & path, ggml_backend_t backend, DSparkDrafter & out) { + g_dspark_err.clear(); ggml_context * meta = nullptr; gguf_init_params gip{}; gip.no_alloc = true; @@ -164,26 +216,110 @@ bool load_deepseek4_dspark_drafter(const std::string & path, out.vocab_size = (int)kv_u32(g, P + "dflash.dspark.vocab_size", w.n_vocab); out.confidence_dim = (int)kv_u32(g, P + "dflash.dspark.confidence_dim", 0); + const bool valid_arch = + w.n_layer > 0 && w.n_layer <= 64 && + w.n_embd > 0 && w.n_embd <= 32768 && + w.n_vocab > 0 && w.n_vocab <= 1000000 && + w.n_head > 0 && w.n_head <= 512 && + w.head_dim > 0 && w.head_dim <= 4096 && + w.n_lora_q > 0 && w.n_lora_q <= 32768 && + w.n_lora_o > 0 && w.n_lora_o <= 32768 && + w.n_out_group > 0 && w.n_out_group <= w.n_head && + w.n_head % w.n_out_group == 0 && + w.n_expert > 0 && w.n_expert <= 4096 && + w.n_expert_used > 0 && w.n_expert_used <= w.n_expert && + w.n_ff_exp > 0 && w.n_ff_exp <= 131072 && + w.n_hc > 0 && w.n_hc <= 64 && + out.n_target_layers > 0 && out.n_target_layers <= 64 && + out.block_size >= 2 && out.block_size <= 16 && + out.markov_rank > 0 && out.markov_rank <= 32768 && + out.vocab_size == w.n_vocab; + if (!valid_arch) { + set_err("invalid DSpark architecture metadata"); + gguf_free(g); + if (meta) ggml_free(meta); + out = DSparkDrafter{}; + return false; + } + w.layers.resize(n_layer); w.backend = backend; - // ── Allocate all tensors from the meta ctx into one backend buffer ── + // Validate the complete tensor table before allocating device memory or + // forming an absolute file offset. The drafter contract intentionally has + // no embedding/lm-head tensors; rejecting unknown tensors prevents a stray + // full-model tensor from consuming the entire GPU allocation. + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + set_err("open failed: " + path); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } + struct stat st{}; + if (::fstat(fd, &st) != 0 || st.st_size < 0) { + set_err("fstat failed: " + path); + ::close(fd); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } + const size_t file_size = (size_t) st.st_size; + const size_t data_off = gguf_get_data_offset(g); + const int64_t n_tensors = gguf_get_n_tensors(g); + bool ok = true; + for (int64_t ti = 0; ti < n_tensors; ++ti) { + const char * tname = gguf_get_tensor_name(g, ti); + ggml_tensor * tensor = ggml_get_tensor(meta, tname); + if (!tensor) { + set_err(std::string("meta tensor missing: ") + tname); + ok = false; + break; + } + if (!recognized_dspark_tensor(tname, w.n_layer)) { + set_err(std::string("unexpected DSpark tensor: ") + tname); + ok = false; + break; + } + const size_t tensor_off = gguf_get_tensor_offset(g, ti); + const size_t size = gguf_get_tensor_size(g, ti); + if (!gguf_tensor_in_file(data_off, tensor_off, size, file_size)) { + set_err(gguf_bounds_error("DSpark draft GGUF", tname, + ggml_type_name(gguf_get_tensor_type(g, ti)), + data_off, tensor_off, size, file_size)); + ok = false; + break; + } + } + if (!ok) { + ::close(fd); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } + + // ── Allocate validated tensors into one backend buffer ────────────── ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(meta, backend); - if (!buf) { set_err("ggml_backend_alloc_ctx_tensors failed"); gguf_free(g); ggml_free(meta); return false; } + if (!buf) { + set_err("ggml_backend_alloc_ctx_tensors failed"); + ::close(fd); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); // ── Stream tensor bytes from the file (pread per tensor) ──────────── - const int fd = ::open(path.c_str(), O_RDONLY); - if (fd < 0) { set_err("open failed: " + path); ggml_backend_buffer_free(buf); gguf_free(g); ggml_free(meta); return false; } - const size_t data_off = gguf_get_data_offset(g); - const int64_t n_tensors = gguf_get_n_tensors(g); std::vector staging; - bool ok = true; for (int64_t ti = 0; ti < n_tensors && ok; ti++) { const char * tname = gguf_get_tensor_name(g, ti); ggml_tensor * t = ggml_get_tensor(meta, tname); - if (!t) { set_err(std::string("meta tensor missing: ") + tname); ok = false; break; } - const size_t off = data_off + gguf_get_tensor_offset(g, ti); + const size_t tensor_off = gguf_get_tensor_offset(g, ti); + const size_t off = data_off + tensor_off; // preflight proved no overflow const size_t size = gguf_get_tensor_size(g, ti); staging.resize(size); size_t done = 0; @@ -196,7 +332,13 @@ bool load_deepseek4_dspark_drafter(const std::string & path, ggml_backend_tensor_set(t, staging.data(), 0, size); } ::close(fd); - if (!ok) { ggml_backend_buffer_free(buf); gguf_free(g); ggml_free(meta); return false; } + if (!ok) { + ggml_backend_buffer_free(buf); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } // ── Bind pointers by name ─────────────────────────────────────────── for (int64_t ti = 0; ti < n_tensors; ti++) { @@ -233,7 +375,6 @@ bool load_deepseek4_dspark_drafter(const std::string & path, else if (suffix == "ffn_norm.weight") L.ffn_norm = t; else if (suffix == "ffn_gate_inp.weight") L.ffn_gate_inp = t; else if (suffix == "exp_probs_b.bias") L.ffn_exp_probs_b = t; - else if (suffix == "ffn_gate_tid2eid.weight") L.ffn_gate_tid2eid = t; else if (suffix == "ffn_gate_exps.weight") L.ffn_gate_exps = t; else if (suffix == "ffn_up_exps.weight") L.ffn_up_exps = t; else if (suffix == "ffn_down_exps.weight") L.ffn_down_exps = t; @@ -288,33 +429,83 @@ bool load_deepseek4_dspark_drafter(const std::string & path, need(L.hc_ffn_base, p + "hc_ffn_base.weight"); } + std::string shape_error; + auto expect_shape = [&](const ggml_tensor * tensor, const std::string & name, + std::initializer_list dims) { + if (tensor && shape_error.empty() && !tensor_shape_is(tensor, dims)) { + shape_error = name + " has a shape incompatible with DSpark metadata"; + } + }; + const int64_t n_embd = w.n_embd; + const int64_t hc_dim = n_embd * w.n_hc; + const int64_t mix_dim = 2LL * w.n_hc + (int64_t) w.n_hc * w.n_hc; + const int64_t q_dim = (int64_t) w.n_head * w.head_dim; + const int64_t group_dim = (int64_t) w.head_dim * (w.n_head / w.n_out_group); + const int64_t out_low_dim = (int64_t) w.n_lora_o * w.n_out_group; + + expect_shape(out.main_proj, "dflash.fc.weight", + {(int64_t) out.n_target_layers * n_embd, n_embd}); + expect_shape(out.main_norm, "dflash.hidden_norm.weight", {n_embd}); + expect_shape(out.markov_w1, "dflash.dspark.markov.w1", + {out.markov_rank, out.vocab_size}); + expect_shape(out.markov_w2, "dflash.dspark.markov.w2", + {out.markov_rank, out.vocab_size}); + expect_shape(w.out_norm, "output_norm.weight", {n_embd}); + expect_shape(w.output_hc_fn, "output_hc_fn.weight", {hc_dim, w.n_hc}); + expect_shape(w.output_hc_scale, "output_hc_scale.weight", {1}); + expect_shape(w.output_hc_base, "output_hc_base.weight", {w.n_hc}); + if (out.confidence_w && out.confidence_dim > 0) { + expect_shape(out.confidence_w, "dflash.dspark.confidence.weight", + {out.confidence_dim, 1}); + expect_shape(out.confidence_b, "dflash.dspark.confidence.bias", {1}); + } + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + const std::string p = "blk." + std::to_string(il) + "."; + expect_shape(L.attn_norm, p + "attn_norm.weight", {n_embd}); + expect_shape(L.attn_q_a, p + "attn_q_a.weight", {n_embd, w.n_lora_q}); + expect_shape(L.attn_q_a_norm, p + "attn_q_a_norm.weight", {w.n_lora_q}); + expect_shape(L.attn_q_b, p + "attn_q_b.weight", {w.n_lora_q, q_dim}); + expect_shape(L.attn_kv, p + "attn_kv.weight", {n_embd, w.head_dim}); + expect_shape(L.attn_kv_a_norm, p + "attn_kv_a_norm.weight", {w.head_dim}); + if (L.attn_sinks) expect_shape(L.attn_sinks, p + "attn_sinks.weight", {w.n_head}); + expect_shape(L.attn_output_a, p + "attn_output_a.weight", + {group_dim, out_low_dim}); + expect_shape(L.attn_output_b, p + "attn_output_b.weight", + {out_low_dim, n_embd}); + expect_shape(L.hc_attn_fn, p + "hc_attn_fn.weight", {hc_dim, mix_dim}); + expect_shape(L.hc_attn_scale, p + "hc_attn_scale.weight", {3}); + expect_shape(L.hc_attn_base, p + "hc_attn_base.weight", {mix_dim}); + expect_shape(L.ffn_norm, p + "ffn_norm.weight", {n_embd}); + expect_shape(L.ffn_gate_inp, p + "ffn_gate_inp.weight", {n_embd, w.n_expert}); + if (L.ffn_exp_probs_b) { + expect_shape(L.ffn_exp_probs_b, p + "exp_probs_b.bias", {w.n_expert}); + } + expect_shape(L.ffn_gate_exps, p + "ffn_gate_exps.weight", + {n_embd, w.n_ff_exp, w.n_expert}); + expect_shape(L.ffn_up_exps, p + "ffn_up_exps.weight", + {n_embd, w.n_ff_exp, w.n_expert}); + expect_shape(L.ffn_down_exps, p + "ffn_down_exps.weight", + {w.n_ff_exp, n_embd, w.n_expert}); + expect_shape(L.ffn_gate_shexp, p + "ffn_gate_shexp.weight", {n_embd, w.n_ff_exp}); + expect_shape(L.ffn_up_shexp, p + "ffn_up_shexp.weight", {n_embd, w.n_ff_exp}); + expect_shape(L.ffn_down_shexp, p + "ffn_down_shexp.weight", {w.n_ff_exp, n_embd}); + expect_shape(L.hc_ffn_fn, p + "hc_ffn_fn.weight", {hc_dim, mix_dim}); + expect_shape(L.hc_ffn_scale, p + "hc_ffn_scale.weight", {3}); + expect_shape(L.hc_ffn_base, p + "hc_ffn_base.weight", {mix_dim}); + } + std::string contract_error; if (!out.dspark_enabled) { contract_error = "dflash.dspark.enabled is false"; - } else if (w.n_layer <= 0 || w.n_embd <= 0 || w.n_vocab <= 0 || - w.n_hc <= 0 || out.n_target_layers <= 0 || - out.block_size < 2 || out.block_size > 16) { - contract_error = "invalid DSpark architecture metadata"; } else if ((int) out.capture_layer_ids.size() != out.n_target_layers) { contract_error = "capture_layer_ids count does not match dflash.n_target_layers"; - } else if (out.main_proj && - (out.main_proj->ne[0] != (int64_t) out.n_target_layers * w.n_embd || - out.main_proj->ne[1] != w.n_embd)) { - contract_error = "dflash.fc.weight shape does not match target-layer metadata"; - } else if (out.markov_w1 && out.markov_w2 && - (out.markov_w1->ne[0] != out.markov_rank || - out.markov_w1->ne[1] != out.vocab_size || - out.markov_w2->ne[0] != out.markov_rank || - out.markov_w2->ne[1] != out.vocab_size)) { - contract_error = "DSpark Markov tensor shape does not match metadata"; } else if ((out.confidence_w == nullptr) != (out.confidence_b == nullptr)) { contract_error = "incomplete DSpark confidence tensor pair"; - } else if (out.confidence_w && - (out.confidence_dim <= 0 || - out.confidence_w->ne[0] != out.confidence_dim || - out.confidence_w->ne[1] != 1 || - ggml_nelements(out.confidence_b) != 1)) { - contract_error = "DSpark confidence tensor shape does not match metadata"; + } else if (out.confidence_w && out.confidence_dim <= 0) { + contract_error = "DSpark confidence metadata is invalid"; + } else if (!shape_error.empty()) { + contract_error = shape_error; } if (!missing.empty() || !contract_error.empty()) { std::string message = "invalid DSpark drafter GGUF"; diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 9cddf3beb..83dc4d056 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -123,7 +123,7 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, std::vector * logits_out, std::vector & capture_out, DeepSeek4StepTelemetry * telemetry = nullptr, - bool allow_graph_reuse = true); + bool allow_graph_reuse = false); // Run DSpark speculative decode: draft block_size candidates with `drafter`, // verify against the DS4 target in one batched forward, accept the matching diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index cc5fa4fa0..ffb099114 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -522,11 +522,7 @@ bool run_deepseek4_dspark_spec_decode( // boundary except at its last token (state rows stay distinct and the // comp emission matches AR). Boundaries sit at p % 4 == 3. The // sequential verify has no such limit. - static const bool fused_verify_mode = [] { - const char * v = std::getenv("DFLASH_DS4_FUSED_VERIFY"); - return v && *v && *v != '0'; - }(); - int q_step_cap = (seq_verify_mode || fused_verify_mode) + int q_step_cap = seq_verify_mode ? std::min(q_cap, 4) : std::min(q_cap, 4 - (pos & 3)); if (adaptive_width && !use_confidence_width && !seq_verify_mode) { @@ -610,7 +606,9 @@ bool run_deepseek4_dspark_spec_decode( int verify_last = -1; if (!target.verify_batch(draft_tok, pos, verify_last, &tgt_am)) { if (full_snap) { - target.restore_kv(); + if (!target.restore_kv()) { + std::fprintf(stderr, "[ds4-spec] restore after verify failure failed\n"); + } } else { spec_rollback_apply(rollback, target_w, target_cache, pos, boundary_crossed); } @@ -651,7 +649,11 @@ bool run_deepseek4_dspark_spec_decode( std::vector kv_toks; kv_toks.push_back(lt); for (int i = 1; i < accept; i++) kv_toks.push_back(draft_tok[i]); - target.restore_kv(); + if (!target.restore_kv()) { + std::fprintf(stderr, "[ds4-spec] snapshot restore failed\n"); + ok = false; + break; + } int replay_last = -1; std::vector replay_am; if (!target.verify_batch(kv_toks, pos, replay_last, &replay_am)) { diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 824d609e0..3aa8d013a 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -123,8 +123,8 @@ static bool ds4_build_fused_verify_graph( hc_cur[(size_t) t] = ggml_repeat_4d(ctx, col, n_embd, n_hc, 1, 1); } - std::vector capture_pieces; // order [ci][t] - capture_pieces.reserve(capture_ids.size() * (size_t) q); + std::vector capture_pieces( + capture_ids.size() * (size_t) q, nullptr); // order [ci][t] for (int il = 0; il < w.n_layer; ++il) { const DeepSeek4Layer & L = w.layers[(size_t) il]; @@ -213,13 +213,14 @@ static bool ds4_build_fused_verify_graph( have_token_ids && hash_tables[(size_t) il].loaded; if (hash_routed) { // ds4_build_hash_routed_ffn is single-token; run per token, concat. - ggml_tensor * hids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, w.n_expert_used, q); + const int n_used = ds4_effective_expert_count(w); + ggml_tensor * hids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_used, q); ggml_set_input(hids); fg.hash_ids[(size_t) il] = hids; for (int t = 0; t < q; ++t) { ggml_tensor * fcol = ggml_view_2d(ctx, ffn_normed, n_embd, 1, ffn_normed->nb[1], (size_t) t * ffn_normed->nb[1]); - ggml_tensor * hcol = ggml_view_2d(ctx, hids, w.n_expert_used, 1, + ggml_tensor * hcol = ggml_view_2d(ctx, hids, n_used, 1, hids->nb[1], (size_t) t * hids->nb[1]); ggml_tensor * fo = ds4_build_hash_routed_ffn(ctx, w, L, ggml_cont(ctx, fcol), hcol); if (!fo) { ffn_out = nullptr; break; } @@ -246,7 +247,8 @@ static bool ds4_build_fused_verify_graph( ggml_tensor * hsT = ggml_cont(ctx, ggml_transpose(ctx, hs)); // [n_hc, n_embd] ggml_tensor * summed = ggml_sum_rows(ctx, hsT); // [1, n_embd] ggml_tensor * mean = ggml_scale(ctx, summed, 1.0f / (float) n_hc); - capture_pieces.push_back(ggml_reshape_1d(ctx, ggml_cont(ctx, mean), n_embd)); + capture_pieces[ci * (size_t) q + (size_t) t] = + ggml_reshape_1d(ctx, ggml_cont(ctx, mean), n_embd); } } } @@ -271,6 +273,12 @@ static bool ds4_build_fused_verify_graph( ggml_build_forward_expand(gf, fg.logits); if (!capture_pieces.empty()) { + for (ggml_tensor * piece : capture_pieces) { + if (!piece) { + std::fprintf(stderr, "[ds4-fused-verify] capture layer id is invalid\n"); + return false; + } + } ggml_tensor * cap = capture_pieces[0]; for (size_t i = 1; i < capture_pieces.size(); ++i) { cap = ggml_concat(ctx, cap, capture_pieces[i], 0); @@ -340,10 +348,12 @@ static int ds4_try_fused_verify_step( const int q = n_tokens; const int token_pos = kv_start + q - 1; std::vector key; - key.reserve((size_t) w.n_layer + 3); + key.reserve((size_t) w.n_layer + hooks->capture_layer_ids->size() + 4); key.push_back(w.n_swa); key.push_back(q); key.push_back(token_ids ? 1 : 0); + key.push_back((int64_t) hooks->capture_layer_ids->size()); + for (int layer : *hooks->capture_layer_ids) key.push_back(layer); for (int il = 0; il < w.n_layer; ++il) { const int ratio = (int) w.compress_ratios[il]; DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; @@ -469,12 +479,20 @@ static int ds4_try_fused_verify_step( for (int il = 0; il < w.n_layer; ++il) { ggml_tensor * hids = fg->hash_ids[(size_t) il]; if (!hids) continue; - const auto & ht = hash_tables[(size_t) il].ids; - const int n_used = w.n_expert_used; + const HashRoutingTableCpu & table = hash_tables[(size_t) il]; + const int n_used = ds4_effective_expert_count(w); hash_scratch.resize((size_t) n_used * q); for (int i = 0; i < q; ++i) { + const int32_t * row = hash_routing_row( + table, token_ids[i], w.n_expert_used); + if (!row) { + std::fprintf(stderr, + "[ds4-fused-verify] token id %d outside hash table for layer %d\n", + token_ids[i], il); + return -1; + } std::memcpy(hash_scratch.data() + (size_t) i * n_used, - ht.data() + (size_t) token_ids[i] * (size_t) n_used, + row, (size_t) n_used * sizeof(int32_t)); } ggml_backend_tensor_set(hids, hash_scratch.data(), 0, diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index e34b4bc61..80e70edac 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -5341,6 +5341,7 @@ struct DsparkDraftCache { int ctx_len = -1; int block = -1; const void * drafter = nullptr; + ggml_backend_t backend = nullptr; std::vector arena; ggml_context * ctx = nullptr; ggml_gallocr_t alloc = nullptr; @@ -5393,7 +5394,7 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, DsparkDraftCache & C = g_dspark_draft_cache; const bool DS4_DBG = std::getenv("DFLASH_DS4_DSPARK_DEBUG") != nullptr; - if (C.drafter != (const void *) &d) { + if (C.drafter != (const void *) &d || C.backend != backend) { // HC scales (host) per layer + output — immutable, read once per drafter. C.s_attn.assign((size_t) w.n_layer, {0.0f, 0.0f, 0.0f}); C.s_ffn.assign((size_t) w.n_layer, {0.0f, 0.0f, 0.0f}); @@ -5406,8 +5407,13 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, C.s_out = so[0]; } - if (!C.ctx || C.ctx_len != ctx_len || C.block != block || C.drafter != (const void *) &d) { + if (!C.ctx || C.ctx_len != ctx_len || C.block != block || + C.drafter != (const void *) &d || C.backend != backend) { // ── (Re)build the graph ───────────────────────────────────────── + if (C.backend != backend && C.alloc) { + ggml_gallocr_free(C.alloc); + C.alloc = nullptr; + } if (C.ctx) { ggml_free(C.ctx); C.ctx = nullptr; } C.gf = nullptr; C.dbg_taps.clear(); @@ -5557,6 +5563,7 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, C.ctx_len = ctx_len; C.block = block; C.drafter = (const void *) &d; + C.backend = backend; } // ── Set inputs + compute (cached graph) ───────────────────────────── diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 7f4bd6084..42259ae36 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -9,6 +9,7 @@ #include "common/backend_ipc.h" #include "common/layer_split_utils.h" +#include "deepseek4/deepseek4_dspark.h" #include #include @@ -198,6 +199,113 @@ static std::string write_deepseek4_tensor_fixture() { return path; } +struct DSparkFixtureOptions { + bool wrong_main_norm_shape = false; + bool add_unknown_tensor = false; +}; + +static std::string write_dspark_loader_fixture(const DSparkFixtureOptions & opts = {}) { + ggml_init_params ip{}; + ip.mem_size = 4u << 20; + ip.mem_buffer = nullptr; + ip.no_alloc = false; + ggml_context * ctx = ggml_init(ip); + gguf_context * g = gguf_init_empty(); + + const char * arch = "deepseek4-dflash-draft"; + const std::string p = std::string(arch) + "."; + gguf_set_val_str(g, "general.architecture", arch); + gguf_set_val_u32(g, (p + "block_count").c_str(), 1); + gguf_set_val_u32(g, (p + "embedding_length").c_str(), 4); + gguf_set_val_u32(g, (p + "vocab_size").c_str(), 8); + gguf_set_val_u32(g, (p + "attention.head_count").c_str(), 1); + gguf_set_val_u32(g, (p + "attention.head_count_kv").c_str(), 1); + gguf_set_val_u32(g, (p + "attention.key_length").c_str(), 4); + gguf_set_val_u32(g, (p + "rope.dimension_count").c_str(), 4); + gguf_set_val_u32(g, (p + "attention.q_lora_rank").c_str(), 4); + gguf_set_val_u32(g, (p + "attention.output_lora_rank").c_str(), 4); + gguf_set_val_u32(g, (p + "attention.output_group_count").c_str(), 1); + gguf_set_val_u32(g, (p + "expert_count").c_str(), 2); + gguf_set_val_u32(g, (p + "expert_used_count").c_str(), 1); + gguf_set_val_u32(g, (p + "expert_shared_count").c_str(), 1); + gguf_set_val_u32(g, (p + "expert_feed_forward_length").c_str(), 8); + gguf_set_val_u32(g, (p + "hash_layer_count").c_str(), 0); + gguf_set_val_u32(g, (p + "attention.sliding_window").c_str(), 8); + gguf_set_val_u32(g, (p + "attention.indexer.head_count").c_str(), 1); + gguf_set_val_u32(g, (p + "attention.indexer.key_length").c_str(), 4); + gguf_set_val_u32(g, (p + "attention.indexer.top_k").c_str(), 1); + gguf_set_val_u32(g, (p + "hyper_connection.count").c_str(), 1); + gguf_set_val_u32(g, (p + "hyper_connection.sinkhorn_iterations").c_str(), 1); + gguf_set_val_u32(g, (p + "dflash.n_target_layers").c_str(), 1); + gguf_set_val_u32(g, (p + "dflash.block_size").c_str(), 2); + gguf_set_val_u32(g, (p + "dflash.mask_token_id").c_str(), 7); + gguf_set_val_u32(g, (p + "dflash.head_hc_enabled").c_str(), 1); + gguf_set_val_u32(g, (p + "dflash.dspark.enabled").c_str(), 1); + gguf_set_val_u32(g, (p + "dflash.dspark.markov_rank").c_str(), 2); + gguf_set_val_u32(g, (p + "dflash.dspark.vocab_size").c_str(), 8); + const int32_t capture_ids[] = {0}; + gguf_set_arr_data(g, (p + "dflash.capture_layer_ids").c_str(), + GGUF_TYPE_INT32, capture_ids, 1); + + auto add1 = [&](const char * name, int64_t n0) { + ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n0); + ggml_set_name(t, name); + std::memset(t->data, 0, ggml_nbytes(t)); + gguf_add_tensor(g, t); + }; + auto add2 = [&](const char * name, int64_t n0, int64_t n1) { + ggml_tensor * t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n0, n1); + ggml_set_name(t, name); + std::memset(t->data, 0, ggml_nbytes(t)); + gguf_add_tensor(g, t); + }; + auto add3 = [&](const char * name, int64_t n0, int64_t n1, int64_t n2) { + ggml_tensor * t = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n0, n1, n2); + ggml_set_name(t, name); + std::memset(t->data, 0, ggml_nbytes(t)); + gguf_add_tensor(g, t); + }; + + add1("output_norm.weight", 4); + add2("output_hc_fn.weight", 4, 1); + add1("output_hc_scale.weight", 1); + add1("output_hc_base.weight", 1); + add2("dflash.fc.weight", 4, 4); + add1("dflash.hidden_norm.weight", opts.wrong_main_norm_shape ? 3 : 4); + add2("dflash.dspark.markov.w1", 2, 8); + add2("dflash.dspark.markov.w2", 2, 8); + + add1("blk.0.attn_norm.weight", 4); + add2("blk.0.attn_q_a.weight", 4, 4); + add1("blk.0.attn_q_a_norm.weight", 4); + add2("blk.0.attn_q_b.weight", 4, 4); + add2("blk.0.attn_kv.weight", 4, 4); + add1("blk.0.attn_kv_a_norm.weight", 4); + add2("blk.0.attn_output_a.weight", 4, 4); + add2("blk.0.attn_output_b.weight", 4, 4); + add2("blk.0.hc_attn_fn.weight", 4, 3); + add1("blk.0.hc_attn_scale.weight", 3); + add1("blk.0.hc_attn_base.weight", 3); + add1("blk.0.ffn_norm.weight", 4); + add2("blk.0.ffn_gate_inp.weight", 4, 2); + add3("blk.0.ffn_gate_exps.weight", 4, 8, 2); + add3("blk.0.ffn_up_exps.weight", 4, 8, 2); + add3("blk.0.ffn_down_exps.weight", 8, 4, 2); + add2("blk.0.ffn_gate_shexp.weight", 4, 8); + add2("blk.0.ffn_up_shexp.weight", 4, 8); + add2("blk.0.ffn_down_shexp.weight", 8, 4); + add2("blk.0.hc_ffn_fn.weight", 4, 3); + add1("blk.0.hc_ffn_scale.weight", 3); + add1("blk.0.hc_ffn_base.weight", 3); + if (opts.add_unknown_tensor) add2("token_embd.weight", 4, 8); + + const std::string path = make_temp_gguf_path("dspark"); + gguf_write_to_file(g, path.c_str(), /*only_meta=*/false); + gguf_free(g); + ggml_free(ctx); + return path; +} + static ggml_context * make_test_context(size_t mem_size = 1u << 20) { ggml_init_params params = {}; params.mem_size = mem_size; @@ -808,6 +916,65 @@ static void test_loader_rejects_truncated_tensor_data(ggml_backend_t backend) { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_dspark_loader_contract_and_bounds(ggml_backend_t backend) { + std::fprintf(stderr, " test_dspark_loader_contract_and_bounds ..."); + + { + const std::string path = write_dspark_loader_fixture(); + DSparkDrafter drafter; + const bool ok = load_deepseek4_dspark_drafter(path, backend, drafter); + TEST_ASSERT_MSG(ok, deepseek4_dspark_last_error()); + free_deepseek4_dspark_drafter(drafter); + unlink(path.c_str()); + } + + { + DSparkFixtureOptions opts; + opts.wrong_main_norm_shape = true; + const std::string path = write_dspark_loader_fixture(opts); + DSparkDrafter drafter; + const bool ok = load_deepseek4_dspark_drafter(path, backend, drafter); + TEST_ASSERT(!ok); + TEST_ASSERT_MSG(std::string(deepseek4_dspark_last_error()).find( + "dflash.hidden_norm.weight") != std::string::npos, + deepseek4_dspark_last_error()); + free_deepseek4_dspark_drafter(drafter); + unlink(path.c_str()); + } + + { + DSparkFixtureOptions opts; + opts.add_unknown_tensor = true; + const std::string path = write_dspark_loader_fixture(opts); + DSparkDrafter drafter; + const bool ok = load_deepseek4_dspark_drafter(path, backend, drafter); + TEST_ASSERT(!ok); + TEST_ASSERT_MSG(std::string(deepseek4_dspark_last_error()).find( + "unexpected DSpark tensor: token_embd.weight") != std::string::npos, + deepseek4_dspark_last_error()); + free_deepseek4_dspark_drafter(drafter); + unlink(path.c_str()); + } + + { + const std::string path = write_dspark_loader_fixture(); + struct stat st{}; + TEST_ASSERT(stat(path.c_str(), &st) == 0); + TEST_ASSERT(st.st_size > 128); + TEST_ASSERT(truncate(path.c_str(), st.st_size - 128) == 0); + DSparkDrafter drafter; + const bool ok = load_deepseek4_dspark_drafter(path, backend, drafter); + TEST_ASSERT(!ok); + TEST_ASSERT_MSG(std::string(deepseek4_dspark_last_error()).find( + "truncated or corrupt") != std::string::npos, + deepseek4_dspark_last_error()); + free_deepseek4_dspark_drafter(drafter); + unlink(path.c_str()); + } + + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_snapshot_save_restore() { std::fprintf(stderr, " test_snapshot_save_restore ..."); @@ -1680,6 +1847,7 @@ int main() { test_loader_rejects_zero_vocab_size(backend); test_loader_reads_tokenizer_special_ids(backend); test_loader_rejects_truncated_tensor_data(backend); + test_dspark_loader_contract_and_bounds(backend); test_snapshot_save_restore(); test_reset_request_state(); test_reset_deepseek4_cache(backend); From 8d8f5694419a0f9807af9a93a32d174fd2c3d8c5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:35:45 -0700 Subject: [PATCH 09/12] fix(ds4): harden compressor and drafter lifecycle --- server/src/deepseek4/deepseek4_backend.cpp | 179 +++++++++++------- server/src/deepseek4/deepseek4_backend.h | 5 + server/src/deepseek4/deepseek4_dspark.h | 4 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 37 +++- .../src/deepseek4/deepseek4_fused_verify.inc | 12 +- server/src/deepseek4/deepseek4_graph.cpp | 98 +++++++++- server/src/deepseek4/deepseek4_internal.h | 7 + server/tests/test_deepseek4_unit.cpp | 45 +++++ 8 files changed, 308 insertions(+), 79 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 9ba2f12bf..0e2c13e67 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -397,6 +397,59 @@ bool DeepSeek4Backend::load_model() { return true; } +bool DeepSeek4Backend::load_spec_drafter() { + if (spec_draft_path_.empty()) return true; + if (parked_ || moe_hybrid_) { + std::fprintf(stderr, + "[deepseek4] cannot load DSpark drafter without a resident " + "monolithic target\n"); + return false; + } + + auto drafter = std::make_unique(); + if (!load_deepseek4_dspark_drafter(spec_draft_path_, backend_, *drafter)) { + std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", + deepseek4_dspark_last_error()); + return false; + } + + const DSparkDrafter & d = *drafter; + bool compatible = d.core.n_embd == w_.n_embd && + d.core.n_vocab == w_.n_vocab && + d.vocab_size == w_.n_vocab && + d.mask_token_id >= 0 && d.mask_token_id < w_.n_vocab && + (int) d.capture_layer_ids.size() == d.n_target_layers; + for (int layer : d.capture_layer_ids) { + compatible = compatible && layer >= 0 && layer < w_.n_layer; + } + if (!compatible) { + std::fprintf(stderr, + "[deepseek4] DSpark drafter is incompatible with target " + "(target embd/vocab/layers=%d/%d/%d, draft=%d/%d)\n", + w_.n_embd, w_.n_vocab, w_.n_layer, + d.core.n_embd, d.vocab_size); + free_deepseek4_dspark_drafter(*drafter); + return false; + } + + spec_drafter_ = std::move(drafter); + spec_enabled_ = true; + spec_drafter_parked_ = false; + std::fprintf(stderr, "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", + spec_draft_path_.c_str()); + return true; +} + +void DeepSeek4Backend::release_spec_drafter(bool mark_parked) { + if (spec_drafter_) { + free_deepseek4_dspark_drafter(*spec_drafter_); + } + spec_drafter_.reset(); + spec_enabled_ = false; + spec_feat_window_.clear(); + spec_drafter_parked_ = mark_parked && !spec_draft_path_.empty(); +} + bool DeepSeek4Backend::init() { // The shared MMVQ/MMQ crossover defaults to q=3 for NVIDIA. On gfx1151, // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged, @@ -436,42 +489,16 @@ bool DeepSeek4Backend::init() { w_.fused_decode ? "on" : "off", moe_hybrid_ ? " [hybrid]" : ""); - if (env_flag_enabled("DFLASH_DS4_SPEC") && moe_hybrid_) { - std::fprintf(stderr, - "[deepseek4] DSpark spec-decode requires monolithic model placement; " - "disabled for hybrid expert placement\n"); - } else if (env_flag_enabled("DFLASH_DS4_SPEC")) { + if (env_flag_enabled("DFLASH_DS4_SPEC")) { const char * dp = std::getenv("DFLASH_DS4_DRAFT"); if (dp && *dp) { - spec_drafter_ = std::make_unique(); - if (load_deepseek4_dspark_drafter(dp, backend_, *spec_drafter_)) { - const DSparkDrafter & d = *spec_drafter_; - bool compatible = d.core.n_embd == w_.n_embd && - d.core.n_vocab == w_.n_vocab && - d.vocab_size == w_.n_vocab && - d.mask_token_id >= 0 && d.mask_token_id < w_.n_vocab && - (int) d.capture_layer_ids.size() == d.n_target_layers; - for (int layer : d.capture_layer_ids) { - compatible = compatible && layer >= 0 && layer < w_.n_layer; - } - if (!compatible) { - std::fprintf(stderr, - "[deepseek4] DSpark drafter is incompatible with target " - "(target embd/vocab/layers=%d/%d/%d, draft=%d/%d)\n", - w_.n_embd, w_.n_vocab, w_.n_layer, - d.core.n_embd, d.vocab_size); - free_deepseek4_dspark_drafter(*spec_drafter_); - spec_drafter_.reset(); - } else { - spec_enabled_ = true; - std::fprintf(stderr, - "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", - dp); - } + spec_draft_path_ = dp; + if (moe_hybrid_) { + std::fprintf(stderr, + "[deepseek4] DSpark spec-decode requires monolithic model " + "placement; disabled for hybrid expert placement\n"); } else { - std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", - deepseek4_dspark_last_error()); - spec_drafter_.reset(); + (void) load_spec_drafter(); } } else { std::fprintf(stderr, "[deepseek4] DFLASH_DS4_SPEC set but DFLASH_DS4_DRAFT gguf missing\n"); @@ -585,7 +612,14 @@ void DeepSeek4Backend::print_ready_banner() const { } bool DeepSeek4Backend::park(const std::string & what) { + const bool want_draft = (what.empty() || what == "all" || what == "draft"); const bool want_target = (what.empty() || what == "all" || what == "target"); + + if (want_draft && spec_drafter_) { + release_spec_drafter(/*mark_parked=*/true); + std::printf("[deepseek4] DSpark drafter parked (VRAM released)\n"); + std::fflush(stdout); + } if (!want_target || parked_) return true; maybe_save_routing_stats(); @@ -599,47 +633,68 @@ bool DeepSeek4Backend::park(const std::string & what) { moe_placement_ = {}; free_deepseek4_weights(w_); parked_ = true; - std::printf("[deepseek4] parked (VRAM released)\n"); + if (spec_drafter_) { + std::printf("[deepseek4] target parked (target VRAM released; " + "DSpark drafter retained)\n"); + } else { + std::printf("[deepseek4] target parked (target VRAM released)\n"); + } std::fflush(stdout); return true; } bool DeepSeek4Backend::unpark(const std::string & what) { + const bool want_draft = (what.empty() || what == "all" || what == "draft"); const bool want_target = (what.empty() || what == "all" || what == "target"); - if (!want_target || !parked_) return true; - if (!load_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; - } + if (want_target && parked_) { + if (!load_model()) { + std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); + free_deepseek4_weights(w_); + stream_engine_.destroy(); + moe_hybrid_.reset(); + moe_placement_ = {}; + return false; + } - const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { - std::fprintf(stderr, "[deepseek4] unpark: failed to recreate KV cache (ctx=%d)\n", max_ctx); - free_deepseek4_cache(cache_); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; + const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; + if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { + std::fprintf(stderr, + "[deepseek4] unpark: failed to recreate KV cache (ctx=%d)\n", + max_ctx); + free_deepseek4_cache(cache_); + free_deepseek4_weights(w_); + stream_engine_.destroy(); + moe_hybrid_.reset(); + moe_placement_ = {}; + return false; + } + + parked_ = false; + std::printf("[deepseek4] target unparked (VRAM restored)\n"); + std::fflush(stdout); } - parked_ = false; - std::printf("[deepseek4] unparked (VRAM restored)\n"); - std::fflush(stdout); + if (want_draft && spec_drafter_parked_) { + if (parked_) { + std::fprintf(stderr, + "[deepseek4] unpark: restore target before DSpark drafter\n"); + return false; + } + if (!load_spec_drafter()) { + std::fprintf(stderr, "[deepseek4] unpark: failed to restore DSpark drafter\n"); + return false; + } + } return true; } int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset) { - // The current batched graph updates only the chunk's final compressor - // window. Until batched compressor state is sequentially updated, process - // prefill one token at a time to preserve long-prompt correctness. + // Keep server prefill token-at-a-time for established numerics and bounded + // dynamic-graph workspace. The lower-level layer-range API independently + // splits arbitrary multi-token callers at every compressor boundary. constexpr int chunk = 1; const int n_total = (int)tokens.size(); int pos = kv_offset; @@ -948,12 +1003,8 @@ bool DeepSeek4Backend::handle_compress(const std::string & line, } void DeepSeek4Backend::free_drafter() { - if (spec_drafter_) { - free_deepseek4_dspark_drafter(*spec_drafter_); - } - spec_drafter_.reset(); - spec_enabled_ = false; - spec_feat_window_.clear(); + release_spec_drafter(/*mark_parked=*/false); + spec_draft_path_.clear(); } void DeepSeek4Backend::maybe_save_routing_stats() { diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 64ce91464..5c37c77a1 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -79,9 +79,14 @@ class DeepSeek4Backend : public ModelBackend { // DSpark speculative decode (opt-in: DFLASH_DS4_SPEC=1 + DFLASH_DS4_DRAFT=). bool spec_enabled_ = false; + bool spec_drafter_parked_ = false; + std::string spec_draft_path_; std::unique_ptr spec_drafter_; std::vector spec_feat_window_; + bool load_spec_drafter(); + void release_spec_drafter(bool mark_parked); + // Prefill prompt tokens in chunks, return absolute committed position. int do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset = 0); diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 83dc4d056..55fd610bb 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -82,8 +82,8 @@ const char * deepseek4_dspark_last_error(); // features. All host-side f32 for a simple v1 (GPU feature-ring plumbing can // come later). // -// seed_tok : the committed anchor token (block position 0's embedding) -// noise_embed : [n_embd * block_size] embeds of [seed]+[MASK]*(block_size-1) +// noise_embed : [n_embd * block_size] embeds of [seed]+[MASK]*(block_size-1); +// its first block element is the committed seed embedding // ctx_features : [n_target_layers*n_embd * ctx_len] captured features, // ordered oldest..newest, absolute positions // [committed-ctx_len .. committed-1] diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index ffb099114..8d61581c2 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -406,7 +406,9 @@ bool run_deepseek4_dspark_spec_decode( // Fast path caps the verify at the compression ratio (4): one boundary max, // no rolling-state row aliasing -> snapshot-free rollback stays exact. - // The full-snapshot A/B path may use the whole block. + // Full snapshots change rollback strategy but not the compressor-window + // limit below. The legacy sequential measurement path is validated only + // through q=4. int q_cap = full_snap ? block + 1 : 4; if (const char * qs = std::getenv("DFLASH_DS4_SPEC_Q")) { const int v = std::atoi(qs); @@ -422,6 +424,13 @@ bool run_deepseek4_dspark_spec_decode( } std::fclose(qf); } + if (seq_verify_mode && q_cap > 4) { + std::fprintf(stderr, + "[ds4-spec] sequential verify supports q<=4; " + "capping requested q=%d to 4\n", + q_cap); + q_cap = 4; + } // Snapshot backend for the legacy full-snapshot rollback path. ggml_backend_t snap_backend = ggml_backend_cpu_init(); @@ -464,7 +473,7 @@ bool run_deepseek4_dspark_spec_decode( int lt = last_tok; int pos = committed; // absolute position of the seed (block slot 0) int n_generated = 0; - long accept_sum = 0, steps = 0; + long accept_sum = 0, offered_sum = 0, steps = 0; bool ok = true; bool stop_requested = false; @@ -520,8 +529,8 @@ bool run_deepseek4_dspark_spec_decode( bool ds_ok = false; // Batched-verify exactness: the batch must not cross a ratio-4 // boundary except at its last token (state rows stay distinct and the - // comp emission matches AR). Boundaries sit at p % 4 == 3. The - // sequential verify has no such limit. + // comp emission matches AR). Boundaries sit at p % 4 == 3. The legacy + // sequential measurement path is independently capped to q=4 above. int q_step_cap = seq_verify_mode ? std::min(q_cap, 4) : std::min(q_cap, 4 - (pos & 3)); @@ -694,6 +703,7 @@ bool run_deepseek4_dspark_spec_decode( pos = commit_pos; // seed + accepted candidates now in KV lt = bonus; // deferred bonus becomes next seed accept_sum += matched; + offered_sum += q - 1; ewma_accept = 0.7 * ewma_accept + 0.3 * (double) matched; steps++; if (timing && (steps <= 4 || (steps & 31) == 0)) { @@ -708,12 +718,17 @@ bool run_deepseek4_dspark_spec_decode( } const double total_ms = spec_ms_since(run_t0); - if (accept_rate_out && steps > 0) { - const int denom = q_cap > 1 ? q_cap - 1 : 1; - *accept_rate_out = (float) accept_sum / (float) (steps * denom); + if (accept_rate_out) { + *accept_rate_out = offered_sum > 0 + ? (float) accept_sum / (float) offered_sum + : 0.0f; } - std::fprintf(stderr, "[ds4-spec] gen=%d steps=%ld mean_accept=%.2f/%d q_cap=%d full_snap=%d\n", - n_generated, steps, steps ? (double) accept_sum / steps : 0.0, q_cap - 1, q_cap, + std::fprintf(stderr, + "[ds4-spec] gen=%d steps=%ld mean_accept=%.2f/%.2f " + "q_cap=%d full_snap=%d\n", + n_generated, steps, + steps ? (double) accept_sum / steps : 0.0, + steps ? (double) offered_sum / steps : 0.0, q_cap, (int) full_snap); if (steps > 0) { std::fprintf(stderr, @@ -730,7 +745,7 @@ bool run_deepseek4_dspark_spec_decode( "[ds4-spec-t] verify tel/step: hc_pre_a=%.1f attn_b=%.1f attn_c=%.1f attn_r=%.1f " "hc_post_a=%.1f hc_pre_f=%.1f route(b/c/r/s)=%.1f/%.1f/%.1f/%.1f " "ffn(b/c/r)=%.1f/%.1f/%.1f eval=%.1f hot=%.1f cold=%.1f comb=%.1f part=%.1f " - "ghits=%llu gbuilds=%llu ms\n", + "full(b/s/c/r)=%.1f/%.1f/%.1f/%.1f ghits=%llu gbuilds=%llu ms\n", tel.hc_pre_attn_us / s, tel.attn_build_us / s, tel.attn_compute_us / s, tel.attn_read_us / s, tel.hc_post_attn_us / s, tel.hc_pre_ffn_us / s, tel.route_build_us / s, tel.route_compute_us / s, tel.route_read_us / s, @@ -738,6 +753,8 @@ bool run_deepseek4_dspark_spec_decode( tel.ffn_build_us / s, tel.ffn_compute_us / s, tel.ffn_read_us / s, tel.ffn_eval_us / s, tel.ffn_hot_us / s, tel.ffn_cold_us / s, tel.ffn_combine_us / s, tel.ffn_partition_us / s, + tel.full_graph_build_us / s, tel.full_graph_set_us / s, + tel.full_graph_compute_us / s, tel.full_graph_read_us / s, (unsigned long long) tel.ffn_hot_graph_hits, (unsigned long long) tel.ffn_hot_graph_builds); } diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 3aa8d013a..83fda9bdf 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -402,6 +402,7 @@ static int ds4_try_fused_verify_step( fg->last_use = vc.counter; // ── Fill inputs ── + const auto set_t0 = Ds4TimingClock::now(); ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); std::vector iv(q); std::vector lv(q); @@ -499,6 +500,9 @@ static int ds4_try_fused_verify_step( sizeof(int32_t) * (size_t) n_used * q); } } + if (telemetry) { + telemetry->full_graph_set_us += ds4_elapsed_us(set_t0, Ds4TimingClock::now()); + } // ── Compute + read back ── const auto compute_t0 = Ds4TimingClock::now(); @@ -506,8 +510,11 @@ static int ds4_try_fused_verify_step( std::fprintf(stderr, "[ds4-fused-verify] compute failed\n"); return -1; } - if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + } + const auto read_t0 = Ds4TimingClock::now(); const int ncap = (int) hooks->capture_layer_ids->size(); if (hooks->all_logits_out) { hooks->all_logits_out->resize((size_t) w.n_vocab * q); @@ -531,5 +538,8 @@ static int ds4_try_fused_verify_step( } } } + if (telemetry) { + telemetry->full_graph_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); + } return 1; } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 80e70edac..1c0769ac5 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -591,6 +591,21 @@ static ggml_tensor * build_tail_rope_2d(ggml_context * ctx, // ─── KV Compressor Step ──────────────────────────────────────────────── +int deepseek4_safe_compressor_batch_tokens(const DeepSeek4Weights & w, + int kv_start, + int n_tokens) { + if (n_tokens <= 0) return 0; + int safe = n_tokens; + for (uint32_t raw_ratio : w.compress_ratios) { + const int ratio = (int) raw_ratio; + if (ratio <= 0) continue; + int pos_mod = kv_start % ratio; + if (pos_mod < 0) pos_mod += ratio; + safe = std::min(safe, ratio - pos_mod); + } + return std::max(1, safe); +} + static void build_compressor_step( ggml_context * ctx, ggml_cgraph * gf, @@ -641,8 +656,9 @@ static void build_compressor_step( ggml_tensor * comp_cache_source = comp_cache; // Causal-batch verify: every token's contribution lands in its - // position-addressed state row (the rows are distinct because the batch - // never crosses a ratio window; the boundary may only be the last token). + // position-addressed state row. deepseek4_step_layer_range splits dynamic + // batches at compressor boundaries, so the boundary can only be the last + // token and no post-boundary write can precede pooling/rotation. const bool batched_state = (cur_all != nullptr && n_tokens_all > 1 && !state_rows_inp && kv_start_all >= 0); if (batched_state) { @@ -4158,6 +4174,84 @@ bool deepseek4_step_layer_range( const int hc_dim = n_hc * n_embd; const bool is_last_shard = (layer_end >= w.n_layer); + // A dynamic batch may be supplied by callers other than the DSpark + // verifier. Split it whenever it spans a learned-compressor boundary: + // each sub-forward then writes at most one window and, if present, its + // boundary is the final token. This preserves the same pool/rotate order + // as sequential execution while retaining safe batched prefixes. + const int first_chunk = deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens); + if (first_chunk > 0 && first_chunk < n_tokens) { + const int input_width = layer_begin == 0 ? n_embd : hc_dim; + std::vector hc_all; + std::vector shard_out_all; + std::vector capture_all; + std::vector logits_all; + std::vector last_out; + hc_all.reserve((size_t) hc_dim * n_tokens); + if (out_logits && !is_last_shard) { + shard_out_all.reserve((size_t) hc_dim * n_tokens); + } + if (verify_hooks && verify_hooks->capture_out && + verify_hooks->capture_layer_ids) { + capture_all.reserve((size_t) verify_hooks->capture_layer_ids->size() * + n_embd * n_tokens); + } + if (verify_hooks && verify_hooks->all_logits_out) { + logits_all.reserve((size_t) w.n_vocab * n_tokens); + } + + for (int off = 0; off < n_tokens;) { + const int chunk = deepseek4_safe_compressor_batch_tokens( + w, kv_start + off, n_tokens - off); + std::vector chunk_hc; + std::vector chunk_out; + std::vector chunk_capture; + std::vector chunk_logits; + Ds4VerifyHooks chunk_hooks; + Ds4VerifyHooks * chunk_hooks_ptr = nullptr; + if (verify_hooks) { + chunk_hooks.capture_layer_ids = verify_hooks->capture_layer_ids; + chunk_hooks.capture_out = verify_hooks->capture_out ? &chunk_capture : nullptr; + chunk_hooks.all_logits_out = verify_hooks->all_logits_out ? &chunk_logits : nullptr; + chunk_hooks_ptr = &chunk_hooks; + } + if (!deepseek4_step_layer_range( + backend, w, cache, chunk_hc, + embed + (size_t) off * input_width, + chunk, kv_start + off, layer_begin, layer_end, + out_logits ? &chunk_out : nullptr, + token_ids ? token_ids + off : nullptr, + telemetry, allow_decode_graph_reuse, chunk_hooks_ptr)) { + return false; + } + hc_all.insert(hc_all.end(), chunk_hc.begin(), chunk_hc.end()); + if (out_logits) { + if (is_last_shard) { + last_out = std::move(chunk_out); + } else { + shard_out_all.insert(shard_out_all.end(), + chunk_out.begin(), chunk_out.end()); + } + } + capture_all.insert(capture_all.end(), + chunk_capture.begin(), chunk_capture.end()); + logits_all.insert(logits_all.end(), chunk_logits.begin(), chunk_logits.end()); + off += chunk; + } + + hc_state = std::move(hc_all); + if (out_logits) { + *out_logits = is_last_shard ? std::move(last_out) : std::move(shard_out_all); + } + if (verify_hooks && verify_hooks->capture_out) { + *verify_hooks->capture_out = std::move(capture_all); + } + if (verify_hooks && verify_hooks->all_logits_out) { + *verify_hooks->all_logits_out = std::move(logits_all); + } + return true; + } + // Initialize HC state. // First shard (layer_begin=0): embed is token embeddings [n_embd × n_tokens], // replicate into n_hc streams. diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index b52ca002b..a09cd1f62 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -321,6 +321,13 @@ bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, DeepSeek4Cache & cache); +// Largest prefix of [kv_start, kv_start + n_tokens) that reaches at most the +// next learned-compressor boundary. Multi-token dynamic forwards split on +// this prefix so state writes after a boundary cannot race ahead of pooling. +int deepseek4_safe_compressor_batch_tokens(const DeepSeek4Weights & w, + int kv_start, + int n_tokens); + // Forward: single step (prefill chunk or decode token). // embed: [n_embd, n_tokens] input embeddings (post-embedding lookup). // hc_state: [n_hc * n_embd] persistent HC residual (updated in-place). diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 42259ae36..5ff9be0a3 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -29,6 +29,7 @@ #include #define private public +#include "deepseek4/deepseek4_backend.h" #include "deepseek4/deepseek4_layer_split_adapter.h" #undef private @@ -975,6 +976,48 @@ static void test_dspark_loader_contract_and_bounds(ggml_backend_t backend) { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_safe_compressor_batch_tokens() { + std::fprintf(stderr, " test_safe_compressor_batch_tokens ..."); + DeepSeek4Weights w; + w.compress_ratios = {0, 4, 128}; + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 0, 1) == 1); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 0, 4) == 4); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 0, 5) == 4); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 1, 8) == 3); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 4, 8) == 4); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 125, 8) == 3); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 128, 8) == 4); + + w.compress_ratios = {128}; + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 0, 129) == 128); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 127, 4) == 1); + + w.compress_ratios.clear(); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 17, 9) == 9); + TEST_ASSERT(deepseek4_safe_compressor_batch_tokens(w, 17, 0) == 0); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_dspark_park_all_releases_drafter() { + std::fprintf(stderr, " test_dspark_park_all_releases_drafter ..."); + + DeepSeek4BackendConfig cfg; + DeepSeek4Backend backend(cfg); + backend.spec_draft_path_ = "/tmp/ds4-dspark-fixture.gguf"; + backend.spec_drafter_ = std::make_unique(); + backend.spec_enabled_ = true; + + TEST_ASSERT(backend.park("all")); + TEST_ASSERT(backend.parked_); + TEST_ASSERT(backend.spec_drafter_ == nullptr); + TEST_ASSERT(!backend.spec_enabled_); + TEST_ASSERT(backend.spec_drafter_parked_); + + // Avoid leaving a synthetic reload request for the destructor. + backend.free_drafter(); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_snapshot_save_restore() { std::fprintf(stderr, " test_snapshot_save_restore ..."); @@ -1848,6 +1891,8 @@ int main() { test_loader_reads_tokenizer_special_ids(backend); test_loader_rejects_truncated_tensor_data(backend); test_dspark_loader_contract_and_bounds(backend); + test_safe_compressor_batch_tokens(); + test_dspark_park_all_releases_drafter(); test_snapshot_save_restore(); test_reset_request_state(); test_reset_deepseek4_cache(backend); From 61d2c91d4af758522247af4160ece5e2a3fc1d63 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:54:44 -0700 Subject: [PATCH 10/12] fix(ds4): preserve callback and draft reload contracts --- server/src/deepseek4/deepseek4_backend.cpp | 24 ++++++++++++---------- server/tests/test_deepseek4_unit.cpp | 5 ++++- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 0e2c13e67..33cc81541 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -876,17 +876,18 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, const DaemonIO & io) { GenerateResult result; + DaemonIO out_io = io.with_token_callback(req.on_token); auto t0 = Clock::now(); // Prefill - int committed = do_prefill(req.prompt, io); + int committed = do_prefill(req.prompt, out_io); if (committed < 0) { result.fail(GenerateErrorCode::PrefillFailed); return result; } result.prefill_s = elapsed_s(t0); - if (io.cancelled) { + if (out_io.cancelled) { result.succeed(); maybe_save_routing_stats(); return result; @@ -912,10 +913,10 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, for (int i = 1; i < w_.n_vocab; i++) if (last_logits_[i] > mv) { mv = last_logits_[i]; seed = i; } } std::vector gen; gen.push_back(seed); - io.emit(seed); + out_io.emit(seed); float accept_rate = 0.0f; bool spec_ran = false; - if (!io.cancelled && !deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { + if (!out_io.cancelled && !deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; std::vector spec_toks; @@ -925,10 +926,10 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, req.n_gen - 1, win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, spec_toks, &accept_rate, - [&io](int32_t tok) { - if (io.cancelled) return false; - io.emit(tok); - return !io.cancelled; + [&out_io](int32_t tok) { + if (out_io.cancelled) return false; + out_io.emit(tok); + return !out_io.cancelled; })) { result.fail(GenerateErrorCode::DecodeFailed, "DSpark speculative decode failed"); @@ -951,7 +952,7 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, gen_tokens.reserve(req.n_gen); bool forced_close = false; - if (!do_decode(committed, req.n_gen, gen_tokens, io, + if (!do_decode(committed, req.n_gen, gen_tokens, out_io, req.budget_hook, &forced_close)) { result.fail(GenerateErrorCode::DecodeFailed); return result; @@ -1003,8 +1004,9 @@ bool DeepSeek4Backend::handle_compress(const std::string & line, } void DeepSeek4Backend::free_drafter() { - release_spec_drafter(/*mark_parked=*/false); - spec_draft_path_.clear(); + // Keep the configured path so request-scoped residency and an explicit + // later `unpark draft` can restore the DSpark model. + release_spec_drafter(/*mark_parked=*/true); } void DeepSeek4Backend::maybe_save_routing_stats() { diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 5ff9be0a3..218148762 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1013,8 +1013,11 @@ static void test_dspark_park_all_releases_drafter() { TEST_ASSERT(!backend.spec_enabled_); TEST_ASSERT(backend.spec_drafter_parked_); - // Avoid leaving a synthetic reload request for the destructor. backend.free_drafter(); + TEST_ASSERT(backend.spec_drafter_ == nullptr); + TEST_ASSERT(!backend.spec_enabled_); + TEST_ASSERT(backend.spec_drafter_parked_); + TEST_ASSERT(backend.spec_draft_path_ == "/tmp/ds4-dspark-fixture.gguf"); std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } From 4766f1fe5c9526d3d3d212c72acc6dea864923b1 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:10:44 -0700 Subject: [PATCH 11/12] fix(ds4): restore rejected SWA rows after wrap --- server/docs/DS4.md | 8 ++ server/src/deepseek4/deepseek4_dspark.h | 29 ++++++ .../src/deepseek4/deepseek4_dspark_spec.cpp | 93 +++++++++++++------ server/tests/test_deepseek4_unit.cpp | 71 ++++++++++++++ 4 files changed, 175 insertions(+), 26 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index fcd230f8f..aae2b62b1 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -183,6 +183,14 @@ export DFLASH_DS4_SPEC_Q=4 --ds4-expert-top-k 4 ``` +`DFLASH_DS4_FUSED_VERIFY=1` is the opt-in throughput profile. Its persistent +whole-model GPU graph uses stable padded reduction shapes, so near-tied greedy +logits can select a different token than the normal causal verifier even at +temperature 0. Leave it unset when comparing against the normal verifier, or +set `DFLASH_DS4_SEQ_VERIFY=1` for the slower token-at-a-time verification +diagnostic. Neither fused verification nor the separate +`--ds4-expert-top-k 4` approximation should be presented as byte-identical AR. + DSpark currently requires monolithic target placement. On HIP, `--ds4-fused-decode` selects that placement; if the target falls back to hybrid expert placement, the server logs that DSpark is disabled and continues with diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 55fd610bb..a45bea280 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -29,6 +29,8 @@ #include "ggml.h" #include "ggml-backend.h" +#include +#include #include #include #include @@ -125,6 +127,33 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, DeepSeek4StepTelemetry * telemetry = nullptr, bool allow_graph_reuse = false); +// Minimal speculative-decode rollback state. Rejected positions must restore +// the physical SWA rows they overwrote after the ring wraps; otherwise a later +// causal verify reads rejected-token KV as if it were older committed history. +// This remains much smaller than a full target-cache snapshot because q <= 4. +struct DeepSeek4SpecRollback { + int raw_pos = 0; + int raw_count = 0; + struct Layer { + std::vector attn_kv, attn_sc, idx_kv, idx_sc; + std::size_t raw_row_bytes = 0; + std::vector raw_rows; + }; + std::vector layers; + std::vector hc_state; +}; + +void deepseek4_spec_rollback_save(const DeepSeek4Cache & cache, + DeepSeek4SpecRollback & rollback, + int raw_pos, + int raw_count); + +void deepseek4_spec_rollback_apply(const DeepSeek4SpecRollback & rollback, + const DeepSeek4Weights & weights, + DeepSeek4Cache & cache, + int commit_pos, + bool restore_prev); + // Run DSpark speculative decode: draft block_size candidates with `drafter`, // verify against the DS4 target in one batched forward, accept the matching // prefix, and loop. Returns generated tokens via `io.emit`. Mirrors the laguna diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 8d61581c2..2b69509a5 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -9,11 +9,12 @@ // Fast spec loop (default): ONE batched verify per step, verify width capped // at DS4_SPEC_Q=4 tokens (seed + 3 candidates). With q <= ratio(4) the verify // crosses at most one compression boundary and never aliases rolling-state -// rows, so rejection rollback needs no KV snapshot at all: -// - raw ring rows are position-addressed (pos % n_swa) -> idempotent, +// rows, so rejection rollback needs no full KV snapshot: +// - at-risk raw ring rows are saved before the verify and rejected rows are +// restored after wrap (accepted rows remain committed), // - comp rows are index-addressed (pos / ratio) -> idempotent, // - n_comp / n_index_comp are pure functions of commit position, -// - the ONLY non-idempotent state is the ratio-4 compressor prev-half +// - the other non-idempotent state is the ratio-4 compressor prev-half // (4 rows/state, flushed cur->prev at chunk boundaries), a few KB/layer, // saved host-side before the verify and restored only when the flush // happened at-or-past the rollback point. @@ -236,19 +237,6 @@ constexpr float kConfidenceQ3Threshold = 0.30f; constexpr float kConfidenceQ4Threshold = 0.25f; // ── Light rollback state ──────────────────────────────────────────────── -// Only the non-position-idempotent cache pieces: the prev-half rows of every -// ratio-4 rolling compressor state (attn + indexer) plus hc_state. Everything -// else (raw ring, comp rows, counters) is position/index-addressed and either -// overwritten idempotently or recomputable from the commit position. -struct SpecRollback { - int cur_pos = 0; - struct Layer { - std::vector attn_kv, attn_sc, idx_kv, idx_sc; - }; - std::vector layers; - std::vector hc_state; -}; - // prev-half = first 4 rows of a [comp_width, 8] ratio-4 rolling state. // ratio-128 states ([comp_width, 128]) are pure position rings -> skip. void save_prev_half(ggml_tensor * t, std::vector & buf) { @@ -263,16 +251,33 @@ void restore_prev_half(ggml_tensor * t, const std::vector & buf) { ggml_backend_tensor_set(t, buf.data(), 0, buf.size()); } -void spec_rollback_save(const DeepSeek4Cache & cache, SpecRollback & rb) { - rb.cur_pos = cache.cur_pos; +void spec_rollback_save_impl(const DeepSeek4Cache & cache, + DeepSeek4SpecRollback & rb, + int raw_pos, int raw_count) { + rb.raw_pos = raw_pos; + rb.raw_count = std::max(0, raw_count); rb.layers.resize(cache.layers.size()); for (size_t il = 0; il < cache.layers.size(); ++il) { const DeepSeek4LayerCache & lc = cache.layers[il]; - SpecRollback::Layer & s = rb.layers[il]; + DeepSeek4SpecRollback::Layer & s = rb.layers[il]; save_prev_half(lc.attn_compressor.state_kv, s.attn_kv); save_prev_half(lc.attn_compressor.state_score, s.attn_sc); save_prev_half(lc.indexer_compressor.state_kv, s.idx_kv); save_prev_half(lc.indexer_compressor.state_score, s.idx_sc); + s.raw_row_bytes = 0; + s.raw_rows.clear(); + if (lc.raw_kv && lc.raw_kv->ne[1] > 0 && rb.raw_count > 0) { + s.raw_row_bytes = ggml_row_size(lc.raw_kv->type, lc.raw_kv->ne[0]); + s.raw_rows.resize(s.raw_row_bytes * (size_t) rb.raw_count); + for (int t = 0; t < rb.raw_count; ++t) { + int row = (rb.raw_pos + t) % (int) lc.raw_kv->ne[1]; + if (row < 0) row += (int) lc.raw_kv->ne[1]; + ggml_backend_tensor_get( + lc.raw_kv, + s.raw_rows.data() + (size_t) t * s.raw_row_bytes, + (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); + } + } } if (cache.hc_state) { const size_t bytes = ggml_nbytes(cache.hc_state); @@ -286,8 +291,11 @@ void spec_rollback_save(const DeepSeek4Cache & cache, SpecRollback & rb) { // prev-half rows with a chunk containing rejected tokens, so put the // pre-verify rows back. (A boundary strictly inside the committed range is a // legitimate flush and must be kept.) -void spec_rollback_apply(const SpecRollback & rb, const DeepSeek4Weights & w, - DeepSeek4Cache & cache, int commit_pos, bool restore_prev) { +void spec_rollback_apply_impl(const DeepSeek4SpecRollback & rb, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + int commit_pos, + bool restore_prev) { cache.cur_pos = commit_pos; for (size_t il = 0; il < cache.layers.size(); ++il) { DeepSeek4LayerCache & lc = cache.layers[il]; @@ -295,12 +303,28 @@ void spec_rollback_apply(const SpecRollback & rb, const DeepSeek4Weights & w, if (ratio > 0) lc.n_comp = commit_pos / (int) ratio; if (ratio == 4) lc.n_index_comp = commit_pos / 4; if (restore_prev && il < rb.layers.size()) { - const SpecRollback::Layer & s = rb.layers[il]; + const DeepSeek4SpecRollback::Layer & s = rb.layers[il]; restore_prev_half(lc.attn_compressor.state_kv, s.attn_kv); restore_prev_half(lc.attn_compressor.state_score, s.attn_sc); restore_prev_half(lc.indexer_compressor.state_kv, s.idx_kv); restore_prev_half(lc.indexer_compressor.state_score, s.idx_sc); } + if (il < rb.layers.size() && lc.raw_kv && lc.raw_kv->ne[1] > 0) { + const DeepSeek4SpecRollback::Layer & s = rb.layers[il]; + const int first_rejected = std::max(0, commit_pos - rb.raw_pos); + for (int t = first_rejected; + t < rb.raw_count && + s.raw_row_bytes > 0 && + s.raw_rows.size() >= (size_t) (t + 1) * s.raw_row_bytes; + ++t) { + int row = (rb.raw_pos + t) % (int) lc.raw_kv->ne[1]; + if (row < 0) row += (int) lc.raw_kv->ne[1]; + ggml_backend_tensor_set( + lc.raw_kv, + s.raw_rows.data() + (size_t) t * s.raw_row_bytes, + (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); + } + } } if (restore_prev && cache.hc_state && !rb.hc_state.empty()) { ggml_backend_tensor_set(cache.hc_state, rb.hc_state.data(), 0, rb.hc_state.size()); @@ -315,6 +339,21 @@ double spec_ms_since(SpecClock::time_point t0) { } // namespace +void deepseek4_spec_rollback_save(const DeepSeek4Cache & cache, + DeepSeek4SpecRollback & rollback, + int raw_pos, + int raw_count) { + spec_rollback_save_impl(cache, rollback, raw_pos, raw_count); +} + +void deepseek4_spec_rollback_apply(const DeepSeek4SpecRollback & rollback, + const DeepSeek4Weights & weights, + DeepSeek4Cache & cache, + int commit_pos, + bool restore_prev) { + spec_rollback_apply_impl(rollback, weights, cache, commit_pos, restore_prev); +} + // Batched target verify + capture: wraps the existing multi-token // deepseek4_step_layer_range (dynamic attention + batched HC), which never // touches the fused single-token 23 tok/s path, with the Ds4VerifyHooks that @@ -439,7 +478,7 @@ bool run_deepseek4_dspark_spec_decode( DeepSeek4DFlashTarget target(target_w, target_cache, backend, snap_backend, drafter.capture_layer_ids, drafter.mask_token_id); DraftWeights dw = make_dspark_shim(drafter); - SpecRollback rollback; + DeepSeek4SpecRollback rollback; DeepSeek4StepTelemetry tel{}; if (timing) target.set_telemetry(&tel); @@ -602,7 +641,7 @@ bool run_deepseek4_dspark_spec_decode( break; } } else { - spec_rollback_save(target_cache, rollback); + deepseek4_spec_rollback_save(target_cache, rollback, pos, q); } tm_save += spec_ms_since(t0); @@ -619,7 +658,8 @@ bool run_deepseek4_dspark_spec_decode( std::fprintf(stderr, "[ds4-spec] restore after verify failure failed\n"); } } else { - spec_rollback_apply(rollback, target_w, target_cache, pos, boundary_crossed); + deepseek4_spec_rollback_apply( + rollback, target_w, target_cache, pos, boundary_crossed); } std::fprintf(stderr, "[ds4-spec] verify failed\n"); ok = false; @@ -674,7 +714,8 @@ bool run_deepseek4_dspark_spec_decode( // The prev-half flush is bad only if the boundary sits at-or-past // the commit point (its chunk then contains rejected tokens). const bool restore_prev = boundary_crossed && first_boundary >= commit_pos; - spec_rollback_apply(rollback, target_w, target_cache, commit_pos, restore_prev); + deepseek4_spec_rollback_apply( + rollback, target_w, target_cache, commit_pos, restore_prev); } // accept == q on the fast path: cur_pos/n_comp already exact, keep. tm_apply += spec_ms_since(t0); diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 218148762..d29cc9a5f 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1021,6 +1021,76 @@ static void test_dspark_park_all_releases_drafter() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_dspark_raw_ring_rollback_after_wrap(ggml_backend_t backend) { + std::fprintf(stderr, " test_dspark_raw_ring_rollback_after_wrap ..."); + + DeepSeek4Weights weights; + weights.n_layer = 1; + weights.n_embd = 4; + weights.n_hc = 1; + weights.head_dim = 4; + weights.n_swa = 8; + weights.n_indexer_head_dim = 2; + weights.compress_ratios = {4}; + + DeepSeek4Cache cache; + TEST_ASSERT(create_deepseek4_cache(backend, weights, 16, cache)); + if (!cache.buf || cache.layers.empty() || !cache.layers[0].raw_kv) { + free_deepseek4_cache(cache); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); + return; + } + + auto & layer = cache.layers[0]; + write_tensor_pattern(layer.raw_kv, 11); + const std::vector original = read_tensor_bytes(layer.raw_kv); + std::vector expected = original; + const size_t row_bytes = ggml_row_size(layer.raw_kv->type, layer.raw_kv->ne[0]); + + cache.cur_pos = 10; + layer.n_comp = 2; + layer.n_index_comp = 2; + DeepSeek4SpecRollback rollback; + deepseek4_spec_rollback_save(cache, rollback, 10, 4); + + auto overwrite_row = [&](int absolute_pos, uint8_t value) { + const int row = absolute_pos % weights.n_swa; + std::vector bytes(row_bytes, value); + ggml_backend_tensor_set(layer.raw_kv, bytes.data(), + (size_t) row * layer.raw_kv->nb[1], row_bytes); + }; + auto expect_overwritten_row = [&](int absolute_pos, uint8_t value) { + const int row = absolute_pos % weights.n_swa; + std::fill(expected.begin() + (size_t) row * layer.raw_kv->nb[1], + expected.begin() + (size_t) row * layer.raw_kv->nb[1] + row_bytes, + value); + }; + for (int t = 0; t < 4; ++t) { + overwrite_row(10 + t, (uint8_t) (0xa0 + t)); + } + + // Commit positions 10 and 11. Their new rows stay in the ring, while the + // rejected positions 12 and 13 must reveal the older history they replaced. + expect_overwritten_row(10, 0xa0); + expect_overwritten_row(11, 0xa1); + deepseek4_spec_rollback_apply(rollback, weights, cache, 12, false); + TEST_ASSERT(read_tensor_bytes(layer.raw_kv) == expected); + TEST_ASSERT(cache.cur_pos == 12); + TEST_ASSERT(layer.n_comp == 3); + TEST_ASSERT(layer.n_index_comp == 3); + + // Restoring to the pre-verify position is used by replay/diagnostic paths + // and must put every overwritten physical row back. + deepseek4_spec_rollback_apply(rollback, weights, cache, 10, false); + TEST_ASSERT(read_tensor_bytes(layer.raw_kv) == original); + TEST_ASSERT(cache.cur_pos == 10); + TEST_ASSERT(layer.n_comp == 2); + TEST_ASSERT(layer.n_index_comp == 2); + + free_deepseek4_cache(cache); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_snapshot_save_restore() { std::fprintf(stderr, " test_snapshot_save_restore ..."); @@ -1896,6 +1966,7 @@ int main() { test_dspark_loader_contract_and_bounds(backend); test_safe_compressor_batch_tokens(); test_dspark_park_all_releases_drafter(); + test_dspark_raw_ring_rollback_after_wrap(backend); test_snapshot_save_restore(); test_reset_request_state(); test_reset_deepseek4_cache(backend); From 27b93d1af145c0507da70e978971d6dd5d4e25e5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:14:23 -0700 Subject: [PATCH 12/12] fix(ds4): feed confidence head pre-norm hidden --- server/src/common/dspark_head.cpp | 29 ++++-- server/src/common/dspark_head.h | 8 +- server/src/deepseek4/deepseek4_dspark.h | 10 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 30 ++++-- server/src/deepseek4/deepseek4_graph.cpp | 19 +++- server/test/test_ds4_dspark_load.cpp | 12 ++- server/tests/test_deepseek4_unit.cpp | 92 +++++++++++++++++++ 7 files changed, 174 insertions(+), 26 deletions(-) diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index 2e1cca95d..f0df52c10 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -178,6 +178,7 @@ struct MarkovChainGraph { ggml_context * ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_confidence_hidden = nullptr; ggml_tensor * inp_seed = nullptr; ggml_tensor * base = nullptr; // [vocab, n_positions] std::vector toks; // corrected argmax per depth @@ -226,6 +227,11 @@ bool build_markov_chain_graph(const DraftWeights & dw, const int vocab = (int)lm_head->ne[1]; const int n_corr = n_positions - first_corrected; if (n_positions <= 0 || n_corr <= 0) return false; + const bool have_confidence = confidence_are_outputs && + dw.dspark.confidence_w != nullptr && + dw.dspark.confidence_b != nullptr && + (dw.dspark.confidence_dim == hdim || + dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); const size_t arena_size = ggml_tensor_overhead() * (size_t)(64 + 16 * n_corr) + ggml_graph_overhead_custom(512, false) + 2 * 1024 * 1024; @@ -240,6 +246,11 @@ bool build_markov_chain_graph(const DraftWeights & dw, out.gf = ggml_new_graph_custom(out.ctx, 512, false); out.inp_hidden = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); + if (have_confidence) { + out.inp_confidence_hidden = + ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); + ggml_set_input(out.inp_confidence_hidden); + } out.inp_seed = ggml_new_tensor_1d(out.ctx, GGML_TYPE_I32, 1); ggml_set_input(out.inp_hidden); ggml_set_input(out.inp_seed); @@ -255,11 +266,6 @@ bool build_markov_chain_graph(const DraftWeights & dw, out.toks.assign((size_t)n_corr, nullptr); out.corrected.assign((size_t)n_corr, nullptr); out.confidence.assign((size_t)n_corr, nullptr); - const bool have_confidence = confidence_are_outputs && - dw.dspark.confidence_w != nullptr && - dw.dspark.confidence_b != nullptr && - (dw.dspark.confidence_dim == hdim || - dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); for (int i = 0; i < n_corr; ++i) { const int row = first_corrected + i; ggml_tensor * prev_emb = ggml_get_rows(out.ctx, dw.dspark.markov_w1, prev_ids); @@ -278,8 +284,9 @@ bool build_markov_chain_graph(const DraftWeights & dw, out.toks[(size_t)i] = tok; if (have_confidence) { ggml_tensor * hidden_i = ggml_view_2d( - out.ctx, out.inp_hidden, hdim, 1, out.inp_hidden->nb[1], - (size_t)row * out.inp_hidden->nb[1]); + out.ctx, out.inp_confidence_hidden, hdim, 1, + out.inp_confidence_hidden->nb[1], + (size_t)row * out.inp_confidence_hidden->nb[1]); ggml_tensor * conf_in = hidden_i; if (dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank) { conf_in = ggml_concat(out.ctx, hidden_i, prev_emb, 0); @@ -307,7 +314,8 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, int q_len, int32_t last_tok, std::vector & draft_tok, - std::vector * confidence_out) { + std::vector * confidence_out, + const float * confidence_hidden) { if (q_len <= 1) return false; if (!dspark_fused_usable(dw, backend, lm_head, local_hidden, "dspark_fused")) return false; const int hdim = dw.n_embd; @@ -337,6 +345,11 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, // Candidate hidden states start at position 1 (position 0 is the seed). ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0, sizeof(float) * (size_t)hdim * (size_t)n_cand); + if (want_confidence && g.inp_confidence_hidden) { + const float * conf_src = confidence_hidden ? confidence_hidden : local_hidden; + ggml_backend_tensor_set(g.inp_confidence_hidden, conf_src + (size_t)hdim, 0, + sizeof(float) * (size_t)hdim * (size_t)n_cand); + } ggml_backend_tensor_set(g.inp_seed, &last_tok, 0, sizeof(int32_t)); if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index ad778e897..9b97b261d 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -22,7 +22,10 @@ bool dspark_markov_correct_greedy_chain(const DraftWeights & dw, // step's get_rows, all in ONE graph on the draft backend. No host logits // round-trip. When confidence_out is non-null and the checkpoint has a // compatible confidence head, returns one score per candidate from the same -// graph and host synchronization as the token ids. +// graph and host synchronization as the token ids. `confidence_hidden`, when +// non-null, has the same padded layout as `local_hidden` and supplies the +// pre-output-norm state expected by the confidence head. Callers without a +// separate state retain the legacy behavior by leaving it null. bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, @@ -30,7 +33,8 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, int q_len, int32_t last_tok, std::vector & draft_tok, - std::vector * confidence_out = nullptr); + std::vector * confidence_out = nullptr, + const float * confidence_hidden = nullptr); // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index a45bea280..adc08167c 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -81,8 +81,9 @@ const char * deepseek4_dspark_last_error(); // One drafter forward. Produces block_size normed hidden states (the input to // the tied lm_head + Markov head), conditioned on a window of captured target -// features. All host-side f32 for a simple v1 (GPU feature-ring plumbing can -// come later). +// features. When requested, also returns the HC-collapsed state before the +// output RMSNorm, which is the trained input to the DSpark confidence head. +// All host-side f32 for a simple v1 (GPU feature-ring plumbing can come later). // // noise_embed : [n_embd * block_size] embeds of [seed]+[MASK]*(block_size-1); // its first block element is the committed seed embedding @@ -92,6 +93,8 @@ const char * deepseek4_dspark_last_error(); // ctx_len : number of context feature columns (<= n_swa) // committed : absolute position of the seed (block position 0) // out_hidden : filled with [n_embd * block_size] = out_norm(hc_head(block)) +// confidence_hidden: optional [n_embd * block_size] = hc_head(block), before +// out_norm; never use it for the tied lm_head // // Defined in deepseek4_graph.cpp (needs the static DS4 sub-builders). bool deepseek4_dspark_draft_forward(ggml_backend_t backend, @@ -100,7 +103,8 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, const float * ctx_features, int ctx_len, int committed, - std::vector & out_hidden); + std::vector & out_hidden, + std::vector * confidence_hidden = nullptr); // Batched target verify forward WITH feature capture (defined in // deepseek4_graph.cpp so it can reuse the target sub-builders). Runs the DS4 diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 2b69509a5..d22bf0ac9 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -233,8 +233,8 @@ bool spec_env_flag(const char * name) { // Calibrated cumulative-confidence thresholds for widening the DS4 verify. // They are part of the policy, not deployment knobs: artifacts without a // compatible confidence head transparently retain the existing EWMA policy. -constexpr float kConfidenceQ3Threshold = 0.30f; -constexpr float kConfidenceQ4Threshold = 0.25f; +constexpr float kConfidenceQ3Threshold = 0.40f; +constexpr float kConfidenceQ4Threshold = 0.30f; // ── Light rollback state ──────────────────────────────────────────────── // prev-half = first 4 rows of a [comp_width, 8] ratio-4 rolling state. @@ -518,7 +518,9 @@ bool run_deepseek4_dspark_spec_decode( std::vector noise_embed((size_t) n_embd * block); std::vector noise_ids(block); - std::vector local_hidden, padded_hidden((size_t) n_embd * (block + 1), 0.0f); + std::vector local_hidden, confidence_hidden; + std::vector padded_hidden((size_t) n_embd * (block + 1), 0.0f); + std::vector padded_confidence_hidden((size_t) n_embd * (block + 1), 0.0f); std::vector draft_tok, tgt_am; std::vector draft_confidence; @@ -539,10 +541,12 @@ bool run_deepseek4_dspark_spec_decode( break; } - // Drafter forward -> block normed hidden states. + // Drafter forward -> normalized states for token logits plus the + // pre-output-norm states expected by the confidence head. if (!deepseek4_dspark_draft_forward(backend, drafter, noise_embed.data(), ctx_len > 0 ? feat_win.data() : nullptr, - ctx_len, pos, local_hidden)) { + ctx_len, pos, local_hidden, + use_confidence_width ? &confidence_hidden : nullptr)) { std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); ok = false; break; @@ -580,10 +584,16 @@ bool run_deepseek4_dspark_spec_decode( if (q_step_cap >= 2) { std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), sizeof(float) * (size_t) n_embd * block); + if (use_confidence_width) { + std::memcpy(padded_confidence_hidden.data() + n_embd, + confidence_hidden.data(), + sizeof(float) * (size_t) n_embd * block); + } ds_ok = dspark_markov_correct_greedy_chain_fused( dw, backend, target.lm_head_tensor(), padded_hidden.data(), q_step_cap, lt, draft_tok, - use_confidence_width ? &draft_confidence : nullptr); + use_confidence_width ? &draft_confidence : nullptr, + use_confidence_width ? padded_confidence_hidden.data() : nullptr); if (!ds_ok) { ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, padded_hidden.data(), q_step_cap, lt, 0.0f, draft_tok); @@ -608,11 +618,11 @@ bool run_deepseek4_dspark_spec_decode( // traces and keep q=4 for high-confidence prefixes while avoiding its // extra verify cost on low-acceptance prompts. if (use_confidence_width && draft_confidence.size() >= 2 && draft_tok.size() >= 3) { - const float p2 = draft_confidence[0] * draft_confidence[1]; - int selected_q = p2 >= kConfidenceQ3Threshold ? 3 : 2; + const float confidence_p2 = draft_confidence[0] * draft_confidence[1]; + int selected_q = confidence_p2 >= kConfidenceQ3Threshold ? 3 : 2; if (selected_q == 3 && draft_confidence.size() >= 3 && draft_tok.size() >= 4) { - const float p3 = p2 * draft_confidence[2]; - if (p3 >= kConfidenceQ4Threshold) selected_q = 4; + const float confidence_p3 = confidence_p2 * draft_confidence[2]; + if (confidence_p3 >= kConfidenceQ4Threshold) selected_q = 4; } if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); } else if (use_confidence_width && !seq_verify_mode) { diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 1c0769ac5..3ef4a25c9 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -5446,6 +5446,7 @@ struct DsparkDraftCache { ggml_tensor * neg_block = nullptr; ggml_tensor * pos_ctx = nullptr; ggml_tensor * out = nullptr; + ggml_tensor * confidence_out = nullptr; std::vector> dbg_taps; // HC scales are immutable weights: read from the backend once. std::vector> s_attn, s_ffn; @@ -5475,7 +5476,8 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, const float * ctx_features, int ctx_len, int committed, - std::vector & out_hidden) { + std::vector & out_hidden, + std::vector * confidence_hidden) { const DeepSeek4Weights & w = d.core; const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -5633,7 +5635,11 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, } // ── Tail: hc_head collapse -> out_norm, per block position ────── + // The tied lm_head consumes the normalized state. The confidence head + // was trained on the HC-collapsed state before this output RMSNorm, so + // keep both instead of reusing the normalized state for confidence. ggml_tensor * out = nullptr; + ggml_tensor * confidence_out = nullptr; for (int p = 0; p < block; p++) { ggml_tensor * hcf = hc_col(hc_cur, p); ggml_tensor * onorm = ggml_rms_norm(ctx, hcf, hc_eps); @@ -5644,10 +5650,16 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, 1); ggml_tensor * hidden_p = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); out = out ? ggml_concat(ctx, out, hidden_p, 1) : hidden_p; + confidence_out = confidence_out + ? ggml_concat(ctx, confidence_out, final_2d, 1) + : final_2d; } ggml_set_output(out); + ggml_set_output(confidence_out); ggml_build_forward_expand(gf, out); + ggml_build_forward_expand(gf, confidence_out); C.out = out; + C.confidence_out = confidence_out; if (!C.alloc) C.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); if (!C.alloc || !ggml_gallocr_alloc_graph(C.alloc, gf)) { @@ -5681,6 +5693,11 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, } out_hidden.resize((size_t) n_embd * block); ggml_backend_tensor_get(C.out, out_hidden.data(), 0, sizeof(float) * out_hidden.size()); + if (confidence_hidden) { + confidence_hidden->resize((size_t) n_embd * block); + ggml_backend_tensor_get(C.confidence_out, confidence_hidden->data(), 0, + sizeof(float) * confidence_hidden->size()); + } if (DS4_DBG) { for (auto & tp : C.dbg_taps) { diff --git a/server/test/test_ds4_dspark_load.cpp b/server/test/test_ds4_dspark_load.cpp index 5340e4bb4..bd1d2b5ec 100644 --- a/server/test/test_ds4_dspark_load.cpp +++ b/server/test/test_ds4_dspark_load.cpp @@ -124,11 +124,12 @@ int main(int argc, char ** argv) { std::vector feats((size_t) fc_in * ctx_len); for (size_t i = 0; i < noise.size(); i++) noise[i] = 0.06f * std::sin(0.31f * (float) i + 1.3f); for (size_t i = 0; i < feats.size(); i++) feats[i] = 0.05f * std::cos(0.17f * (float) i + 0.4f); - std::vector hidden; + std::vector hidden, confidence_hidden; const int committed = 64; // arbitrary >= ctx_len std::fprintf(stderr, "\n── drafter forward (ctx_len=%d block=%d) ──\n", ctx_len, block); if (!deepseek4_dspark_draft_forward(backend, d, noise.data(), feats.data(), - ctx_len, committed, hidden)) { + ctx_len, committed, hidden, + &confidence_hidden)) { std::fprintf(stderr, "FAIL: drafter forward returned false\n"); rc = 1; } else { @@ -142,6 +143,13 @@ int main(int argc, char ** argv) { hidden.size() > 2 ? hidden[2] : 0.0f); need(finite, "finite forward output"); need(hidden.size() == (size_t) n_embd * block, "forward output size"); + bool confidence_finite = true; + for (float v : confidence_hidden) { + if (!std::isfinite(v)) confidence_finite = false; + } + need(confidence_finite, "finite pre-norm confidence hidden"); + need(confidence_hidden.size() == (size_t) n_embd * block, + "pre-norm confidence hidden size"); } } diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index d29cc9a5f..ec7c33bd5 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -8,6 +8,7 @@ #endif #include "common/backend_ipc.h" +#include "common/dspark_head.h" #include "common/layer_split_utils.h" #include "deepseek4/deepseek4_dspark.h" @@ -315,6 +316,96 @@ static ggml_context * make_test_context(size_t mem_size = 1u << 20) { return ggml_init(params); } +static void test_dspark_confidence_uses_separate_hidden(ggml_backend_t backend) { + std::fprintf(stderr, " test_dspark_confidence_uses_separate_hidden ..."); + + constexpr int hidden = 2; + constexpr int rank = 1; + constexpr int vocab = 3; + constexpr int q_len = 2; // dummy seed row + one candidate row + + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * lm_head = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, vocab); + ggml_tensor * markov_w1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rank, vocab); + ggml_tensor * markov_w2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rank, vocab); + ggml_tensor * confidence_w = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden + rank, 1); + ggml_tensor * confidence_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + TEST_ASSERT_MSG(buf != nullptr, "weight allocation failed"); + if (!buf) { + ggml_free(ctx); + std::fprintf(stderr, " FAIL\n"); + return; + } + + const std::vector zeros_lm((size_t) hidden * vocab, 0.0f); + const std::vector zeros_markov((size_t) rank * vocab, 0.0f); + const std::vector confidence_weight = {1.0f, 0.0f, 0.0f}; + const float zero = 0.0f; + ggml_backend_tensor_set(lm_head, zeros_lm.data(), 0, + zeros_lm.size() * sizeof(float)); + ggml_backend_tensor_set(markov_w1, zeros_markov.data(), 0, + zeros_markov.size() * sizeof(float)); + ggml_backend_tensor_set(markov_w2, zeros_markov.data(), 0, + zeros_markov.size() * sizeof(float)); + ggml_backend_tensor_set(confidence_w, confidence_weight.data(), 0, + confidence_weight.size() * sizeof(float)); + ggml_backend_tensor_set(confidence_b, &zero, 0, sizeof(zero)); + + DraftWeights dw{}; + dw.n_embd = hidden; + dw.dspark.enabled = true; + dw.dspark.markov_rank = rank; + dw.dspark.vocab_size = vocab; + dw.dspark.confidence_dim = hidden + rank; + dw.dspark.markov_w1 = markov_w1; + dw.dspark.markov_w2 = markov_w2; + dw.dspark.confidence_w = confidence_w; + dw.dspark.confidence_b = confidence_b; + + // Both calls use the same normalized candidate hidden (all zero), so token + // logits and Markov correction are identical. Only the reference-faithful + // confidence input differs: its candidate row starts with 2.0. + const std::vector normalized_hidden((size_t) hidden * q_len, 0.0f); + const std::vector confidence_hidden = {0.0f, 0.0f, 2.0f, -3.0f}; + std::vector separate_tokens; + std::vector separate_confidence; + const bool separate_ok = dspark_markov_correct_greedy_chain_fused( + dw, backend, lm_head, normalized_hidden.data(), q_len, 0, + separate_tokens, &separate_confidence, confidence_hidden.data()); + + std::vector legacy_tokens; + std::vector legacy_confidence; + const bool legacy_ok = dspark_markov_correct_greedy_chain_fused( + dw, backend, lm_head, normalized_hidden.data(), q_len, 0, + legacy_tokens, &legacy_confidence); + + TEST_ASSERT_MSG(separate_ok, "separate-confidence fused graph failed"); + TEST_ASSERT_MSG(legacy_ok, "legacy-confidence fused graph failed"); + TEST_ASSERT(separate_tokens == legacy_tokens); + TEST_ASSERT(separate_confidence.size() == 1); + TEST_ASSERT(legacy_confidence.size() == 1); + if (separate_confidence.size() == 1 && legacy_confidence.size() == 1) { + const float expected_separate = 1.0f / (1.0f + std::exp(-2.0f)); + TEST_ASSERT_MSG(nearly_equal(separate_confidence[0], expected_separate), + "confidence did not use the separate pre-norm hidden"); + TEST_ASSERT_MSG(nearly_equal(legacy_confidence[0], 0.5f), + "legacy confidence fallback changed"); + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static float softplus_stable(float x) { if (x > 20.0f) { return x; @@ -1964,6 +2055,7 @@ int main() { test_loader_reads_tokenizer_special_ids(backend); test_loader_rejects_truncated_tensor_data(backend); test_dspark_loader_contract_and_bounds(backend); + test_dspark_confidence_uses_separate_hidden(backend); test_safe_compressor_batch_tokens(); test_dspark_park_all_releases_drafter(); test_dspark_raw_ring_rollback_after_wrap(backend);