diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 805d91c09..687938351 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 @@ -801,6 +803,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/docs/DS4.md b/server/docs/DS4.md index c3e32012f..b4c259d8e 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -110,6 +110,62 @@ 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 Speculative Decode + +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: + +| 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 +``` + +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 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(): 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_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 5a3edbf9b..4afe7731c 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -356,6 +356,23 @@ bool DeepSeek4Backend::init() { std::fprintf(stderr, "[deepseek4] initialized: %d layers, ctx=%d, %d experts (%d used)%s\n", w_.n_layer, max_ctx, w_.n_expert, w_.n_expert_used, moe_hybrid_ ? " [hybrid]" : ""); + + 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; } @@ -556,6 +573,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, ggml_backend_buffer_clear(cache_.buf, 0); } 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; @@ -586,11 +604,27 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, routing_stats_.get()); } else { std::vector hc_state; + 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; + } ok = deepseek4_step_layer_range(backend_, w_, cache_, hc_state, embed.data(), n_tok, pos, 0, w_.n_layer, &logits, tokens.data() + i, - timing ? &step_tel : nullptr); + timing ? &step_tel : nullptr, + /*allow_decode_graph_reuse=*/true, hp); + if (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); @@ -724,6 +758,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); @@ -780,7 +843,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 508637867..0641d9231 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..48527e391 --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -0,0 +1,722 @@ +// 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'; +} + +// 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 +// 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); + } + 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, + // 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; + 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; + 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(); + 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 + // 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 && !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, + 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); + } + 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 + } + // 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); + + 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 3f81ba9aa..0eeea9a0d 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -220,6 +220,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, @@ -676,7 +681,10 @@ static void build_compressor_step( std::vector & i64_array_inputs, std::vector & i32_array_inputs, ggml_tensor ** comp_cache_source_out = nullptr, - ggml_tensor * flush_rows_inp = 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; @@ -695,18 +703,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); @@ -722,7 +801,11 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_cur, sc_slot)); } - if (!flush_rows_inp && ((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. @@ -814,6 +897,39 @@ static void build_compressor_step( *comp_cache_source_out = comp_cache_source; } + 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 @@ -869,7 +985,10 @@ static void build_indexer_compressor_step( ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, std::vector & i32_array_inputs, - ggml_tensor * flush_rows_inp = nullptr) { + 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, @@ -894,7 +1013,10 @@ static void build_indexer_compressor_step( i64_array_inputs, i32_array_inputs, nullptr, - flush_rows_inp); + 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) { @@ -1016,7 +1138,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; @@ -1076,6 +1199,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; @@ -1127,7 +1295,10 @@ static ggml_tensor * build_mla_attention( i64_array_inputs, i32_array_inputs, &comp_kv_source, - cached_inputs ? cached_inputs->flush_rows : nullptr); + cached_inputs ? cached_inputs->flush_rows : nullptr, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start); } if (ratio == 4 && L.indexer_compressor_kv) { @@ -1138,7 +1309,10 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, - cached_inputs ? cached_inputs->flush_rows : nullptr); + 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); } @@ -1147,14 +1321,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 @@ -1207,6 +1381,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 @@ -1216,9 +1393,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; @@ -2438,6 +2653,17 @@ static void cpu_matvec_f16_pooled(float * out, const uint16_t * mat, const float 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]; @@ -2672,7 +2898,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) { @@ -2735,7 +2961,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, @@ -2757,7 +2983,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); @@ -4030,6 +4256,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( @@ -4240,7 +4468,8 @@ bool deepseek4_step_layer_range( std::vector * out_logits, const int32_t * token_ids, DeepSeek4StepTelemetry * telemetry, - bool allow_decode_graph_reuse) { + bool allow_decode_graph_reuse, + Ds4VerifyHooks * verify_hooks) { const auto step_t0 = Ds4TimingClock::now(); // NOTE: The old deepseek4_step() lacks HC implementation. @@ -4341,6 +4570,27 @@ bool deepseek4_step_layer_range( static thread_local DeepSeek4LayerRangeScratch scratch; scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used); + if (n_tokens >= 2 && n_tokens <= 4 && verify_hooks && layer_begin == 0 && is_last_shard && + out_logits && ds4_backend_is_gpu(backend) && 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) && ds4_fused_decode_enabled()) { @@ -4605,6 +4855,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); @@ -4612,7 +4863,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); @@ -4640,6 +4891,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(); @@ -4848,6 +5101,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; + } + } + } + } } } } @@ -4926,6 +5198,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()); @@ -5253,3 +5530,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 a5d28d23c..66d71c471 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -323,6 +323,16 @@ bool deepseek4_step( DeepSeek4StepTelemetry * telemetry = nullptr, MoeHybridRoutingStats * routing_stats = 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, const DeepSeek4Weights & w, @@ -336,7 +346,8 @@ bool deepseek4_step_layer_range( std::vector * out_logits, const int32_t * token_ids = nullptr, DeepSeek4StepTelemetry * telemetry = nullptr, - bool allow_decode_graph_reuse = true); + 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; +}