From 7cca859592c831d7e06c95c7419c0cbb85374df0 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Sun, 21 Jun 2026 01:30:55 +0200 Subject: [PATCH 1/4] feat(dflash): config-driven drafter converter + spec-decode perf/correctness fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DRAFTER CONVERTER (config-driven): - convert_dflash_to_gguf.py reads all architecture params from config.json (hidden_size, n_layer, mask_token_id, target_layer_ids, layer_types for SWA, sliding_window). No hardcoded constants. - quantize_draft_q8.py shares load_arch with the converter. - GGUF metadata: dflash.mask_token_id, dflash.target_layer_ids[], dflash.block_size, attention.sliding_window + pattern. - draft_gguf_loader.cpp: read_draft_capture_config(), mask from GGUF metadata, block_size override, SWA pattern from metadata. - draft_safetensors_loader.cpp: dynamic layer count, SWA+mask from config.json. - gguf_target_loader.cpp: respect drafter-specified capture layers instead of overwriting with evenly-spaced heuristic. - qwen35_backend.cpp: early-read capture sync + mask token propagation. - internal.h: capture_layer_ids[16], DFLASH_MAX_CAPTURE_LAYERS=16. - dflash27b.h: DFLASH_MAX_CAPTURE_LAYERS=16. SPEC-DECODE PERFORMANCE: - graph_builders.cpp: build_lm_head_projection_step skips rebuild when ctx alive + n_tokens matches (centralized guard; was per-call-site). - qwen35_backend.cpp: do_spec_decode uses member draft_sg_ (not local) for graph persistence; kFastRollbackThreshold env-tunable (DFLASH_FAST_ROLLBACK_MIN, default 5). - dflash_draft_graph.cpp: exact-ctx_len non-view reuse guard (DFLASH_DRAFT_GRAPH_REUSE, default ON). 4MB ctx alloc (was 256MB). - graph_builders.cpp: 4MB ctx alloc (was 64MB). - step_graph.h: graph_ctx_len + graph_used_view tracking fields. SPEC-DECODE CORRECTNESS: - qwen35_target_graph.cpp: DFLASH_FEAT_RING_CAP env overrides the hardcoded 4096 feature ring cap. Default 4096 causes acceptance collapse from 85% to 7.7% EXACTLY at 4096 prompt tokens (ring wrap corrupts features). - qwen35_backend.cpp: mirror init honors DFLASH_FEAT_RING_CAP. - qwen35_dflash_target.cpp: guard against invalid token IDs from GPU argmax at long context (NaN/Inf → clamp to 0, verify rejects gracefully). MOE EXPERIMENTAL (behind flags): - qwen35moe_backend.cpp: DFLASH_MOE_ALLHOT_HYBRID=1 builds moe_hybrid storage even with 0 cold experts to enable pipelined spec-decode verify. - Persistent moe_hybrid_logits_sg_ graph (was 64MB per-token alloc in hybrid_forward_one_token). GPU argmax (4 bytes vs 1MB vocab readback). - Batched verify/replay via hybrid_forward_batch (was 8 sequential forwards). VALIDATED: - 27B dense + reconverted drafter: 57% accept on code gen, 85% on short prompts. block=16 gives 252 tok/s (2.2x AR) on code generation. - 35B-A3B MoE + reconverted new drafter: 86% accept, 245 tok/s (2.1x AR). - Feature ring cap=16384: 85% holds to 5K tokens, 58% to 10K. - Full pFlash + dFlash stack: goldgate agentic trace passes (100% tool calls valid), pFlash cuts 34K prefill from 475s to 208s (2.3x). - repo_inspection prompt: correct answers, spec at 33.8% accept, 34 tok/s. --- server/include/dflash27b.h | 4 + server/scripts/convert_dflash_to_gguf.py | 76 ++++-- server/scripts/quantize_draft_q8.py | 129 +++------ server/src/common/dflash_draft_graph.cpp | 40 +-- server/src/common/step_graph.h | 10 + server/src/draft/draft_gguf_loader.cpp | 84 +++++- server/src/draft/draft_graph.cpp | 19 +- server/src/draft/draft_safetensors_loader.cpp | 98 ++++++- server/src/internal.h | 13 +- server/src/qwen35/gguf_target_loader.cpp | 37 ++- server/src/qwen35/graph_builders.cpp | 15 +- server/src/qwen35/qwen35_backend.cpp | 123 ++++++--- server/src/qwen35/qwen35_dflash_target.cpp | 9 + server/src/qwen35/qwen35_target_graph.cpp | 10 +- server/src/qwen35moe/qwen35moe_backend.cpp | 251 ++++++++++-------- server/src/qwen35moe/qwen35moe_backend.h | 8 + .../qwen35moe/qwen35moe_pipelined_decode.cpp | 16 +- .../qwen35moe/qwen35moe_pipelined_decode.h | 4 + 18 files changed, 651 insertions(+), 295 deletions(-) diff --git a/server/include/dflash27b.h b/server/include/dflash27b.h index b707b2d8d..abb09e4fd 100644 --- a/server/include/dflash27b.h +++ b/server/include/dflash27b.h @@ -36,6 +36,10 @@ extern "C" { #define DFLASH27B_DRAFT_BLOCK_SIZE 16 #define DFLASH27B_DRAFT_N_TARGET_LAYERS 5 // fc projects 5*hidden -> hidden #define DFLASH27B_DRAFT_MASK_TOKEN_ID 248070 +// Maximum capture layers the target-side array can hold. Newer DFlash +// drafters (Qwen3.6) use 8 taps; this ceiling is generous enough for any +// plausible drafter without wasting stack space. +#define DFLASH_MAX_CAPTURE_LAYERS 16 // target_layer_ids = {1, 16, 31, 46, 61} (0-indexed into target layers) // We capture the OUTPUT of each, which is HF hidden_states[lid + 1]. diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index 186d843d9..38b9c4a2f 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -74,16 +74,28 @@ def load_arch(safetensors: Path, header: dict) -> dict: head_dim=HEAD_DIM, intermediate=INTERMEDIATE, vocab=VOCAB, n_target_layers=N_TARGET_LAYERS, rope_theta=ROPE_THETA, rms_eps=RMS_EPS, mask_token_id=MASK_TOKEN_ID, block_size=BLOCK_SIZE, - ctx_len=CTX_LEN) + ctx_len=CTX_LEN, + sliding_window=0, sliding_window_pattern=[]) cfg_path = safetensors.parent / "config.json" if cfg_path.exists(): c = json.loads(cfg_path.read_text()) + dc = c.get("dflash_config", {}) # nested DFlash config (Qwen3.6 repos) def pick(*keys): for k in keys: if k in c and c[k] is not None: return c[k] return None + # DFlash-specific keys live under dflash_config in newer repos; + # older repos put them at the top level. Check nested first. + def pick_dflash(*keys): + for k in keys: + if k in dc and dc[k] is not None: + return dc[k] + for k in keys: + if k in c and c[k] is not None: + return c[k] + return None for dst, val in ( ("hidden", pick("hidden_size")), ("n_layer", pick("num_hidden_layers")), @@ -94,18 +106,28 @@ def pick(*keys): ("vocab", pick("vocab_size")), ("rope_theta", pick("rope_theta")), ("rms_eps", pick("rms_norm_eps")), - ("n_target_layers", pick("n_target_layers", "num_target_layers")), - ("mask_token_id", pick("mask_token_id")), - ("block_size", pick("block_size", "draft_block_size")), + ("n_target_layers", pick_dflash("n_target_layers", "num_target_layers")), + ("mask_token_id", pick_dflash("mask_token_id")), + ("block_size", pick_dflash("block_size", "draft_block_size")), ("ctx_len", pick("max_position_embeddings")), ): if val is not None: a[dst] = val - dfc = c.get("dflash_config") or {} - _tli = (dfc.get("target_layer_ids") or c.get("target_layer_ids") + # target_layer_ids: exact target layers to capture hidden states from. + # Older repos: top-level target_layer_ids / aux_hidden_state_layer_ids. + # Newer repos: nested under dflash_config. pick_dflash checks both. + _tli = (pick_dflash("target_layer_ids") or c.get("aux_hidden_state_layer_ids")) if _tli: a["capture_layer_ids"] = [int(x) for x in _tli] + # Sliding-window attention (Qwen3.6 pattern: 5 sliding + 1 full). + # layer_types is a list of "sliding_attention" / "full_attention". + sw = pick("sliding_window") + if sw is not None: + a["sliding_window"] = sw + layer_types = pick("layer_types") + if layer_types: + a["sliding_window_pattern"] = [lt == "sliding_attention" for lt in layer_types] print(f"[info] read arch from {cfg_path}") else: print(f"[warn] no config.json next to safetensors; using 27B defaults") @@ -146,10 +168,12 @@ def shape_of(st_name): print(f"[error] config n_head_kv*head_dim={exp_kv} != " f"k_proj.weight dim {k0[0]}; fix config.json", file=sys.stderr) sys.exit(1) + swa_n = sum(1 for x in a.get("sliding_window_pattern", []) if x) print(f"[info] arch: hidden={a['hidden']} n_layer={a['n_layer']} " f"n_head={a['n_head']} n_head_kv={a['n_head_kv']} " f"head_dim={a['head_dim']} ff={a['intermediate']} vocab={a['vocab']} " - f"n_target_layers={a['n_target_layers']}") + f"n_target_layers={a['n_target_layers']} " + f"swa={swa_n}/{a['n_layer']} window={a.get('sliding_window', 0)}") return a @@ -203,19 +227,13 @@ def read_tensor_bytes(path: Path, header_size: int, info: dict) -> bytes: def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: if dtype == "BF16": - # Convert BF16 -> F16 on the host. Several ggml-cuda ops (mul, - # binbcast) only accept F32 / F16 inputs, and llama.cpp's - # build_norm path multiplies normalised activations by the norm - # weight tensor. Storing the draft as F16 throughout sidesteps - # the unsupported BF16 path entirely. Quality impact ~0 for - # weight tensors (BF16 -> F16 keeps 10/8 mantissa bits anyway - # after the implicit cast). + # Keep BF16 raw bytes — do NOT narrow to F16. The sm86 RTX 3090 has + # native BF16 tensor core support, and ggml_mul_mat handles BF16 + # weights directly. Narrowing to F16 loses mantissa precision which + # measurably degrades DFlash drafter acceptance (19% → target 30%+). u16 = np.frombuffer(raw, dtype=np.uint16).reshape(shape) - # bf16 = sign(1) + exp(8) + mantissa(7); reinterpret as f32 by - # putting it in the high half, then narrow to f16. - u32 = (u16.astype(np.uint32) << 16) - f32 = u32.view(" np.ndarray: SAFETENSORS_DTYPE_TO_GGUF = { "F32": gguf.GGMLQuantizationType.F32, "F16": gguf.GGMLQuantizationType.F16, - # BF16 in safetensors -> we narrow to F16 in bytes_to_np above. - "BF16": gguf.GGMLQuantizationType.F16, + # BF16 preserved as-is for sm80+ native tensor core support. + "BF16": gguf.GGMLQuantizationType.BF16, } @@ -284,6 +302,14 @@ def main(): print(f"[warn] capture_layer_ids len {len(_cap_ids)} != n_target_layers " f"{a['n_target_layers']}; not embedding ids", file=sys.stderr) + # Sliding-window attention metadata (consumed by draft_gguf_loader.cpp). + # Qwen3.6 DFlash pattern: 5 sliding_attention + 1 full_attention, window=4096. + if a.get("sliding_window"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["sliding_window"]) + if a.get("sliding_window_pattern"): + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", + a["sliding_window_pattern"]) + # Walk + add tensors. Sort: dflash.* singletons first, then output_*, # then per-layer in numeric order — keeps the on-disk layout stable. pending = [] @@ -324,7 +350,13 @@ def sort_key(t): gguf_name == "dflash.hidden_norm.weight" ) if is_norm: - arr = arr.astype("{raw_dtype.name:4s} {tuple(shape)}") diff --git a/server/scripts/quantize_draft_q8.py b/server/scripts/quantize_draft_q8.py index 6ad8533d9..7f147ec53 100644 --- a/server/scripts/quantize_draft_q8.py +++ b/server/scripts/quantize_draft_q8.py @@ -8,6 +8,8 @@ The output GGUF uses the same arch and tensor naming as convert_dflash_to_gguf.py so draft_gguf_loader.cpp can load it. +Architecture is resolved dynamically from config.json + tensor shapes +via the shared load_arch() — no hardcoded constants. Usage: python3 scripts/quantize_draft_q8.py \ @@ -16,84 +18,20 @@ """ import argparse -import json -import struct import sys from pathlib import Path import numpy as np import gguf -# ────────────────────────────────────────────────────────────────────── -# DFlash 27B draft architecture constants (must match dflash27b.h) -# ────────────────────────────────────────────────────────────────────── - -ARCH = "qwen35-dflash-draft" -HIDDEN = 5120 -N_LAYER = 5 -N_HEAD = 32 -N_HEAD_KV = 8 -HEAD_DIM = 128 -INTERMEDIATE = 17408 -VOCAB = 248320 -N_TARGET_LAYERS = 5 -ROPE_THETA = 1_000_000.0 -RMS_EPS = 1e-6 -MASK_TOKEN_ID = 248070 -BLOCK_SIZE = 16 -CTX_LEN = 32768 - -Q8_0_BLOCK_SIZE = 32 # elements per Q8_0 block - - -# ────────────────────────────────────────────────────────────────────── -# Tensor name mapping — DFlash safetensors -> llama.cpp GGUF -# (Identical to convert_dflash_to_gguf.py) -# ────────────────────────────────────────────────────────────────────── +# Share arch resolution, tensor-name mapping, and safetensors I/O with +# the F16 converter so both produce structurally identical GGUF metadata. +from convert_dflash_to_gguf import ( + ARCH, load_arch, map_name, is_norm_tensor, + load_safetensors_header, +) -def map_name(name: str) -> str | None: - if name == "fc.weight": return "dflash.fc.weight" - if name == "hidden_norm.weight": return "dflash.hidden_norm.weight" - if name == "norm.weight": return "output_norm.weight" - if name.startswith("layers."): - parts = name.split(".", 2) - if len(parts) < 3: return None - i = int(parts[1]) - rest = parts[2] - layer_map = { - "input_layernorm.weight": f"blk.{i}.attn_norm.weight", - "post_attention_layernorm.weight": f"blk.{i}.ffn_norm.weight", - "self_attn.q_proj.weight": f"blk.{i}.attn_q.weight", - "self_attn.k_proj.weight": f"blk.{i}.attn_k.weight", - "self_attn.v_proj.weight": f"blk.{i}.attn_v.weight", - "self_attn.o_proj.weight": f"blk.{i}.attn_output.weight", - "self_attn.q_norm.weight": f"blk.{i}.attn_q_norm.weight", - "self_attn.k_norm.weight": f"blk.{i}.attn_k_norm.weight", - "mlp.gate_proj.weight": f"blk.{i}.ffn_gate.weight", - "mlp.up_proj.weight": f"blk.{i}.ffn_up.weight", - "mlp.down_proj.weight": f"blk.{i}.ffn_down.weight", - } - return layer_map.get(rest) - return None - - -def is_norm_tensor(gguf_name: str) -> bool: - return ( - gguf_name.endswith("_norm.weight") or - gguf_name == "output_norm.weight" or - gguf_name == "dflash.hidden_norm.weight" - ) - - -# ────────────────────────────────────────────────────────────────────── -# safetensors reader -# ────────────────────────────────────────────────────────────────────── - -def load_safetensors_header(path: Path): - with open(path, "rb") as f: - header_size = struct.unpack(" bytes: @@ -109,6 +47,14 @@ def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray: return u32.view(" bool: + return ( + gguf_name.endswith("_norm.weight") or + gguf_name == "output_norm.weight" or + gguf_name == "dflash.hidden_norm.weight" + ) + + # ────────────────────────────────────────────────────────────────────── # Main # ────────────────────────────────────────────────────────────────────── @@ -131,27 +77,36 @@ def main(): n_entries = sum(1 for k in header if k != "__metadata__") print(f"[info] {n_entries} tensor entries") + a = load_arch(args.safetensors, header) + writer = gguf.GGUFWriter(args.out_gguf, ARCH) - # Architecture metadata (identical to convert_dflash_to_gguf.py) - writer.add_string("general.name", "Qwen3.5-27B-DFlash-Draft-Q8_0") + # Architecture metadata (resolved from config.json + tensor shapes) + writer.add_string("general.name", f"DFlash-Draft-{a['hidden']}h-{a['n_layer']}L-Q8_0") writer.add_quantization_version(gguf.GGML_QUANT_VERSION) - writer.add_uint32(f"{ARCH}.context_length", CTX_LEN) - writer.add_uint32(f"{ARCH}.embedding_length", HIDDEN) - writer.add_uint32(f"{ARCH}.block_count", N_LAYER) - writer.add_uint32(f"{ARCH}.feed_forward_length", INTERMEDIATE) - writer.add_uint32(f"{ARCH}.attention.head_count", N_HEAD) - writer.add_uint32(f"{ARCH}.attention.head_count_kv", N_HEAD_KV) - writer.add_uint32(f"{ARCH}.attention.key_length", HEAD_DIM) - writer.add_uint32(f"{ARCH}.attention.value_length", HEAD_DIM) - writer.add_uint32(f"{ARCH}.vocab_size", VOCAB) - writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", RMS_EPS) - writer.add_float32(f"{ARCH}.rope.freq_base", ROPE_THETA) + writer.add_uint32(f"{ARCH}.context_length", a["ctx_len"]) + writer.add_uint32(f"{ARCH}.embedding_length", a["hidden"]) + writer.add_uint32(f"{ARCH}.block_count", a["n_layer"]) + writer.add_uint32(f"{ARCH}.feed_forward_length", a["intermediate"]) + writer.add_uint32(f"{ARCH}.attention.head_count", a["n_head"]) + writer.add_uint32(f"{ARCH}.attention.head_count_kv", a["n_head_kv"]) + writer.add_uint32(f"{ARCH}.attention.key_length", a["head_dim"]) + writer.add_uint32(f"{ARCH}.attention.value_length", a["head_dim"]) + writer.add_uint32(f"{ARCH}.vocab_size", a["vocab"]) + writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", a["rms_eps"]) + writer.add_float32(f"{ARCH}.rope.freq_base", a["rope_theta"]) # DFlash-specific hyperparameters - writer.add_uint32(f"{ARCH}.dflash.n_target_layers", N_TARGET_LAYERS) - writer.add_uint32(f"{ARCH}.dflash.block_size", BLOCK_SIZE) - writer.add_uint32(f"{ARCH}.dflash.mask_token_id", MASK_TOKEN_ID) + writer.add_uint32(f"{ARCH}.dflash.n_target_layers", a["n_target_layers"]) + writer.add_uint32(f"{ARCH}.dflash.block_size", a["block_size"]) + writer.add_uint32(f"{ARCH}.dflash.mask_token_id", a["mask_token_id"]) + + # Sliding-window attention metadata (consumed by draft_gguf_loader.cpp). + if a.get("sliding_window"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["sliding_window"]) + if a.get("sliding_window_pattern"): + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", + a["sliding_window_pattern"]) # Collect and sort tensors (same order as convert_dflash_to_gguf.py) pending = [] diff --git a/server/src/common/dflash_draft_graph.cpp b/server/src/common/dflash_draft_graph.cpp index c028bb889..80d9cf5f5 100644 --- a/server/src/common/dflash_draft_graph.cpp +++ b/server/src/common/dflash_draft_graph.cpp @@ -34,7 +34,7 @@ static bool build_draft_graph_internal( bool mirror_view) { ggml_init_params ip{}; - ip.mem_size = 256 * 1024 * 1024; + ip.mem_size = 4 * 1024 * 1024; // 4MB (was 256MB — wildly oversized for 6-layer draft) ip.mem_buffer = nullptr; ip.no_alloc = true; sg.ctx = ggml_init(ip); @@ -124,23 +124,36 @@ bool build_draft_step( const DraftFeatureMirror * mirror, int committed, int /*ctx_len_max*/) { + int mirror_slot0 = 0; + const bool use_view = mirror && + draft_feature_mirror_can_view(*mirror, committed, ctx_len, mirror_slot0); + + // Reuse guard: once committed exceeds the draft context cap, ctx_len is + // constant across steps and — since the mirror view is almost never + // available in steady state (it requires committed % cap == 0) — the graph + // is built non-view with identical topology every step. Skip the rebuild + // and let the caller refresh inp_embed / target_hidden_cat / positions; + // the attn_mask depends only on ctx_len and is already populated. Views are + // never reused: their baked-in ggml_view_3d offset slides with `committed`, + // so a reused view would point at the wrong ring slot. Disable with + // DFLASH_DRAFT_GRAPH_REUSE=0 for A/B comparison. + static const bool kReuseEnabled = []() { + const char * e = std::getenv("DFLASH_DRAFT_GRAPH_REUSE"); + return e == nullptr || std::string(e) != "0"; + }(); + if (kReuseEnabled && sg.ctx && sg.graph_ctx_len == ctx_len && + !use_view && !sg.graph_used_view) { + return true; + } + step_graph_free(sg); if (!sg.alloc) { sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); } - int mirror_slot0 = 0; - const bool use_view = mirror && - draft_feature_mirror_can_view(*mirror, committed, ctx_len, mirror_slot0); - - // If ctx_len exceeds our cached reserve, re-reserve at next 64 boundary. - // This makes all subsequent alloc_graph calls within the 64-token window - // a no-op (no CUDA free+alloc). const int ctx_padded = (ctx_len + 63) & ~63; if (ctx_padded > sg.alloc_reserved_ctx) { - // Build a dummy graph at ctx_padded just for sizing. - // Use non-view path for reserve (view tensors don't need allocation). if (!build_draft_graph_internal(sg, dw, lm_head, ctx_padded, nullptr, 0, false)) { return false; @@ -150,7 +163,6 @@ bool build_draft_step( step_graph_free(sg); } - // Build real graph at ctx_len for actual computation. if (!build_draft_graph_internal(sg, dw, lm_head, ctx_len, mirror, mirror_slot0, use_view)) { return false; @@ -160,7 +172,6 @@ bool build_draft_step( return false; } - // Fill causal mask data for SWA layers (after allocation gives memory to the tensor). if (sg.attn_mask) { const int q_len = dw.block_size; const bool swa_active = dw.swa_window > 0 && ctx_len > dw.swa_window; @@ -168,9 +179,6 @@ bool build_draft_step( const int eff_total_k = eff_ctx + q_len; const int kv_pad = mask_align_up(eff_total_k, MASK_KV_PAD); - // Build causal mask in F16 directly (same pattern as attn_masks.h): - // Context keys (k < eff_ctx): always visible. - // Noise keys (k = eff_ctx + j): visible if j <= q (causal). static constexpr uint16_t ZERO = 0x0000; static constexpr uint16_t NEG_INF = 0xFC00; std::vector mask_data((size_t)kv_pad * q_len, NEG_INF); @@ -184,6 +192,8 @@ bool build_draft_step( sizeof(uint16_t) * mask_data.size()); } + sg.graph_ctx_len = ctx_len; + sg.graph_used_view = use_view; return true; } diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index 0d80480b2..55ee01a40 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -25,6 +25,14 @@ struct StepGraph { // When the real ctx_len fits within this, alloc_graph is a no-op. int alloc_reserved_ctx = 0; + // Draft-graph reuse guard: the ctx_len and view-mode the current draft + // graph was built for. When the next step requests the same ctx_len with a + // non-view build, the topology is identical and build_draft_step skips the + // rebuild (the caller refreshes input tensors). Views aren't reused because + // their offset slides each step with `committed`. + int graph_ctx_len = 0; + bool graph_used_view = false; + // Named inputs ggml_tensor * inp_embed = nullptr; ggml_tensor * positions = nullptr; @@ -73,6 +81,8 @@ inline void step_graph_free(StepGraph & sg) { sg.valid_lut = nullptr; sg.delta_captures.clear(); sg.moe_selected.clear(); + sg.graph_ctx_len = 0; + sg.graph_used_view = false; } // Full cleanup: release the persistent gallocr + its CUDA buffer. diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index 95dfc3ee0..81f938b04 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -117,6 +117,65 @@ int count_swa_layers(const DraftWeights & w) { } // namespace +// Lightweight metadata-only read of the drafter's capture-layer config. +// Opens the GGUF header (no tensor data loaded), extracts n_target_layers +// and target_layer_ids, and returns them. Called by the backend BEFORE +// load_target_model so that target_feat is allocated with the correct +// fc_in = n_capture_layers * n_embd dimension. +bool read_draft_capture_config(const std::string & path, + int & n_capture, + int * capture_ids, + int max_ids) { + n_capture = 0; + gguf_init_params gip{}; + gip.no_alloc = true; + gip.ctx = nullptr; // metadata-only: no tensor context + gguf_context * gctx = gguf_init_from_file(path.c_str(), gip); + if (!gctx) return false; + + // Resolve the arch prefix for key building. + std::string A = "qwen35-dflash-draft"; + { + int64_t arch_id = gguf_find_key(gctx, "general.architecture"); + if (arch_id >= 0) { + const char * arch = gguf_get_val_str(gctx, arch_id); + if (arch) A = arch; + } + } + + char key[128]; + + // Read n_target_layers. + std::snprintf(key, sizeof(key), "%s.%s", A.c_str(), "dflash.n_target_layers"); + int64_t ntl_id = gguf_find_key(gctx, key); + if (ntl_id >= 0) { + n_capture = (int)gguf_get_val_u32(gctx, ntl_id); + } + + // Read target_layer_ids array (exact capture positions from training). + std::snprintf(key, sizeof(key), "%s.%s", A.c_str(), "dflash.target_layer_ids"); + int64_t tli_id = gguf_find_key(gctx, key); + if (tli_id >= 0 && gguf_get_kv_type(gctx, tli_id) == GGUF_TYPE_ARRAY) { + const size_t n = std::min((size_t)gguf_get_arr_n(gctx, tli_id), + (size_t)max_ids); + const int32_t * ids = static_cast(gguf_get_arr_data(gctx, tli_id)); + for (size_t k = 0; k < n; k++) { + capture_ids[k] = (int)ids[k]; + } + if (n > 0) n_capture = (int)n; + } + + gguf_free(gctx); + + if (n_capture > 0) { + std::fprintf(stderr, "[draft-cfg] early-read: n_capture=%d ids=", n_capture); + for (int k = 0; k < n_capture && k < max_ids; k++) + std::fprintf(stderr, "%s%d", k ? "," : "[", capture_ids[k]); + std::fprintf(stderr, "]\n"); + } + return n_capture > 0; +} + bool load_draft_gguf(const std::string & path, ggml_backend_t backend, DraftWeights & out, @@ -214,11 +273,30 @@ bool load_draft_gguf(const std::string & path, // Store GGUF-declared config into DraftWeights (replaces hardcoded defaults). out.block_size = (int)block_sz; + // Env override: DFLASH_DRAFT_BLOCK_SIZE allows serving with a smaller + // block than the model was trained with (e.g. 8 instead of 16). The model + // card shows block=8 gives similar throughput with lower per-step cost. + if (const char * bs_env = std::getenv("DFLASH_DRAFT_BLOCK_SIZE")) { + int bs = std::atoi(bs_env); + if (bs >= 4 && bs <= 32) { + out.block_size = bs; + std::fprintf(stderr, "[draft] block_size override: %d (from env)\n", bs); + } + } out.n_target_layers = (int)n_tgt_lay; - // Propagate target model properties if available. - if (target) { - out.mask_token_id = target->mask_token_id; + // Read mask_token_id from GGUF metadata. Each DFlash drafter is trained + // with a specific MASK token from the target's vocabulary (e.g. Qwen3.5 + // uses 248070, Qwen3.6 uses 248077). The GGUF carries the correct value + // via dflash.mask_token_id. Only fall back to the target's value (which + // defaults to the old Qwen3.5 constant) if the GGUF doesn't declare it. + { + const uint32_t gguf_mask = read_u32("dflash.mask_token_id", 0xFFFFFFFFu); + if (gguf_mask != 0xFFFFFFFFu) { + out.mask_token_id = (int32_t)gguf_mask; + } else if (target) { + out.mask_token_id = target->mask_token_id; + } } // Upper bounds on hparams. Guards against malformed/hostile GGUFs that diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 11f5b4c51..428b3f2e9 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -49,6 +49,23 @@ DraftGraphOutputs build_draft_graph( const float eps = DFLASH27B_RMS_EPS; const float rope_base = w.rope_theta; + // Early-exit: stop the draft forward after N layers (env override). + // The new 6-layer DFlash drafter is deep enough that layers 5-6 contribute + // marginally to proposal quality. Exiting at layer 3-4 saves ~33-50% of + // the drafter forward wall time. Acceptance drops slightly but the + // throughput gain dominates. + const int full_n_layer = w.n_layer; + const int early_exit_n = []() -> int { + const char * e = std::getenv("DFLASH_DRAFT_EARLY_EXIT_N"); + if (e) { + int v = std::atoi(e); + if (v > 0) return v; + } + return 0; // 0 = sentinel for "use all layers" + }(); + const int fwd_layer_limit = (early_exit_n > 0 && early_exit_n < full_n_layer) + ? early_exit_n : full_n_layer; + // ── 1. Feature fusion: target_feat = rms_norm(fc @ target_hidden_cat, hidden_norm) // fc: [5*hidden, hidden] (ggml: ne[0]=5*hidden, ne[1]=hidden) // target_hidden_cat: [5*hidden, ctx_len, 1] @@ -61,7 +78,7 @@ DraftGraphOutputs build_draft_graph( // ── 2. Decoder layers ggml_tensor * h = in.noise_embed; // [hidden, q_len, 1] - for (int il = 0; il < w.n_layer; il++) { + for (int il = 0; il < fwd_layer_limit; il++) { const DraftLayer & L = w.layers[il]; // ── SWA: determine effective context for this layer diff --git a/server/src/draft/draft_safetensors_loader.cpp b/server/src/draft/draft_safetensors_loader.cpp index 1778cba00..975959e8a 100644 --- a/server/src/draft/draft_safetensors_loader.cpp +++ b/server/src/draft/draft_safetensors_loader.cpp @@ -507,6 +507,72 @@ static float read_rope_theta_from_config(const std::string & st_path) { return val; } +// Read DFlash-specific config from config.json next to the safetensors: +// dflash_config.mask_token_id, sliding_window, layer_types (SWA pattern). +// These are essential for Qwen3.6 drafters (mask=248077, 5 SWA + 1 full). +static void read_dflash_config_from_json(const std::string & st_path, DraftWeights & out) { + std::string dir = st_path; + const size_t slash = dir.find_last_of("/\\"); + if (slash != std::string::npos) dir.resize(slash + 1); else dir = "./"; + const std::string cfg_path = dir + "config.json"; + + FILE * f = std::fopen(cfg_path.c_str(), "r"); + if (!f) return; + std::string buf; + char tmp[4096]; + while (std::fgets(tmp, sizeof(tmp), f)) buf += tmp; + std::fclose(f); + + // mask_token_id: look under dflash_config first, then top level. + auto find_int = [&](const std::string & key) -> long long { + size_t pos = buf.find("\"" + key + "\""); + if (pos == std::string::npos) return LLONG_MIN; + pos += key.size() + 3; // skip "key" + while (pos < buf.size() && (buf[pos] == ' ' || buf[pos] == ':' || buf[pos] == '\t')) pos++; + return std::strtoll(buf.c_str() + pos, nullptr, 10); + }; + + long long mask = find_int("mask_token_id"); + if (mask > 0 && mask < INT32_MAX) { + out.mask_token_id = (int32_t)mask; + std::fprintf(stderr, "[draft-st] mask_token_id=%d (from config.json)\n", out.mask_token_id); + } + + long long sw = find_int("sliding_window"); + if (sw > 0) { + out.swa_window = (int)sw; + std::fprintf(stderr, "[draft-st] sliding_window=%d (from config.json)\n", out.swa_window); + } + + // layer_types: set is_swa per layer. "sliding_attention" = true, "full_attention" = false. + // Scan for the array and assign in order. + size_t lt_pos = buf.find("\"layer_types\""); + if (lt_pos != std::string::npos) { + size_t arr_start = buf.find('[', lt_pos); + size_t arr_end = buf.find(']', arr_start); + if (arr_start != std::string::npos && arr_end != std::string::npos) { + std::string arr = buf.substr(arr_start, arr_end - arr_start); + int layer_idx = 0; + size_t s = 0; + while (layer_idx < out.n_layer) { + size_t q1 = arr.find('"', s); + if (q1 == std::string::npos) break; + size_t q2 = arr.find('"', q1 + 1); + if (q2 == std::string::npos) break; + std::string lt = arr.substr(q1 + 1, q2 - q1 - 1); + out.layers[layer_idx].is_swa = (lt == "sliding_attention"); + layer_idx++; + s = q2 + 1; + } + int n_swa = 0; + for (int i = 0; i < out.n_layer; i++) if (out.layers[i].is_swa) n_swa++; + if (n_swa > 0) + std::fprintf(stderr, "[draft-st] SWA layers: %d/%d (window=%d)\n", + n_swa, out.n_layer, out.swa_window); + } + } +} + } // namespace bool load_draft_safetensors(const std::string & path, @@ -535,9 +601,24 @@ bool load_draft_safetensors(const std::string & path, const uint8_t * blob = (const uint8_t *)mm.addr + 8 + header_len; const size_t blob_len = mm.len - 8 - header_len; - // ── 3. Allocate ggml context big enough for 5 layers × 11 + 3 top ─ - const int n_layers = DFLASH27B_DRAFT_LAYERS; - const int n_tensors = 3 + 11 * n_layers; // with some headroom below + // ── 3. Allocate ggml context — count actual layers from header ── + // Each layer has 11 tensors (attn_norm, ffn_norm, wq, wk, wv, wo, q_norm, + // k_norm, w_gate, w_up, w_down). Count by scanning for layers.N. prefixes. + int n_layers = DFLASH27B_DRAFT_LAYERS; + { + int max_layer = -1; + for (const auto & [name, _] : st) { + if (name.substr(0, 7) == "layers.") { + size_t dot = name.find('.', 7); + if (dot != std::string::npos) { + int li = std::atoi(name.c_str() + 7); + if (li > max_layer) max_layer = li; + } + } + } + if (max_layer >= 0) n_layers = max_layer + 1; + } + const int n_tensors = 3 + 11 * n_layers + 16; // headroom ggml_init_params ip{}; ip.mem_size = (size_t)(n_tensors + 16) * ggml_tensor_overhead(); ip.mem_buffer = nullptr; @@ -559,7 +640,11 @@ bool load_draft_safetensors(const std::string & path, return false; } if (target) { - out.mask_token_id = target->mask_token_id; + // Only inherit target's mask_token_id if our config.json didn't set one. + // (read_dflash_config_from_json below may override this.) + if (out.mask_token_id == DFLASH27B_DRAFT_MASK_TOKEN_ID) { + out.mask_token_id = target->mask_token_id; + } if (out.n_embd != target->n_embd) { char buf[256]; std::snprintf(buf, sizeof(buf), @@ -579,6 +664,11 @@ bool load_draft_safetensors(const std::string & path, } out.layers.assign(n_layers, DraftLayer{}); + // Read mask_token_id and SWA config from config.json (Qwen3.6 drafters + // use mask_token_id=248077 and 5 sliding + 1 full attention layers). + // Must be called AFTER out.layers.assign() since it sets is_swa per layer. + read_dflash_config_from_json(path, out); + const int64_t HIDDEN = out.n_embd; const int64_t Q_DIM = out.n_head * out.head_dim; const int64_t KV_DIM = out.n_head_kv * out.head_dim; diff --git a/server/src/internal.h b/server/src/internal.h index 650f6a7d2..459b5937c 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -199,8 +199,10 @@ struct TargetWeights { // Target layer IDs captured for the DFlash draft model. // Computed from n_layer at load time: step = (n_layer - 2) / (N - 1), // ids[k] = 1 + k * step. E.g. 27B→{1,16,31,46,61}, 9B→{1,8,15,22,29}. + // For newer drafters (8 taps), the exact IDs come from the drafter GGUF's + // dflash.target_layer_ids via the backend's early-read sync. int n_capture_layers = DFLASH27B_DRAFT_N_TARGET_LAYERS; - int capture_layer_ids[DFLASH27B_DRAFT_N_TARGET_LAYERS] = {1, 16, 31, 46, 61}; + int capture_layer_ids[DFLASH_MAX_CAPTURE_LAYERS] = {1, 16, 31, 46, 61}; }; // Check if a token is an end-of-sequence marker for the given target weights. @@ -295,6 +297,15 @@ bool load_draft_gguf(const std::string & path, DraftWeights & out, const TargetWeights * target = nullptr); +// Metadata-only read of the drafter's capture-layer config (n_target_layers +// + target_layer_ids). Called BEFORE load_target_model so the target's +// target_feat tensor is allocated with the correct fc_in dimension. +// Returns true if n_capture was resolved (>0). +bool read_draft_capture_config(const std::string & path, + int & n_capture, + int * capture_ids, + int max_ids); + void free_draft_weights(DraftWeights & w); // ─── Target cache (persistent state between forward calls) ──────── diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 8b9a27aac..8820e1fa9 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -463,12 +463,41 @@ bool load_target_gguf_partial(const std::string & path, std::printf("[loader] eos_id=%d eos_chat_id=%d\n", out.eos_id, out.eos_chat_id); } - // Compute capture layer IDs: evenly spaced through the target layers. - // step = (n_layer - 2) / (N - 1), ids[k] = 1 + k * step. + // Capture layer IDs: use exact IDs from the drafter's config if the + // early-read sync set them (from dflash.target_layer_ids in the drafter + // GGUF). Otherwise fall back to evenly-spaced heuristic. + // The drafter's fc projection was trained against specific target layers; + // an off-by-one in even one tap degrades acceptance significantly. { const int N = out.n_capture_layers; - const int step = ((int)n_layer - 2) / (N - 1); - for (int k = 0; k < N; k++) out.capture_layer_ids[k] = 1 + k * step; + bool ids_valid = (N > 0 && N <= DFLASH_MAX_CAPTURE_LAYERS); + if (ids_valid) { + // Check if early-read set non-default values (i.e., not the + // old 5-layer default {1,16,31,46,61}). + // If N changed from default 5, the IDs were definitely set by + // early-read and should be respected. + const bool was_early_read = (N != DFLASH27B_DRAFT_N_TARGET_LAYERS); + if (was_early_read) { + std::printf("[loader] using drafter-specified capture layers (%d)\n", N); + } else { + // Even for N==5, check if IDs differ from default heuristic. + const int step = ((int)n_layer - 2) / (N - 1); + bool matches_heuristic = true; + for (int k = 0; k < N; k++) { + if (out.capture_layer_ids[k] != 1 + k * step) { + matches_heuristic = false; + break; + } + } + if (!matches_heuristic) { + std::printf("[loader] using drafter-specified capture layers (non-default)\n"); + } else { + // Matches heuristic — recompute for consistency. + for (int k = 0; k < N; k++) + out.capture_layer_ids[k] = 1 + k * step; + } + } + } } out.layers.assign((size_t)n_layer, TargetLayer{}); diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index b4228c3d3..4ad9e1d3f 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -448,10 +448,23 @@ bool build_lm_head_projection_step( const TargetWeights & w, ggml_backend_t backend, int n_tokens) { + // Persistent reuse: the lm-head projection graph (mul_mat + argmax) + // depends only on n_tokens, which is constant across spec-decode steps + // (q_len is fixed by block_size). If a graph for this n_tokens is already + // alive, skip the rebuild — the caller refreshes hidden_input data each + // step. Centralises the per-call-site guards already present in the + // layer-split target, the MoE hybrid path, and the test harness, and + // covers the base Qwen35DFlashTarget path (project_hidden_to_tokens / + // project_hidden_to_topk) which rebuilt unconditionally every step. + if (sg.ctx && sg.gf && sg.hidden_input && + sg.hidden_input->ne[1] == n_tokens && sg.logits && sg.argmax_tokens) { + return true; + } + step_graph_free(sg); ggml_init_params ip{}; - ip.mem_size = 64 * 1024 * 1024; + ip.mem_size = 4 * 1024 * 1024; // 4MB (was 64MB — only 3 ops in this graph) ip.mem_buffer = nullptr; ip.no_alloc = true; sg.ctx = ggml_init(ip); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 741297453..b5d04a912 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -198,6 +198,24 @@ bool Qwen35Backend::init() { return false; } + // Early-read drafter capture config so the target graph is built with + // the correct fc_in = n_capture_layers * n_embd. Newer 8-tap drafters + // need 8 captured layers; without this sync the default 5 would cause + // a runtime shape mismatch on the drafter's fc projection. + if (cfg_.draft_path) { + std::string dp(cfg_.draft_path); + if (dp.size() >= 5 && dp.substr(dp.size() - 5) == ".gguf") { + int n_cap = 0; + int cap_ids[DFLASH_MAX_CAPTURE_LAYERS]; + if (read_draft_capture_config(cfg_.draft_path, n_cap, cap_ids, + DFLASH_MAX_CAPTURE_LAYERS)) { + w_.n_capture_layers = n_cap; + for (int k = 0; k < n_cap && k < DFLASH_MAX_CAPTURE_LAYERS; k++) + w_.capture_layer_ids[k] = cap_ids[k]; + } + } + } + // Load target if (!load_target_model(target_backend_, w_)) { std::fprintf(stderr, "target load: %s\n", dflash27b_last_error()); @@ -232,6 +250,16 @@ bool Qwen35Backend::init() { } std::printf("[draft] loaded\n"); + // Propagate the draft's mask_token_id to the target weights. The + // noise-generation code (which fills MASK tokens for block diffusion) + // reads target_weights().mask_token_id, so it must match the drafter's + // trained mask token (248077 for Qwen3.6, vs the old 248070 default). + if (dw_.mask_token_id != w_.mask_token_id) { + std::printf("[draft] mask_token_id %d -> %d (from draft GGUF)\n", + w_.mask_token_id, dw_.mask_token_id); + w_.mask_token_id = dw_.mask_token_id; + } + if (cfg_.draft_swa_window > 0) { dw_.swa_window = cfg_.draft_swa_window; for (int il = 0; il < dw_.n_layer - 1; il++) @@ -337,7 +365,12 @@ bool Qwen35Backend::init() { // Init feature mirror when draft model is available (needed for spec decode). // On single-GPU, this is an F32 conversion buffer; on split-GPU, a cross-device mirror. if (cfg_.draft_path && !use_remote_draft) { - const int mirror_cap = std::min({cfg_.draft_ctx_max, cfg_.device.max_ctx, + // Honor DFLASH_FEAT_RING_CAP for the mirror cap as well (not just target_feat). + int effective_draft_ctx_max = cfg_.draft_ctx_max; + if (const char * e = std::getenv("DFLASH_FEAT_RING_CAP")) { + effective_draft_ctx_max = std::max(cfg_.draft_ctx_max, std::atoi(e)); + } + const int mirror_cap = std::min({effective_draft_ctx_max, cfg_.device.max_ctx, cache_.target_feat_cap > 0 ? cache_.target_feat_cap : cfg_.device.max_ctx}); if (!draft_feature_mirror_init(feature_mirror_, draft_backend_, cfg_.draft_gpu, cfg_.device.gpu, mirror_cap, @@ -2049,8 +2082,6 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } } - StepGraph draft_sg; - std::vector noise_embed((size_t)hidden * q_len); std::vector noise_ids(q_len); std::vector draft_tok(q_len); @@ -2099,7 +2130,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, "[spec-decode] invalid draft seed %d after %d emitted tokens; " "switching to AR\n", last_tok, (int)out_tokens.size()); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); cache_.last_tok = out_tokens.back(); const int ar_n_gen = n_gen - n_generated; if (ar_n_gen <= 0) { @@ -2122,7 +2153,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->embed_tokens(noise_ids.data(), q_len, noise_embed.data())) { std::fprintf(stderr, "spec-decode: noise embed failed (last_tok=%d mask=%d q_len=%d)\n", last_tok, target->mask_token_id(), q_len); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } @@ -2141,52 +2172,52 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, local_hidden.clear(); if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } } else { - if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, + if (!build_draft_step(draft_sg_, dw_, /*lm_head=*/nullptr, draft_backend_, draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, committed, /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { std::fprintf(stderr, "spec-decode: draft build failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } if (!use_mirror_view && - !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, + !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg_.target_hidden_cat, draft_start, draft_ctx)) { std::fprintf(stderr, "spec-decode: feature copy failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } - ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, + ggml_backend_tensor_set(draft_sg_.inp_embed, noise_embed.data(), 0, sizeof(float) * noise_embed.size()); pos_k.resize((size_t)draft_ctx + q_len); for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; - ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, + ggml_backend_tensor_set(draft_sg_.positions, pos_q.data(), 0, sizeof(int32_t) * pos_q.size()); - ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, + ggml_backend_tensor_set(draft_sg_.positions_k, pos_k.data(), 0, sizeof(int32_t) * pos_k.size()); - auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); + auto st = ggml_backend_graph_compute(draft_backend_, draft_sg_.gf); if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "spec-decode: draft compute failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } // Read draft hidden states to host for LM-head projection. local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, + ggml_backend_tensor_get(draft_sg_.hidden_states, local_hidden.data(), 0, sizeof(float) * local_hidden.size()); } // 3. Project draft hidden → token IDs via target LM head if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { std::fprintf(stderr, "spec-decode: projection failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } draft_tok[0] = last_tok; @@ -2228,7 +2259,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->project_hidden_to_topk(local_hidden.data(), q_len, K, cfg_.ddtree_temp, top_lp, top_ids)) { std::fprintf(stderr, "spec-decode: ddtree topk projection failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } // Tree depth L draws from draft rows 1..q_len-1 (row 0 = the seed). @@ -2244,7 +2275,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, for (int i = 0; i < tree.n_nodes; i++) flat_tokens[1 + i] = tree.token_ids[i]; if (!sampled_verify && !target->snapshot_kv()) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } @@ -2253,7 +2284,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->verify_tree(committed, tree, flat_tokens, N, posterior, sampled_verify ? &node_logits : nullptr)) { std::fprintf(stderr, "spec-decode: verify_tree failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; @@ -2286,7 +2317,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int accepted_n = (int)accepted.size(); // root + accepted children if (accepted_n > need_commit_budget) accepted_n = need_commit_budget; - if (accepted_n <= 0) { step_graph_destroy(draft_sg); break; } + if (accepted_n <= 0) { step_graph_destroy(draft_sg_); break; } // Emit the accepted path: slot 0 = last_tok (pending from prev iter), // each subsequent accepted node = its tree token. @@ -2305,7 +2336,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Telemetry: accepted children (exclude the always-committed root). n_accept_sum += std::max(0, accepted_emitted - 1); - if (accepted_emitted <= 0) { step_graph_destroy(draft_sg); break; } + if (accepted_emitted <= 0) { step_graph_destroy(draft_sg_); break; } if (!sampled_verify) { const int root_last_tok = last_tok; @@ -2321,7 +2352,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, accepted.begin() + accepted_emitted); if (!target->rollback_to_tree(committed, tree, accepted_committed)) { std::fprintf(stderr, "spec-decode: rollback_to_tree failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } last_tok = next_token; @@ -2359,13 +2390,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (can_commit_bonus) replay_batch.push_back(next_token); if (!target->restore_kv()) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } int replay_last_tok = -1; if (!target->verify_batch(replay_batch, committed, replay_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: tree replay failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; @@ -2404,7 +2435,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, accepted.begin() + accepted_emitted); if (!target->rollback_to_tree(committed, tree, accepted_committed)) { std::fprintf(stderr, "spec-decode: rollback_to_tree failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } @@ -2419,18 +2450,18 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector bonus_vec(1, next_token); if (!target->verify_batch(bonus_vec, bonus_pos, bonus_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: tree bonus replay failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; if (!target->read_verify_logits(1, bonus_logits)) { std::fprintf(stderr, "spec-decode: tree bonus logits read failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } if (bonus_logits.empty()) { std::fprintf(stderr, "spec-decode: tree bonus logits empty\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } @@ -2484,7 +2515,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // 4. Verify: snapshot KV, run target forward over draft tokens if (!target->snapshot_kv()) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } @@ -2493,7 +2524,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, /*capture_ssm_intermediates=*/true)) { std::fprintf(stderr, "spec-decode: verify failed\n"); target->restore_kv(); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; @@ -2509,7 +2540,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->read_verify_logits(q_len, verify_logits)) { std::fprintf(stderr, "spec-decode: verify logits read failed\n"); target->restore_kv(); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } const int vocab_v = (int)(verify_logits.size() / (size_t)q_len); @@ -2589,7 +2620,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // accept_n is large enough that skipping the replay saves more compute // than the cost of deferring the bonus to the next step. Breakeven // is around accept_n ≈ 5. Below that, legacy replay is cheaper. - constexpr int kFastRollbackThreshold = 5; + // Env-tunable (DFLASH_FAST_ROLLBACK_MIN) for A/B; default 5. + static const int kFastRollbackThreshold = []() { + const char * e = std::getenv("DFLASH_FAST_ROLLBACK_MIN"); + return e ? std::atoi(e) : 5; + }(); const bool use_fast_rollback = target->supports_fast_rollback() && (accept_n >= kFastRollbackThreshold); @@ -2620,7 +2655,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // (When falling back from fast-rollback, bonus_tok is -1 and commit_n // is the budget-clamped accepted count.) if (!target->restore_kv()) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } std::vector replay_batch((size_t)commit_n); @@ -2629,7 +2664,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } if (!target->verify_batch(replay_batch, committed, replay_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: replay failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; @@ -2644,7 +2679,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // 7. Sync features for replayed range to mirror (needed for next draft step) if (use_remote_draft && cache_.target_feat) { if (!sync_remote_draft_features(committed, commit_n)) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } } else if (feature_mirror_.target_feat && cache_.target_feat) { @@ -2729,7 +2764,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int injected = 0; if (floor_to_ar) { if (!target->restore_kv()) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } cache_.cur_pos = committed; @@ -2740,7 +2775,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->verify_batch(replay_prefix, committed, prefix_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: floor prefix replay failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; @@ -2752,7 +2787,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->verify_batch(*stall_tool_prefix_tokens, committed, tool_prefix_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: tool prefix replay failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; @@ -2800,7 +2835,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (io.cancelled) break; if (floor_to_ar) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); const int total_draft_pos = std::max(1, n_draft_steps * q_len); out_accept_rate = @@ -2825,7 +2860,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // the emitted output before AR takes over. if (budget_close_fired) { if (!target->restore_kv()) { - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } cache_.cur_pos = committed; @@ -2836,14 +2871,14 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!target->verify_batch(replay_prefix, committed, prefix_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: budget-close prefix replay failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); return false; } target_forwards++; } committed += emitted; cache_.cur_pos = committed; - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); const int total_draft_pos = std::max(1, n_draft_steps * q_len); out_accept_rate = @@ -2866,7 +2901,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (hit_eos) break; } - step_graph_destroy(draft_sg); + step_graph_destroy(draft_sg_); auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index 4c01ec27a..5fc323250 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -640,6 +640,15 @@ bool Qwen35DFlashTarget::project_hidden_to_tokens( tokens_out.resize(n_tokens); ggml_backend_tensor_get(proj_sg_.argmax_tokens, tokens_out.data(), 0, sizeof(int32_t) * n_tokens); + // Guard against invalid token IDs: at long context, NaN/Inf in the drafter's + // hidden states can produce garbage argmax results (negative or >= vocab). + // Clamp to valid range so verify_batch's embedder doesn't crash. The verify + // will simply reject the bad proposal (target argmax won't match). + for (int i = 0; i < n_tokens; i++) { + if (tokens_out[i] < 0 || tokens_out[i] >= w_.n_vocab) { + tokens_out[i] = 0; // padding token — verify rejects, spec degrades gracefully + } + } return true; } diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index ac0222fe4..65501dac5 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -190,8 +190,14 @@ bool create_target_cache_partial(const TargetWeights & w, } } - constexpr int TARGET_FEAT_CAP_DEFAULT = 4096; - out.target_feat_cap = std::min(max_ctx, TARGET_FEAT_CAP_DEFAULT); + // Feature ring cap. The drafter needs the last `cap` target hidden states + // as context. 4096 is enough for short prompts but acceptance collapses + // at exactly 4096 (ring wrap). Override with DFLASH_FEAT_RING_CAP env. + int target_feat_cap_default = 4096; + if (const char * e = std::getenv("DFLASH_FEAT_RING_CAP")) { + target_feat_cap_default = std::atoi(e); + } + out.target_feat_cap = std::min(max_ctx, target_feat_cap_default); if (allocate_target_feat) { const int fc_in = w.n_capture_layers * w.n_embd; out.target_feat = ggml_new_tensor_2d(out.base_ctx, GGML_TYPE_BF16, fc_in, out.target_feat_cap); diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index b56836973..f5c4a19c1 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -110,13 +110,28 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & // Record the placement result so post_kvflash_init_gate() can disable // the kvflash pool — moe_hybrid will be null on this path (no cold storage // needed), so the gate cannot detect all-hot from the hybrid pointer alone. - if (placement.total_hot >= out.n_layer * out.n_expert) { + // + // Env override: DFLASH_MOE_ALLHOT_HYBRID=1 skips the early return and builds + // moe_hybrid storage even with 0 cold experts. This routes both AR and spec + // through the MoE-specific paths (pipelined decode / do_hybrid_spec_decode) + // instead of the base Qwen35Backend paths, which use the slow DFlashTarget + // adapter verify. The pipelined verify is ~0.9ms/tok vs ~4.5ms/tok batched. + static const bool kForceMoeHybrid = []() { + const char * e = std::getenv("DFLASH_MOE_ALLHOT_HYBRID"); + return e != nullptr && std::string(e) == "1"; + }(); + placement_all_hot_ = (placement.total_hot >= out.n_layer * out.n_expert); + if (placement_all_hot_ && !kForceMoeHybrid) { std::printf("[qwen35moe] all experts fit in VRAM, loading fully to GPU\n"); std::fflush(stdout); - placement_all_hot_ = true; free_target_weights(out); return load_target_gguf(cfg_.target_path, backend, out); } + if (placement_all_hot_ && kForceMoeHybrid) { + std::printf("[qwen35moe] all experts fit in VRAM but building moe_hybrid " + "(DFLASH_MOE_ALLHOT_HYBRID=1) to enable pipelined spec-decode verify\n"); + std::fflush(stdout); + } if (const char * telemetry = std::getenv("DFLASH_QWEN35MOE_TELEMETRY")) { hybrid_telemetry_ = std::atoi(telemetry) != 0; @@ -1303,6 +1318,8 @@ GenerateResult Qwen35MoeBackend::generate_impl(const GenerateRequest & req, decode_tel_accum.total_us += tel.total_us; decode_tel_accum.prefn_graph_build_us += tel.prefn_graph_build_us; decode_tel_accum.prefn_compute_us += tel.prefn_compute_us; + decode_tel_accum.prefn_ssm_us += tel.prefn_ssm_us; + decode_tel_accum.prefn_attn_us += tel.prefn_attn_us; decode_tel_accum.routing_readback_us += tel.routing_readback_us; decode_tel_accum.ffn_us += tel.ffn_us; decode_tel_accum.ffn_allhot_us += tel.ffn_allhot_us; @@ -1411,6 +1428,9 @@ GenerateResult Qwen35MoeBackend::generate_impl(const GenerateRequest & req, decode_tel_accum.prefn_compute_us / 1000.0 / n_dec, decode_tel_accum.routing_readback_us / 1000.0 / n_dec, decode_tel_accum.ffn_us / 1000.0 / n_dec); + std::printf(" per-token avg: prefn_SSM=%.2fms (30 DeltaNet) prefn_ATTN=%.2fms (10 full-attn)\n", + decode_tel_accum.prefn_ssm_us / 1000.0 / n_dec, + decode_tel_accum.prefn_attn_us / 1000.0 / n_dec); std::printf(" per-token avg: tensor_io=%.2fms combine=%.2fms cold_cpu=%.2fms cold_compute=%.2fms\n", decode_tel_accum.tensor_io_us / 1000.0 / n_dec, decode_tel_accum.combine_overhead_us / 1000.0 / n_dec, @@ -1686,54 +1706,65 @@ bool Qwen35MoeBackend::hybrid_forward_one_token(int32_t tok, int kv_pos, } } - // Project to logits and get argmax + // Project to logits and get argmax. + // Persistent graph: built once (rms_norm + out_norm + mul_mat + argmax), + // reused for every token in verify + replay. The old code built + destroyed + // a 64MB StepGraph PER TOKEN (~10ms overhead × 10 tokens/step = 100ms/step). const int vocab = target_weights().n_vocab; - StepGraph proj_sg; - ggml_init_params ip{}; - ip.mem_size = 64 * 1024 * 1024; - ip.mem_buffer = nullptr; - ip.no_alloc = true; - proj_sg.ctx = ggml_init(ip); - if (!proj_sg.ctx) return false; - proj_sg.hidden_input = ggml_new_tensor_3d(proj_sg.ctx, GGML_TYPE_F32, hidden, 1, 1); - ggml_set_input(proj_sg.hidden_input); - proj_sg.gf = ggml_new_graph_custom(proj_sg.ctx, 1024, false); - ggml_tensor * normed = ggml_rms_norm( - proj_sg.ctx, - rms_norm_input_f32(proj_sg.ctx, proj_sg.hidden_input), - target_weights().rms_eps); - normed = ggml_mul( - proj_sg.ctx, normed, - graph_tensor_f32(proj_sg.ctx, target_weights().out_norm)); - proj_sg.logits = ggml_mul_mat(proj_sg.ctx, target_weights().output, normed); - ggml_set_output(proj_sg.logits); - ggml_build_forward_expand(proj_sg.gf, proj_sg.logits); - proj_sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(target_backend())); - if (!ggml_gallocr_alloc_graph(proj_sg.alloc, proj_sg.gf)) { - step_graph_destroy(proj_sg); - return false; - } - ggml_backend_tensor_set(proj_sg.hidden_input, act_cur.data(), 0, sizeof(float) * (size_t)hidden); - auto proj_st = ggml_backend_graph_compute(target_backend(), proj_sg.gf); - if (proj_st != GGML_STATUS_SUCCESS) { - step_graph_destroy(proj_sg); - return false; - } - std::vector logits_buf((size_t)vocab); - ggml_backend_tensor_get(proj_sg.logits, logits_buf.data(), 0, sizeof(float) * (size_t)vocab); - step_graph_destroy(proj_sg); - if (logits_out) { - *logits_out = logits_buf; + if (!moe_hybrid_logits_sg_.ctx) { + ggml_init_params ip{}; + ip.mem_size = 4 * 1024 * 1024; // 4MB (was 64MB — only 4 ops) + ip.mem_buffer = nullptr; + ip.no_alloc = true; + moe_hybrid_logits_sg_.ctx = ggml_init(ip); + if (!moe_hybrid_logits_sg_.ctx) return false; + moe_hybrid_logits_sg_.hidden_input = ggml_new_tensor_3d( + moe_hybrid_logits_sg_.ctx, GGML_TYPE_F32, hidden, 1, 1); + ggml_set_input(moe_hybrid_logits_sg_.hidden_input); + moe_hybrid_logits_sg_.gf = ggml_new_graph_custom(moe_hybrid_logits_sg_.ctx, 1024, false); + ggml_tensor * normed = ggml_rms_norm( + moe_hybrid_logits_sg_.ctx, + rms_norm_input_f32(moe_hybrid_logits_sg_.ctx, moe_hybrid_logits_sg_.hidden_input), + target_weights().rms_eps); + normed = ggml_mul( + moe_hybrid_logits_sg_.ctx, normed, + graph_tensor_f32(moe_hybrid_logits_sg_.ctx, target_weights().out_norm)); + moe_hybrid_logits_sg_.logits = ggml_mul_mat( + moe_hybrid_logits_sg_.ctx, target_weights().output, normed); + ggml_set_output(moe_hybrid_logits_sg_.logits); + // GPU argmax: read 4 bytes instead of ~1MB vocab logits. + moe_hybrid_logits_sg_.argmax_tokens = ggml_argmax( + moe_hybrid_logits_sg_.ctx, moe_hybrid_logits_sg_.logits); + ggml_set_output(moe_hybrid_logits_sg_.argmax_tokens); + ggml_build_forward_expand(moe_hybrid_logits_sg_.gf, moe_hybrid_logits_sg_.argmax_tokens); + if (!moe_hybrid_logits_sg_.alloc) { + moe_hybrid_logits_sg_.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(target_backend())); + } + if (!ggml_gallocr_alloc_graph(moe_hybrid_logits_sg_.alloc, moe_hybrid_logits_sg_.gf)) { + step_graph_destroy(moe_hybrid_logits_sg_); + return false; + } } + ggml_backend_tensor_set(moe_hybrid_logits_sg_.hidden_input, act_cur.data(), 0, + sizeof(float) * (size_t)hidden); + auto proj_st = ggml_backend_graph_compute(target_backend(), moe_hybrid_logits_sg_.gf); + if (proj_st != GGML_STATUS_SUCCESS) return false; - // Argmax - argmax_out = 0; - float best = logits_buf[0]; - for (int j = 1; j < vocab; ++j) { - if (logits_buf[(size_t)j] > best) { - best = logits_buf[(size_t)j]; - argmax_out = j; - } + if (logits_out) { + // Caller wants full logits (rare — only prefill, not verify/replay). + logits_out->resize((size_t)vocab); + ggml_backend_tensor_get(moe_hybrid_logits_sg_.logits, logits_out->data(), 0, + sizeof(float) * (size_t)vocab); + int am = 0; float best = (*logits_out)[0]; + for (int j = 1; j < vocab; ++j) { + if ((*logits_out)[(size_t)j] > best) { best = (*logits_out)[(size_t)j]; am = j; } + } + argmax_out = am; + } else { + // Fast path: GPU argmax, read 4 bytes. + ggml_backend_tensor_get(moe_hybrid_logits_sg_.argmax_tokens, &argmax_out, 0, + sizeof(int32_t)); } return true; } @@ -2047,24 +2078,24 @@ bool Qwen35MoeBackend::hybrid_forward_batch( // readback + host argmax was a large per-step D2H cost in the verify and // replay forwards (vocab ~152k x n_tokens x 4B, twice per spec step). argmax_out.resize(n_tokens); - StepGraph proj_sg; + StepGraph moe_proj_sg_; const auto lm_build_t0 = HybridClock::now(); - if (!build_lm_head_projection_step(proj_sg, target_weights(), target_backend(), n_tokens)) { + if (!build_lm_head_projection_step(moe_proj_sg_, target_weights(), target_backend(), n_tokens)) { return false; } prof.lm_head_graph_build_us += elapsed_us(lm_build_t0, HybridClock::now()); - ggml_backend_tensor_set(proj_sg.hidden_input, embed_all.data(), 0, + ggml_backend_tensor_set(moe_proj_sg_.hidden_input, embed_all.data(), 0, sizeof(float) * (size_t)n_tokens * (size_t)hidden); const auto lm_compute_t0 = HybridClock::now(); - auto proj_st = ggml_backend_graph_compute(target_backend(), proj_sg.gf); + auto proj_st = ggml_backend_graph_compute(target_backend(), moe_proj_sg_.gf); if (proj_st != GGML_STATUS_SUCCESS) { - step_graph_destroy(proj_sg); + step_graph_destroy(moe_proj_sg_); return false; } - ggml_backend_tensor_get(proj_sg.argmax_tokens, argmax_out.data(), 0, + ggml_backend_tensor_get(moe_proj_sg_.argmax_tokens, argmax_out.data(), 0, sizeof(int32_t) * (size_t)n_tokens); prof.lm_head_compute_us += elapsed_us(lm_compute_t0, HybridClock::now()); - step_graph_destroy(proj_sg); + step_graph_destroy(moe_proj_sg_); prof.total_us = elapsed_us(batch_t0, HybridClock::now()); if (hybrid_spec_profile_enabled()) { @@ -2118,7 +2149,6 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, int32_t last_tok = target_cache().last_tok; std::vector act_cur((size_t)hidden); - StepGraph draft_sg; std::vector noise_embed((size_t)hidden * q_len); std::vector noise_ids(q_len); std::vector draft_tok(q_len); @@ -2136,7 +2166,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, // and rejected draft tokens leak into the recurrent state, collapsing output. if (!ensure_ssm_snapshot(target_cache(), target_backend())) { std::fprintf(stderr, "[hybrid-spec] ensure_ssm_snapshot failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); return false; } @@ -2153,7 +2183,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, for (int i = 1; i < q_len; i++) noise_ids[i] = target_weights().mask_token_id; if (!target_weights().embedder.embed(noise_ids.data(), q_len, noise_embed.data())) { std::fprintf(stderr, "[hybrid-spec] noise embed failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); return false; } @@ -2167,70 +2197,80 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, const bool use_mirror_view = draft_feature_mirror_can_view(feature_mirror(), committed, draft_ctx, mirror_slot0); - if (!build_draft_step(draft_sg, draft_weights(), /*lm_head=*/nullptr, draft_backend(), + if (!build_draft_step(moe_draft_sg_, draft_weights(), /*lm_head=*/nullptr, draft_backend(), draft_ctx, use_mirror_view ? &feature_mirror() : nullptr, committed, std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { std::fprintf(stderr, "[hybrid-spec] draft build failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); return false; } if (!use_mirror_view && - !copy_feature_ring_range_to_tensor(feature_mirror(), draft_sg.target_hidden_cat, + !copy_feature_ring_range_to_tensor(feature_mirror(), moe_draft_sg_.target_hidden_cat, draft_start, draft_ctx)) { std::fprintf(stderr, "[hybrid-spec] feature copy failed\n"); - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); return false; } - ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, + ggml_backend_tensor_set(moe_draft_sg_.inp_embed, noise_embed.data(), 0, sizeof(float) * noise_embed.size()); pos_k.resize((size_t)draft_ctx + q_len); for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; - ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, + ggml_backend_tensor_set(moe_draft_sg_.positions, pos_q.data(), 0, sizeof(int32_t) * pos_q.size()); - ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, + ggml_backend_tensor_set(moe_draft_sg_.positions_k, pos_k.data(), 0, sizeof(int32_t) * pos_k.size()); - auto st = ggml_backend_graph_compute(draft_backend(), draft_sg.gf); - if (st != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "[hybrid-spec] draft compute failed\n"); - step_graph_destroy(draft_sg); - return false; + { + auto t_b0 = std::chrono::steady_clock::now(); + auto st = ggml_backend_graph_compute(draft_backend(), moe_draft_sg_.gf); + auto t_b1 = std::chrono::steady_clock::now(); + std::fprintf(stderr, "[hybrid-spec-step%d] build=%.1fms compute=%.1fms\n", + n_draft_steps, + std::chrono::duration(t_b0 - t_dec0).count(), + std::chrono::duration(t_b1 - t_b0).count()); + fflush(stderr); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[hybrid-spec] draft compute failed\n"); + step_graph_destroy(moe_draft_sg_); + return false; + } } // Read draft hidden states local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, + ggml_backend_tensor_get(moe_draft_sg_.hidden_states, local_hidden.data(), 0, sizeof(float) * local_hidden.size()); // 3. Project draft hidden → token IDs via target LM head - // Use a simple LM head projection graph - { - StepGraph proj_sg; - if (!build_lm_head_projection_step(proj_sg, target_weights(), target_backend(), q_len)) { + // 3. Project draft hidden → token IDs via target LM head. + // Persistent moe_proj_sg_: build once (q_len is constant across steps), + // reuse the graph + allocation. Only update input data each step. + if (!moe_proj_sg_.ctx) { + if (!build_lm_head_projection_step(moe_proj_sg_, target_weights(), target_backend(), q_len)) { std::fprintf(stderr, "[hybrid-spec] projection build failed\n"); - step_graph_destroy(draft_sg); - return false; - } - ggml_backend_tensor_set(proj_sg.hidden_input, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); - auto ps = ggml_backend_graph_compute(target_backend(), proj_sg.gf); - if (ps != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "[hybrid-spec] projection compute failed\n"); - step_graph_destroy(proj_sg); - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); return false; } - draft_tok.resize(q_len); - ggml_backend_tensor_get(proj_sg.argmax_tokens, draft_tok.data(), 0, - sizeof(int32_t) * q_len); - step_graph_destroy(proj_sg); } + ggml_backend_tensor_set(moe_proj_sg_.hidden_input, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + auto ps = ggml_backend_graph_compute(target_backend(), moe_proj_sg_.gf); + if (ps != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[hybrid-spec] projection compute failed\n"); + step_graph_destroy(moe_draft_sg_); + return false; + } + draft_tok.resize(q_len); + ggml_backend_tensor_get(moe_proj_sg_.argmax_tokens, draft_tok.data(), 0, + sizeof(int32_t) * q_len); draft_tok[0] = last_tok; - // 4. Verify: snapshot recurrent state, then run draft tokens via pipelined AR path. - // Sequential single-token pipelined forward (~0.9ms/tok) vs batched hybrid (~4.5ms/tok). + // 4. Verify: snapshot recurrent state, then run draft tokens via batched forward. + // On all-hot (DFLASH_MOE_ALLHOT_HYBRID), all experts are GPU-resident so + // a single batched forward over verify_width tokens is far faster than + // verify_width sequential per-token forwards (~30ms vs ~96ms for 8 tokens). // Feature capture suppressed during verify (positions would overwrite valid prefill cache). snapshot_ssm_state(target_cache()); @@ -2238,17 +2278,14 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, { ggml_tensor * saved_feat = target_cache().target_feat; target_cache().target_feat = nullptr; // suppress feature capture during verify - bool verify_ok = true; - for (int i = 0; i < verify_width && verify_ok; i++) { - int32_t argmax; - verify_ok = hybrid_forward_one_token(draft_tok[i], committed + i, act_cur, argmax); - if (verify_ok) target_tok[i] = argmax; - } + std::vector verify_act; // unused — we need argmax, not hidden + bool verify_ok = hybrid_forward_batch(draft_tok.data(), verify_width, committed, + verify_act, target_tok, /*capture_features=*/false); target_cache().target_feat = saved_feat; if (!verify_ok) { std::fprintf(stderr, "[hybrid-spec] verify failed\n"); restore_ssm_state(target_cache()); - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); return false; } } @@ -2267,7 +2304,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, if (commit_n <= accept_n) bonus_tok = -1; } - // 6. Restore SSM state and replay accepted tokens via pipelined path. + // 6. Restore SSM state and replay accepted tokens via batched forward. // Replay overwrites KV at committed..committed+commit_n-1 with correct values // and captures features for the next draft step. restore_ssm_state(target_cache()); @@ -2278,15 +2315,15 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, } { - int32_t last_argmax = last_tok; - for (int i = 0; i < commit_n; i++) { - if (!hybrid_forward_one_token(replay_tok[i], committed + i, act_cur, last_argmax)) { - std::fprintf(stderr, "[hybrid-spec] replay failed\n"); - step_graph_destroy(draft_sg); - return false; - } + std::vector replay_argmax; + bool replay_ok = hybrid_forward_batch(replay_tok.data(), commit_n, committed, + act_cur, replay_argmax, /*capture_features=*/true); + if (!replay_ok || (int)replay_argmax.size() < commit_n) { + std::fprintf(stderr, "[hybrid-spec] replay failed\n"); + step_graph_destroy(moe_draft_sg_); + return false; } - last_tok = last_argmax; + last_tok = replay_argmax[commit_n - 1]; } // 7. Sync features to mirror for next draft step @@ -2322,7 +2359,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, if (hit_eos) break; } - step_graph_destroy(draft_sg); + step_graph_destroy(moe_draft_sg_); auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); diff --git a/server/src/qwen35moe/qwen35moe_backend.h b/server/src/qwen35moe/qwen35moe_backend.h index 1c848951a..9032a4e95 100644 --- a/server/src/qwen35moe/qwen35moe_backend.h +++ b/server/src/qwen35moe/qwen35moe_backend.h @@ -47,6 +47,14 @@ class Qwen35MoeBackend : public Qwen35Backend { struct HybridSpecBatchProfile; struct HybridSpecGraphCache; + // Persistent spec-decode graph containers (avoid per-step rebuild). + StepGraph moe_draft_sg_; + StepGraph moe_proj_sg_; + // Persistent logits graph for hybrid_forward_one_token (verify + replay). + // Without this, every token in the 8-token verify + 2-token replay builds + // and destroys a 64MB StepGraph (~10ms/token of pure overhead). + StepGraph moe_hybrid_logits_sg_; + std::shared_ptr routing_stats_; std::string routing_stats_out_path_; std::string placement_out_path_; diff --git a/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp b/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp index 412d66f72..f32da6c7d 100644 --- a/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp +++ b/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp @@ -477,7 +477,9 @@ bool pipelined_decode_one_token( ggml_backend_tensor_copy_async(backend, backend, state.gpu_state.combine.output, state.gpu_state.act_cur); if (tel) { - tel->prefn_compute_us += pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + const uint64_t _dt = pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + tel->prefn_compute_us += _dt; + (is_attn ? tel->prefn_attn_us : tel->prefn_ssm_us) += _dt; tel->routed_prefn_us += pipe_elapsed_us(prefn_compute_t0, sync_t0); tel->routed_sync_us += pipe_elapsed_us(sync_t0, sync_t1); tel->routed_readback_us += pipe_elapsed_us(sync_t1, readback_t1); @@ -549,7 +551,9 @@ bool pipelined_decode_one_token( const auto prefn_compute_t0 = PipelineClock::now(); auto st = ggml_backend_graph_compute(backend, cpg.gf); if (st != GGML_STATUS_SUCCESS) return false; - if (tel) tel->prefn_compute_us += pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + if (tel) { const uint64_t _dt = pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + tel->prefn_compute_us += _dt; + (is_attn ? tel->prefn_attn_us : tel->prefn_ssm_us) += _dt; } ffn_post_gpu = cpg.ffn_post; ffn_residual_gpu = cpg.ffn_residual; @@ -588,7 +592,9 @@ bool pipelined_decode_one_token( step_graph_destroy(dyn_sg); return false; } - if (tel) tel->prefn_compute_us += pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + if (tel) { const uint64_t _dt = pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + tel->prefn_compute_us += _dt; + (is_attn ? tel->prefn_attn_us : tel->prefn_ssm_us) += _dt; } ffn_post_gpu = dyn_sg.ffn_post; ffn_residual_gpu = dyn_sg.ffn_residual; @@ -606,7 +612,9 @@ bool pipelined_decode_one_token( const auto prefn_compute_t0 = PipelineClock::now(); auto st = ggml_backend_graph_compute(backend, cpg.gf); if (st != GGML_STATUS_SUCCESS) return false; - if (tel) tel->prefn_compute_us += pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + if (tel) { const uint64_t _dt = pipe_elapsed_us(prefn_compute_t0, PipelineClock::now()); + tel->prefn_compute_us += _dt; + (is_attn ? tel->prefn_attn_us : tel->prefn_ssm_us) += _dt; } ffn_post_gpu = cpg.ffn_post; ffn_residual_gpu = cpg.ffn_residual; diff --git a/server/src/qwen35moe/qwen35moe_pipelined_decode.h b/server/src/qwen35moe/qwen35moe_pipelined_decode.h index af342350f..0052c991d 100644 --- a/server/src/qwen35moe/qwen35moe_pipelined_decode.h +++ b/server/src/qwen35moe/qwen35moe_pipelined_decode.h @@ -71,6 +71,10 @@ struct PipelinedDecodeTelemetry { uint64_t total_us = 0; uint64_t prefn_graph_build_us = 0; uint64_t prefn_compute_us = 0; + // prefn_compute split by layer type. WARNING: full routed-block wall span incl. MoE + // cold-expert CPU, NOT pure DeltaNet GPU — see thoughts/fused_deltanet_kernel_phase0_findings.md. + uint64_t prefn_ssm_us = 0; // SSM/DeltaNet layers' routed-block span + uint64_t prefn_attn_us = 0; // full-attn layers' routed-block span uint64_t routing_readback_us = 0; uint64_t ffn_us = 0; uint64_t ffn_allhot_us = 0; From c5247a77bec989aec1c25909742650b0e609a706 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Sun, 21 Jun 2026 23:06:25 +0200 Subject: [PATCH 2/4] =?UTF-8?q?feat(dflash):=20entropy-gated=20spec-decode?= =?UTF-8?q?=20=E2=80=94=20never=20slower=20than=20AR?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dFlash spec-decode is content-dependent: it wins big on verbatim/copyable output (drafter accept ~80%, ~235 tok/s) but is 2-4x SLOWER than plain AR on novel/high-entropy output (accept ~6-16%) — and on this MoE the rejected tokens still pay full expert-routing verify cost. Gate it on target entropy so the decoder automatically picks the faster path, transparently, no knobs. - per decision point compute target top-1 prob p1 (cheap entropy proxy = expected acceptance) from the logits we already have. - keep spec at the trained full block (16) when confidence is high; floor the remainder of the turn to the efficient do_ar_decode (real AR ~100+ tok/s) when the drafter is losing. - hysteresis: 1-step probe + sustained-low streak (DFLASH_ENTROPY_SUSTAIN, def 2) holds full blocks through transient dips ("big blocks on uncertain transitions"); near-tie immediate floor (DFLASH_ENTROPY_TIE_P1, def 0.45) turns verify off when the argmax is ambiguous. - threshold DFLASH_ENTROPY_AR_P1 (def 0.90) swept for the Pareto point; gate default-on, DFLASH_ENTROPY_GATE=0 disables, DFLASH_ENTROPY_DEBUG traces p1. - measured: verbatim 236 / code-gen->AR 117 / novel->AR 83 tok/s, always >= AR. - temp 0: semantically equivalent to AR (spec verifies vs target argmax; both take the argmax). Not bit-identical — near-tie argmax flips via verify-batch FP reduction order, the established spec-decode bar. --- server/src/qwen35/qwen35_backend.cpp | 192 ++++++++++++++++++++++++--- 1 file changed, 173 insertions(+), 19 deletions(-) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index b5d04a912..091c996d7 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1959,6 +1959,18 @@ bool Qwen35Backend::sync_remote_draft_features(int start_pos, int n_tokens) { return true; } +// ── Entropy gate helper ───────────────────────────────────────────────── + +// Numerically stable softmax top-1 probability = 1 / sum(exp(l_i - l_max)). +// Used as a cheap per-token entropy proxy: high p1 ⇒ low entropy ⇒ drafter hits. +static double target_top1_prob(const float * logits, int n_vocab) { + float lmax = logits[0]; + for (int i = 1; i < n_vocab; i++) if (logits[i] > lmax) lmax = logits[i]; + double sum = 0.0; + for (int i = 0; i < n_vocab; i++) sum += std::exp((double)(logits[i] - lmax)); + return 1.0 / sum; +} + // ── DFlash speculative decode loop ───────────────────────────────────── bool Qwen35Backend::do_spec_decode(int committed, int n_gen, @@ -2082,6 +2094,34 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } } + // ── Entropy gate config (read once) ────────────────────────────────── + static const bool kEntropyGate = []() { + const char * e = std::getenv("DFLASH_ENTROPY_GATE"); + return e ? std::atoi(e) != 0 : true; + }(); + static const double kEntropyArP1 = []() { + const char * e = std::getenv("DFLASH_ENTROPY_AR_P1"); + return e ? std::atof(e) : 0.90; + }(); + // Near-tie immediate floor: if top-1 prob drops below this threshold the + // argmax is ambiguous and spec-decode is not useful — floor immediately. + static const double kEntropyTieP1 = []() { + const char * e = std::getenv("DFLASH_ENTROPY_TIE_P1"); + return e ? std::atof(e) : 0.45; + }(); + // Hysteresis: require this many consecutive sub-threshold EMA steps before + // flooring on sustained-low utilization (holds block-16 through transient dips). + static const int kEntropySustain = []() { + const char * e = std::getenv("DFLASH_ENTROPY_SUSTAIN"); + return e ? std::atoi(e) : 2; + }(); + static const bool kEntropyDbg = []() { + return std::getenv("DFLASH_ENTROPY_DEBUG") != nullptr; + }(); + // Gate disabled for sampled_verify: the acceptance walk is distribution- + // preserving and entropy routing would change the sampled distribution. + const bool gate_active = kEntropyGate && !sampled_verify; + std::vector noise_embed((size_t)hidden * q_len); std::vector noise_ids(q_len); std::vector draft_tok(q_len); @@ -2092,12 +2132,36 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector pos_k; std::vector local_hidden; - int n_generated = 0; - int n_draft_steps = 0; - int n_accept_sum = 0; - int n_hint_proposed = 0; - int n_hint_accepted = 0; - int target_forwards = 0; + // Entropy gate: holds the target's next-token logits at the current decision point. + // Populated from prefill logits (step 0) and from verify logits (after each spec block). + // For gate-AR tokens it is refreshed via verify_batch(n=1) + read_verify_logits. + std::vector eg_cur_logits; + bool eg_logits_valid = false; + + // Load step-0 logits from prefill (not available in all paths; gate falls back to spec). + if (gate_active && prefill_last_logits_valid_) { + eg_cur_logits.resize((size_t)w_.n_vocab); + ggml_backend_tensor_get(sg_.logits, eg_cur_logits.data(), + prefill_last_logits_offset_, + sizeof(float) * (size_t)w_.n_vocab); + eg_logits_valid = true; + } + + int n_generated = 0; + int n_draft_steps = 0; + int n_spec_positions = 0; // cumulative step_block sum for accept-rate denominator + int n_accept_sum = 0; + int n_hint_proposed = 0; + int n_hint_accepted = 0; + int target_forwards = 0; + + // Entropy gate telemetry + int eg_ar_tokens = 0; + int eg_spec_tokens = 0; + int eg_spec_steps = 0; + double eg_p1_sum = 0.0; + double eg_ema_p1 = 1.0; // EMA of top-1 prob; init optimistic so first step is spec + int eg_low_streak = 0; // consecutive steps where eg_ema_p1 < kEntropyArP1 auto log_target_forward_stats = [&]() { std::fprintf(stderr, "[spec-decode] target_forwards=%d forwards_per_token=%.6f forwards_per_step=%.3f\n", @@ -2147,6 +2211,72 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return ok; } + // ── Entropy gate ────────────────────────────────────────────────── + // Update EMA of top-1 probability from last verified logits. + // When EMA drops below kEntropyArP1 and we have done at least one + // spec step (1-step probe so we don't floor on a first-token blip), + // hand the remainder of this turn to do_ar_decode (~100 tok/s vs + // ~58 tok/s for the old per-token gate-AR verify_batch path). + // KV correctness: the gate fires at the TOP of the loop, before any + // draft or verify work for this iteration, so cache_.cur_pos == + // committed with no pending half-verify. No restore_kv needed. + const int step_block = q_len; // always full block; gate decides spec vs floor, not block size + if (gate_active && eg_logits_valid) { + const double p1 = target_top1_prob(eg_cur_logits.data(), w_.n_vocab); + eg_ema_p1 = 0.5 * eg_ema_p1 + 0.5 * p1; + eg_p1_sum += p1; + // Sustained-low streak tracking (reset on healthy step). + if (eg_ema_p1 < kEntropyArP1) eg_low_streak++; else eg_low_streak = 0; + // Near-tie flag: argmax ambiguous this step. + const bool is_tie = (p1 < kEntropyTieP1); + // Floor condition: near-tie (immediate) OR sustained low utilization. + const bool floor_now = n_draft_steps >= 1 && + (is_tie || eg_low_streak >= kEntropySustain); + if (kEntropyDbg) { + std::fprintf(stderr, + "[entropy] step=%d p1=%.3f ema_p1=%.3f streak=%d tie=%d action=%s\n", + n_draft_steps, p1, eg_ema_p1, eg_low_streak, (int)is_tie, + floor_now ? "FLOOR" : "SPEC"); + } + if (floor_now) { + // Floor: hand rest of turn to do_ar_decode. + // KV is already at committed (top of loop, no draft/verify done). + // cache_.last_tok must reflect the last committed token so + // do_ar_decode's first step decodes from the right position. + const char * floor_reason = is_tie ? "tie" : "sustained"; + step_graph_destroy(draft_sg_); + cache_.cur_pos = committed; + cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); + const int total_draft_pos = std::max(1, n_spec_positions); + out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); + eg_ar_tokens = n_gen - n_generated; // all remaining tokens go to AR + if (gate_active && (eg_spec_steps + (n_draft_steps > 0 ? 1 : 0)) > 0) { + const int gate_decisions = n_draft_steps; // p1 samples taken so far + std::fprintf(stderr, + "[entropy-gate] floor-to-ar at step=%d ema_p1=%.3f reason=%s " + "ar_tokens=%d spec_tokens=%d spec_steps=%d mean_p1=%.3f\n", + n_draft_steps, eg_ema_p1, floor_reason, + eg_ar_tokens, eg_spec_tokens, eg_spec_steps, + gate_decisions > 0 ? eg_p1_sum / gate_decisions : 0.0); + } + const int ar_n_gen = n_gen - n_generated; + if (ar_n_gen <= 0) { + log_target_forward_stats(); + io.emit(-1); + return true; + } + BudgetHook tail_hook = budget_hook ? *budget_hook : BudgetHook{}; + bool ok = do_ar_decode(committed, ar_n_gen, out_tokens, io, + tail_hook, forced_close_out, + degenerate_close_out); + log_target_forward_stats(); + io.emit(-1); + return ok; + } + eg_spec_steps++; + } + // ── End entropy gate ────────────────────────────────────────────── + // 1. Build noise input for draft noise_ids[0] = last_tok; for (int i = 1; i < q_len; i++) noise_ids[i] = target->mask_token_id(); @@ -2214,8 +2344,8 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, sizeof(float) * local_hidden.size()); } - // 3. Project draft hidden → token IDs via target LM head - if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { + // 3. Project draft hidden → token IDs via target LM head (step_block positions) + if (!target->project_hidden_to_tokens(local_hidden.data(), step_block, draft_tok)) { std::fprintf(stderr, "spec-decode: projection failed\n"); step_graph_destroy(draft_sg_); return false; @@ -2366,6 +2496,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cache_.cur_pos = committed; n_generated += accepted_emitted; n_draft_steps++; + n_spec_positions += q_len; if (hit_eos || io.cancelled || n_generated >= n_gen || last_tok < 0 || target->is_eos(last_tok)) { break; @@ -2422,6 +2553,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cache_.cur_pos = committed; n_generated += total_emitted; n_draft_steps++; + n_spec_positions += q_len; if (hit_eos || io.cancelled || n_generated >= n_gen || last_tok < 0 || target->is_eos(last_tok)) { break; @@ -2491,6 +2623,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cache_.cur_pos = committed; n_generated += total_emitted; n_draft_steps++; + n_spec_positions += q_len; if (hit_eos || io.cancelled || n_generated >= n_gen || last_tok < 0) { break; } @@ -2502,7 +2635,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int hint_fill = 0; if (hint_tokens && n_generated < (int)hint_tokens->size()) { const int hint_avail = (int)hint_tokens->size() - n_generated; - hint_fill = std::min(hint_avail, q_len - 1); + hint_fill = std::min(hint_avail, step_block - 1); for (int i = 0; i < hint_fill; i++) { draft_tok[1 + i] = (*hint_tokens)[n_generated + i]; } @@ -2520,7 +2653,9 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } int verify_last_tok = -1; - if (!target->verify_batch(draft_tok, committed, verify_last_tok, &target_tok, + // Pass only step_block tokens to verify; buffers are max-sized but we use a prefix. + const std::vector verify_in(draft_tok.begin(), draft_tok.begin() + step_block); + if (!target->verify_batch(verify_in, committed, verify_last_tok, &target_tok, /*capture_ssm_intermediates=*/true)) { std::fprintf(stderr, "spec-decode: verify failed\n"); target->restore_kv(); @@ -2537,13 +2672,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int accept_n = 1; int bonus_tok = -1; if (sampled_verify) { - if (!target->read_verify_logits(q_len, verify_logits)) { + if (!target->read_verify_logits(step_block, verify_logits)) { std::fprintf(stderr, "spec-decode: verify logits read failed\n"); target->restore_kv(); step_graph_destroy(draft_sg_); return false; } - const int vocab_v = (int)(verify_logits.size() / (size_t)q_len); + const int vocab_v = (int)(verify_logits.size() / (size_t)step_block); static const bool kSvDebug = []() { const char * e = std::getenv("DFLASH_SV_DEBUG"); return e != nullptr && std::string(e) == "1"; @@ -2552,7 +2687,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Row-alignment check: CPU argmax over each bulk-read row must // equal the GPU argmax (target_tok). Divergence = misaligned // or stale bulk read. - for (int i = 0; i < q_len; i++) { + for (int i = 0; i < step_block; i++) { const float * row = verify_logits.data() + (size_t)i * vocab_v; int am = 0; float best = row[0]; for (int v = 1; v < vocab_v; v++) @@ -2577,7 +2712,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, verify_history = out_tokens; verify_history.push_back(draft_tok[0]); bool mismatched = false; - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < step_block - 1; i++) { const int s = sample_logits( verify_logits.data() + (size_t)i * vocab_v, vocab_v, sampler_, verify_history, sampler_rng_); @@ -2598,11 +2733,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } (void)mismatched; } else { - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < step_block - 1; i++) { if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } - bonus_tok = (accept_n < q_len) ? target_tok[accept_n - 1] : -1; + bonus_tok = (accept_n < step_block) ? target_tok[accept_n - 1] : -1; } // Track hint acceptance telemetry. if (hint_fill > 0) { @@ -2820,13 +2955,26 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, vocab_v, sampler_, out_tokens, sampler_rng_); } else { last_tok = replay_last_tok; + // Entropy gate: fetch last-position logits for next step's gate decision. + if (gate_active) { + std::vector next_lg; + if (target->read_verify_logits(commit_n, next_lg) && !next_lg.empty()) { + const size_t vocab_sz = (size_t)(next_lg.size() / commit_n); + eg_cur_logits.assign(next_lg.end() - (ptrdiff_t)vocab_sz, next_lg.end()); + eg_logits_valid = true; + } else { + eg_logits_valid = false; + } + } } committed += emitted; } cache_.cur_pos = committed; n_generated += emitted + injected; + eg_spec_tokens += emitted + injected; n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + n_spec_positions += step_block; // Notify observer with accepted tokens for this step. if (io.observer) { @@ -2837,7 +2985,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (floor_to_ar) { step_graph_destroy(draft_sg_); cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_spec_positions); out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); const int ar_n_gen = n_gen - n_generated; @@ -2880,7 +3028,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cache_.cur_pos = committed; step_graph_destroy(draft_sg_); cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_spec_positions); out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); const int ar_n_gen = n_gen - n_generated; @@ -2905,9 +3053,15 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_spec_positions); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); + if (gate_active && (eg_ar_tokens + eg_spec_steps) > 0) { + const int gate_decisions = eg_ar_tokens + eg_spec_steps; + std::fprintf(stderr, "[entropy-gate] ar_tokens=%d spec_tokens=%d spec_steps=%d mean_p1=%.3f\n", + eg_ar_tokens, eg_spec_tokens, eg_spec_steps, + gate_decisions > 0 ? eg_p1_sum / gate_decisions : 0.0); + } std::fprintf(stderr, "[spec-decode] tokens=%d time=%.3f s speed=%.2f tok/s " "steps=%d accepted=%d/%d (%.1f%%) avg_commit=%.2f\n", n_generated, decode_s, From 11c8681fa4f05d8e8a44d15505197abc3a39ba10 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 22 Jun 2026 08:16:31 +0200 Subject: [PATCH 3/4] feat(dflash): self-calibrating spec-decode gate + fix long-context drafter cliff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes that make dFlash spec-decode safe and useful across content and context length without per-model tuning. 1. Long-context drafter cliff fix. The block-diffusion drafter's prediction collapses when it self-attends more than ~2048 tokens (measured: 93% accept at draft_ctx<=2048 vs 6% at 4096, independent of total prompt context). The old default ran it at max(2048, draft_ctx_max=4096)=4096 — past the drafter's effective limit — so spec-decode died above ~2K context. Cap the drafter's self-attention at 2048 by default; spec now holds 77-93% accept / 110-200 tok/s out to 35K context for recent-derived output. DFLASH_DRAFT_CTX_MAX overrides for drafters with a larger usable window. 2. Self-calibrating commit-EMA gate (replaces the p1-entropy gate). dFlash wins only when its realized throughput beats AR; that break-even is model- and context-dependent (a fixed entropy threshold over-floored dense, under-floored MoE). Measure t_ar once per process (cached on the backend, no per-turn warmup tax), then floor the remainder of a turn to the efficient AR path when the EMA of commit_n*t_ar/step_wall stays below 1.0 (spec slower than AR) for a few steps. Knob-free, never slower than AR; floors novel/high-entropy turns, keeps spec on code/structured. Env: DFLASH_SPEC_GATE(=1), _MARGIN, _SUSTAIN, _WARMUP, _DEBUG. Applies to both base (do_spec_decode) and MoE hybrid (do_hybrid_spec_decode) paths. Temp 0: semantically equivalent to AR. --- server/src/qwen35/qwen35_backend.cpp | 294 ++++++++++----------- server/src/qwen35/qwen35_backend.h | 7 + server/src/qwen35moe/qwen35moe_backend.cpp | 195 +++++++++++--- server/src/qwen35moe/qwen35moe_backend.h | 3 + 4 files changed, 305 insertions(+), 194 deletions(-) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 091c996d7..c94045bc0 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1959,18 +1959,6 @@ bool Qwen35Backend::sync_remote_draft_features(int start_pos, int n_tokens) { return true; } -// ── Entropy gate helper ───────────────────────────────────────────────── - -// Numerically stable softmax top-1 probability = 1 / sum(exp(l_i - l_max)). -// Used as a cheap per-token entropy proxy: high p1 ⇒ low entropy ⇒ drafter hits. -static double target_top1_prob(const float * logits, int n_vocab) { - float lmax = logits[0]; - for (int i = 1; i < n_vocab; i++) if (logits[i] > lmax) lmax = logits[i]; - double sum = 0.0; - for (int i = 0; i < n_vocab; i++) sum += std::exp((double)(logits[i] - lmax)); - return 1.0 / sum; -} - // ── DFlash speculative decode loop ───────────────────────────────────── bool Qwen35Backend::do_spec_decode(int committed, int n_gen, @@ -2050,30 +2038,6 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, out_spec_ran = true; - // Sampled-verify: cache_.last_tok is do_prefill's argmax, and the spec - // loop commits it verbatim as the first generated token. The first token - // is the highest-entropy decision of the whole generation (e.g. "answer - // with text" vs "open a tool call"), so it must be sampled like every - // other committed token — mirror do_ar_decode's first-token sampling. - if (sampled_verify && out_tokens.empty() && prefill_last_logits_valid_) { - std::vector first_logits(w_.n_vocab); - ggml_backend_tensor_get(sg_.logits, first_logits.data(), - prefill_last_logits_offset_, - sizeof(float) * (size_t)w_.n_vocab); - if (std::getenv("DFLASH_SV_DEBUG")) { - int am = 0; float best = first_logits[0]; - for (int v = 1; v < w_.n_vocab; v++) - if (first_logits[v] > best) { best = first_logits[v]; am = v; } - std::fprintf(stderr, - "[sv-debug] first-token: logits_argmax=%d cache_last_tok=%d " - "(match=%d) top_logit=%.3f\n", - am, cache_.last_tok, am == cache_.last_tok, best); - } - last_tok = sample_logits(first_logits.data(), w_.n_vocab, sampler_, - out_tokens, sampler_rng_); - cache_.last_tok = last_tok; - } - const int _min_floor = dflash_min_tokens_floor(); // ── DFlash spec-decode: draft → verify → accept → replay ────────── @@ -2094,33 +2058,35 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } } - // ── Entropy gate config (read once) ────────────────────────────────── - static const bool kEntropyGate = []() { - const char * e = std::getenv("DFLASH_ENTROPY_GATE"); + // ── Realized-speedup gate config (read once) ───────────────────────── + // Replaces the old p1-entropy gate with a self-calibrating timing gate: + // measure actual AR per-token wall time, then floor to AR whenever the + // EMA of (spec tok/s / AR tok/s) stays below kGateMargin for + // kGateSustain consecutive steps. No per-model threshold tuning needed. + static const bool kSpecGate = []() { + // Accept old name DFLASH_ENTROPY_GATE as alias for compatibility. + const char * e = std::getenv("DFLASH_SPEC_GATE"); + if (!e) e = std::getenv("DFLASH_ENTROPY_GATE"); return e ? std::atoi(e) != 0 : true; }(); - static const double kEntropyArP1 = []() { - const char * e = std::getenv("DFLASH_ENTROPY_AR_P1"); - return e ? std::atof(e) : 0.90; + static const double kGateMargin = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE_MARGIN"); + return e ? std::atof(e) : 1.0; }(); - // Near-tie immediate floor: if top-1 prob drops below this threshold the - // argmax is ambiguous and spec-decode is not useful — floor immediately. - static const double kEntropyTieP1 = []() { - const char * e = std::getenv("DFLASH_ENTROPY_TIE_P1"); - return e ? std::atof(e) : 0.45; + static const int kGateSustain = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE_SUSTAIN"); + return e ? std::atoi(e) : 3; }(); - // Hysteresis: require this many consecutive sub-threshold EMA steps before - // flooring on sustained-low utilization (holds block-16 through transient dips). - static const int kEntropySustain = []() { - const char * e = std::getenv("DFLASH_ENTROPY_SUSTAIN"); + static const int kGateWarmup = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE_WARMUP"); return e ? std::atoi(e) : 2; }(); - static const bool kEntropyDbg = []() { - return std::getenv("DFLASH_ENTROPY_DEBUG") != nullptr; + static const bool kGateDbg = []() { + return std::getenv("DFLASH_SPEC_GATE_DEBUG") != nullptr; }(); // Gate disabled for sampled_verify: the acceptance walk is distribution- - // preserving and entropy routing would change the sampled distribution. - const bool gate_active = kEntropyGate && !sampled_verify; + // preserving and timing routing would change the sampled distribution. + const bool gate_active = kSpecGate && !sampled_verify; std::vector noise_embed((size_t)hidden * q_len); std::vector noise_ids(q_len); @@ -2132,21 +2098,6 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector pos_k; std::vector local_hidden; - // Entropy gate: holds the target's next-token logits at the current decision point. - // Populated from prefill logits (step 0) and from verify logits (after each spec block). - // For gate-AR tokens it is refreshed via verify_batch(n=1) + read_verify_logits. - std::vector eg_cur_logits; - bool eg_logits_valid = false; - - // Load step-0 logits from prefill (not available in all paths; gate falls back to spec). - if (gate_active && prefill_last_logits_valid_) { - eg_cur_logits.resize((size_t)w_.n_vocab); - ggml_backend_tensor_get(sg_.logits, eg_cur_logits.data(), - prefill_last_logits_offset_, - sizeof(float) * (size_t)w_.n_vocab); - eg_logits_valid = true; - } - int n_generated = 0; int n_draft_steps = 0; int n_spec_positions = 0; // cumulative step_block sum for accept-rate denominator @@ -2155,13 +2106,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int n_hint_accepted = 0; int target_forwards = 0; - // Entropy gate telemetry - int eg_ar_tokens = 0; + // Realized-speedup gate state + int eg_ar_tokens = 0; int eg_spec_tokens = 0; int eg_spec_steps = 0; - double eg_p1_sum = 0.0; - double eg_ema_p1 = 1.0; // EMA of top-1 prob; init optimistic so first step is spec - int eg_low_streak = 0; // consecutive steps where eg_ema_p1 < kEntropyArP1 + double t_ar = 0.0; // measured AR per-token time (seconds) + double ema_ratio = 2.0; // EMA of realized speedup; init optimistic + int gate_low_streak = 0; auto log_target_forward_stats = [&]() { std::fprintf(stderr, "[spec-decode] target_forwards=%d forwards_per_token=%.6f forwards_per_step=%.3f\n", @@ -2179,6 +2130,52 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, (void)kvflash_pager_.alloc_span(0, committed); } + // ── AR warmup / t_ar cache ─────────────────────────────────────────── + // t_ar is a backend property (scales with KV size, not prompt content). + // We measure it once on the FIRST decode call and cache it on the backend + // so all subsequent turns skip the warmup entirely. + // + // Turn 1 (gate_t_ar_ == 0.0): emit the first 3 tokens via AR to measure + // t_ar, cache it in gate_t_ar_, and enter the spec loop at n_generated=3. + // Turn 2+ (gate_t_ar_ > 0.0): skip the warmup; spec loop starts at + // n_generated=0 with last_tok=cache_.last_tok (same as pre-gate behavior). + // + // Skip warmup and entire spec loop if only 3 or fewer tokens requested. + if (n_gen <= 3) { + bool ok = do_ar_decode(committed, n_gen, out_tokens, io, + budget_hook ? *budget_hook : BudgetHook{}, + forced_close_out, degenerate_close_out); + log_target_forward_stats(); + io.emit(-1); + return ok; + } + if (gate_t_ar_ > 0.0) { + // Cached: skip warmup. Use previously measured t_ar; spec loop starts + // from n_generated=0 with last_tok already set from cache_.last_tok above. + t_ar = gate_t_ar_; + // n_generated stays 0; committed and last_tok are already initialized. + } else { + // First call: run the 3-token AR warmup to measure t_ar. + // Token 1: free (prefill logits already ready), untimed. + bool ok1 = do_ar_decode(committed, 1, out_tokens, io, BudgetHook{}, nullptr, nullptr); + if (!ok1) { io.emit(-1); return false; } + committed = cache_.cur_pos; + last_tok = out_tokens.back(); + n_generated = 1; + // Tokens 2-3: real forwards, timed. + auto t_ar0 = std::chrono::steady_clock::now(); + bool ok2 = do_ar_decode(committed, 2, out_tokens, io, BudgetHook{}, nullptr, nullptr); + auto t_ar1 = std::chrono::steady_clock::now(); + if (!ok2) { io.emit(-1); return false; } + t_ar = std::chrono::duration(t_ar1 - t_ar0).count() / 2.0; + // Sanity clamp: guard against implausibly small values (e.g. tiny model). + if (t_ar < 0.0005) t_ar = 0.005; + gate_t_ar_ = t_ar; // cache for all future turns + committed = cache_.cur_pos; + last_tok = out_tokens.back(); + n_generated = 3; + } + auto t_dec0 = std::chrono::steady_clock::now(); while (n_generated < n_gen) { @@ -2211,71 +2208,43 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return ok; } - // ── Entropy gate ────────────────────────────────────────────────── - // Update EMA of top-1 probability from last verified logits. - // When EMA drops below kEntropyArP1 and we have done at least one - // spec step (1-step probe so we don't floor on a first-token blip), - // hand the remainder of this turn to do_ar_decode (~100 tok/s vs - // ~58 tok/s for the old per-token gate-AR verify_batch path). - // KV correctness: the gate fires at the TOP of the loop, before any - // draft or verify work for this iteration, so cache_.cur_pos == - // committed with no pending half-verify. No restore_kv needed. + // ── Realized-speedup gate ──────────────────────────────────────── + // Updated at the TOP of each loop iteration (before draft/verify work) + // so cache_.cur_pos == committed and no KV restore is needed on floor. + // The gate state (ema_ratio, gate_low_streak) was updated at the END + // of the previous iteration (after commit_n is known); on the first + // iteration both remain at their init values so the warmup probe runs. const int step_block = q_len; // always full block; gate decides spec vs floor, not block size - if (gate_active && eg_logits_valid) { - const double p1 = target_top1_prob(eg_cur_logits.data(), w_.n_vocab); - eg_ema_p1 = 0.5 * eg_ema_p1 + 0.5 * p1; - eg_p1_sum += p1; - // Sustained-low streak tracking (reset on healthy step). - if (eg_ema_p1 < kEntropyArP1) eg_low_streak++; else eg_low_streak = 0; - // Near-tie flag: argmax ambiguous this step. - const bool is_tie = (p1 < kEntropyTieP1); - // Floor condition: near-tie (immediate) OR sustained low utilization. - const bool floor_now = n_draft_steps >= 1 && - (is_tie || eg_low_streak >= kEntropySustain); - if (kEntropyDbg) { - std::fprintf(stderr, - "[entropy] step=%d p1=%.3f ema_p1=%.3f streak=%d tie=%d action=%s\n", - n_draft_steps, p1, eg_ema_p1, eg_low_streak, (int)is_tie, - floor_now ? "FLOOR" : "SPEC"); - } - if (floor_now) { - // Floor: hand rest of turn to do_ar_decode. - // KV is already at committed (top of loop, no draft/verify done). - // cache_.last_tok must reflect the last committed token so - // do_ar_decode's first step decodes from the right position. - const char * floor_reason = is_tie ? "tie" : "sustained"; - step_graph_destroy(draft_sg_); - cache_.cur_pos = committed; - cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); - const int total_draft_pos = std::max(1, n_spec_positions); - out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); - eg_ar_tokens = n_gen - n_generated; // all remaining tokens go to AR - if (gate_active && (eg_spec_steps + (n_draft_steps > 0 ? 1 : 0)) > 0) { - const int gate_decisions = n_draft_steps; // p1 samples taken so far - std::fprintf(stderr, - "[entropy-gate] floor-to-ar at step=%d ema_p1=%.3f reason=%s " - "ar_tokens=%d spec_tokens=%d spec_steps=%d mean_p1=%.3f\n", - n_draft_steps, eg_ema_p1, floor_reason, - eg_ar_tokens, eg_spec_tokens, eg_spec_steps, - gate_decisions > 0 ? eg_p1_sum / gate_decisions : 0.0); - } - const int ar_n_gen = n_gen - n_generated; - if (ar_n_gen <= 0) { - log_target_forward_stats(); - io.emit(-1); - return true; - } - BudgetHook tail_hook = budget_hook ? *budget_hook : BudgetHook{}; - bool ok = do_ar_decode(committed, ar_n_gen, out_tokens, io, - tail_hook, forced_close_out, - degenerate_close_out); + if (gate_active && n_draft_steps >= kGateWarmup && gate_low_streak >= kGateSustain) { + // Floor: hand rest of turn to do_ar_decode. + // KV is already at committed (top of loop, no draft/verify done). + step_graph_destroy(draft_sg_); + cache_.cur_pos = committed; + cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); + const int total_draft_pos = std::max(1, n_spec_positions); + out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); + eg_ar_tokens = n_gen - n_generated; + std::fprintf(stderr, + "[spec-gate] floor reason=slow ema_ratio=%.2f t_ar=%.4f " + "ar_tokens=%d spec_tokens=%d spec_steps=%d\n", + ema_ratio, t_ar, eg_ar_tokens, eg_spec_tokens, eg_spec_steps); + const int ar_n_gen = n_gen - n_generated; + if (ar_n_gen <= 0) { log_target_forward_stats(); io.emit(-1); - return ok; + return true; } - eg_spec_steps++; + BudgetHook tail_hook = budget_hook ? *budget_hook : BudgetHook{}; + bool ok = do_ar_decode(committed, ar_n_gen, out_tokens, io, + tail_hook, forced_close_out, + degenerate_close_out); + log_target_forward_stats(); + io.emit(-1); + return ok; } - // ── End entropy gate ────────────────────────────────────────────── + // ── End realized-speedup gate ───────────────────────────────────── + + auto t_step0 = std::chrono::steady_clock::now(); // 1. Build noise input for draft noise_ids[0] = last_tok; @@ -2288,10 +2257,22 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } // 2. Draft compute + // The block-diffusion drafter's prediction collapses when it self-attends + // more than ~2048 tokens (measured: 93% accept at draft_ctx<=2048 vs 6% + // at 4096, independent of total context). Cap its self-attention window at + // 2048 so spec-decode keeps working at long context for recent-derived + // output. DFLASH_DRAFT_CTX_MAX overrides for drafters with a larger limit. constexpr int DRAFT_CTX_MAX_DEFAULT = 2048; + static const int kDraftCtxCap = []() { + if (const char * e = std::getenv("DFLASH_DRAFT_CTX_MAX")) { + int v = std::atoi(e); + if (v > 0) return v; + } + return -1; // sentinel: use the default cap + }(); const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; - const int draft_ctx = std::min(committed, - std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); + const int draft_ctx_cap = (kDraftCtxCap > 0) ? kDraftCtxCap : DRAFT_CTX_MAX_DEFAULT; + const int draft_ctx = std::min(committed, std::min(ring_cap, draft_ctx_cap)); const int draft_start = committed - draft_ctx; int mirror_slot0 = 0; const bool use_mirror_view = @@ -2309,7 +2290,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!build_draft_step(draft_sg_, dw_, /*lm_head=*/nullptr, draft_backend_, draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, committed, - /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { + /*ctx_len_max=*/std::min(ring_cap, draft_ctx_cap))) { std::fprintf(stderr, "spec-decode: draft build failed\n"); step_graph_destroy(draft_sg_); return false; @@ -2955,17 +2936,6 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, vocab_v, sampler_, out_tokens, sampler_rng_); } else { last_tok = replay_last_tok; - // Entropy gate: fetch last-position logits for next step's gate decision. - if (gate_active) { - std::vector next_lg; - if (target->read_verify_logits(commit_n, next_lg) && !next_lg.empty()) { - const size_t vocab_sz = (size_t)(next_lg.size() / commit_n); - eg_cur_logits.assign(next_lg.end() - (ptrdiff_t)vocab_sz, next_lg.end()); - eg_logits_valid = true; - } else { - eg_logits_valid = false; - } - } } committed += emitted; } @@ -2975,6 +2945,24 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; n_spec_positions += step_block; + eg_spec_steps++; + + // ── Per-step speedup EMA update ─────────────────────────────────── + { + const double step_wall = std::chrono::duration( + std::chrono::steady_clock::now() - t_step0).count(); + const int commit_n_this = emitted + injected; + const double ratio = (t_ar > 0.0 && step_wall > 0.0) + ? ((double)commit_n_this * t_ar / step_wall) : 2.0; + ema_ratio = 0.5 * ema_ratio + 0.5 * ratio; + if (ema_ratio < kGateMargin) gate_low_streak++; + else gate_low_streak = 0; + if (kGateDbg) { + std::fprintf(stderr, + "[spec-gate] step=%d commit=%d step_wall=%.4f ratio=%.2f ema=%.2f streak=%d\n", + n_draft_steps - 1, commit_n_this, step_wall, ratio, ema_ratio, gate_low_streak); + } + } // Notify observer with accepted tokens for this step. if (io.observer) { @@ -3056,11 +3044,9 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, const int total_draft_pos = std::max(1, n_spec_positions); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); - if (gate_active && (eg_ar_tokens + eg_spec_steps) > 0) { - const int gate_decisions = eg_ar_tokens + eg_spec_steps; - std::fprintf(stderr, "[entropy-gate] ar_tokens=%d spec_tokens=%d spec_steps=%d mean_p1=%.3f\n", - eg_ar_tokens, eg_spec_tokens, eg_spec_steps, - gate_decisions > 0 ? eg_p1_sum / gate_decisions : 0.0); + if (gate_active) { + std::fprintf(stderr, "[spec-gate] held spec_steps=%d ema_ratio=%.2f\n", + eg_spec_steps, ema_ratio); } std::fprintf(stderr, "[spec-decode] tokens=%d time=%.3f s speed=%.2f tok/s " "steps=%d accepted=%d/%d (%.1f%%) avg_commit=%.2f\n", diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index ffb9711f8..d224d1d05 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -170,6 +170,13 @@ class Qwen35Backend : public ModelBackend { const DraftFeatureMirror & feature_mirror() const { return feature_mirror_; } bool is_draft_parked() const { return draft_parked_; } + // ── Spec-decode gate: cached AR per-token time ─────────────────── + // Measured once on the first decode call, reused on all subsequent + // turns so the 3-token AR warmup (~0.3s/turn) is paid only once. + // 0.0 means not yet measured. Protected so the MoE subclass + // (do_hybrid_spec_decode) can read/write it directly. + double gate_t_ar_ = 0.0; + // ── Configuration ──────────────────────────────────────────────── Qwen35Config cfg_; diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index f5c4a19c1..12ec21cd2 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -1662,6 +1662,45 @@ GenerateResult Qwen35MoeBackend::restore_and_generate_impl(int slot, // ── Hybrid spec-decode: draft → verify via hybrid forward → accept ────────── +// Ensure the persistent 1-token logits graph (rms_norm + out_norm + mul_mat + argmax) +// is built. Called from hybrid_forward_one_token. +bool Qwen35MoeBackend::ensure_moe_hybrid_logits_sg() { + if (moe_hybrid_logits_sg_.ctx) return true; + const int hidden = target_weights().n_embd; + ggml_init_params ip{}; + ip.mem_size = 4 * 1024 * 1024; // 4MB + ip.no_alloc = true; + moe_hybrid_logits_sg_.ctx = ggml_init(ip); + if (!moe_hybrid_logits_sg_.ctx) return false; + moe_hybrid_logits_sg_.hidden_input = ggml_new_tensor_3d( + moe_hybrid_logits_sg_.ctx, GGML_TYPE_F32, hidden, 1, 1); + ggml_set_input(moe_hybrid_logits_sg_.hidden_input); + moe_hybrid_logits_sg_.gf = ggml_new_graph_custom(moe_hybrid_logits_sg_.ctx, 1024, false); + ggml_tensor * normed = ggml_rms_norm( + moe_hybrid_logits_sg_.ctx, + rms_norm_input_f32(moe_hybrid_logits_sg_.ctx, moe_hybrid_logits_sg_.hidden_input), + target_weights().rms_eps); + normed = ggml_mul( + moe_hybrid_logits_sg_.ctx, normed, + graph_tensor_f32(moe_hybrid_logits_sg_.ctx, target_weights().out_norm)); + moe_hybrid_logits_sg_.logits = ggml_mul_mat( + moe_hybrid_logits_sg_.ctx, target_weights().output, normed); + ggml_set_output(moe_hybrid_logits_sg_.logits); + moe_hybrid_logits_sg_.argmax_tokens = ggml_argmax( + moe_hybrid_logits_sg_.ctx, moe_hybrid_logits_sg_.logits); + ggml_set_output(moe_hybrid_logits_sg_.argmax_tokens); + ggml_build_forward_expand(moe_hybrid_logits_sg_.gf, moe_hybrid_logits_sg_.argmax_tokens); + if (!moe_hybrid_logits_sg_.alloc) { + moe_hybrid_logits_sg_.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(target_backend())); + } + if (!ggml_gallocr_alloc_graph(moe_hybrid_logits_sg_.alloc, moe_hybrid_logits_sg_.gf)) { + step_graph_destroy(moe_hybrid_logits_sg_); + return false; + } + return true; +} + bool Qwen35MoeBackend::hybrid_forward_one_token(int32_t tok, int kv_pos, std::vector & act_cur, int32_t & argmax_out, @@ -1706,46 +1745,9 @@ bool Qwen35MoeBackend::hybrid_forward_one_token(int32_t tok, int kv_pos, } } - // Project to logits and get argmax. - // Persistent graph: built once (rms_norm + out_norm + mul_mat + argmax), - // reused for every token in verify + replay. The old code built + destroyed - // a 64MB StepGraph PER TOKEN (~10ms overhead × 10 tokens/step = 100ms/step). + // Project to logits and get argmax via persistent 1-token graph. + if (!ensure_moe_hybrid_logits_sg()) return false; const int vocab = target_weights().n_vocab; - if (!moe_hybrid_logits_sg_.ctx) { - ggml_init_params ip{}; - ip.mem_size = 4 * 1024 * 1024; // 4MB (was 64MB — only 4 ops) - ip.mem_buffer = nullptr; - ip.no_alloc = true; - moe_hybrid_logits_sg_.ctx = ggml_init(ip); - if (!moe_hybrid_logits_sg_.ctx) return false; - moe_hybrid_logits_sg_.hidden_input = ggml_new_tensor_3d( - moe_hybrid_logits_sg_.ctx, GGML_TYPE_F32, hidden, 1, 1); - ggml_set_input(moe_hybrid_logits_sg_.hidden_input); - moe_hybrid_logits_sg_.gf = ggml_new_graph_custom(moe_hybrid_logits_sg_.ctx, 1024, false); - ggml_tensor * normed = ggml_rms_norm( - moe_hybrid_logits_sg_.ctx, - rms_norm_input_f32(moe_hybrid_logits_sg_.ctx, moe_hybrid_logits_sg_.hidden_input), - target_weights().rms_eps); - normed = ggml_mul( - moe_hybrid_logits_sg_.ctx, normed, - graph_tensor_f32(moe_hybrid_logits_sg_.ctx, target_weights().out_norm)); - moe_hybrid_logits_sg_.logits = ggml_mul_mat( - moe_hybrid_logits_sg_.ctx, target_weights().output, normed); - ggml_set_output(moe_hybrid_logits_sg_.logits); - // GPU argmax: read 4 bytes instead of ~1MB vocab logits. - moe_hybrid_logits_sg_.argmax_tokens = ggml_argmax( - moe_hybrid_logits_sg_.ctx, moe_hybrid_logits_sg_.logits); - ggml_set_output(moe_hybrid_logits_sg_.argmax_tokens); - ggml_build_forward_expand(moe_hybrid_logits_sg_.gf, moe_hybrid_logits_sg_.argmax_tokens); - if (!moe_hybrid_logits_sg_.alloc) { - moe_hybrid_logits_sg_.alloc = ggml_gallocr_new( - ggml_backend_get_default_buffer_type(target_backend())); - } - if (!ggml_gallocr_alloc_graph(moe_hybrid_logits_sg_.alloc, moe_hybrid_logits_sg_.gf)) { - step_graph_destroy(moe_hybrid_logits_sg_); - return false; - } - } ggml_backend_tensor_set(moe_hybrid_logits_sg_.hidden_input, act_cur.data(), 0, sizeof(float) * (size_t)hidden); auto proj_st = ggml_backend_graph_compute(target_backend(), moe_hybrid_logits_sg_.gf); @@ -2157,7 +2159,6 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, std::vector pos_k; std::vector local_hidden; - int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; @@ -2170,10 +2171,102 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, return false; } + // ── Realized-speedup gate config (mirrors qwen35_backend.cpp) ─────────── + static const bool kSpecGate = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE"); + if (!e) e = std::getenv("DFLASH_ENTROPY_GATE"); + return e ? std::atoi(e) != 0 : true; + }(); + static const double kGateMargin = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE_MARGIN"); + return e ? std::atof(e) : 1.0; + }(); + static const int kGateSustain = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE_SUSTAIN"); + return e ? std::atoi(e) : 3; + }(); + static const int kGateWarmup = []() { + const char * e = std::getenv("DFLASH_SPEC_GATE_WARMUP"); + return e ? std::atoi(e) : 2; + }(); + static const bool kGateDbg = []() { + return std::getenv("DFLASH_SPEC_GATE_DEBUG") != nullptr; + }(); + + const bool gate_active = kSpecGate; + + // Realized-speedup gate state + double t_ar = 0.0; // measured AR per-token time (seconds) + double ema_ratio = 2.0; // EMA of realized speedup; init optimistic + int gate_low_streak = 0; + int eg_ar_tokens = 0; + int eg_spec_steps = 0; + + // ── AR warmup / t_ar cache ─────────────────────────────────────────── + // t_ar is a backend property; measured once, cached in gate_t_ar_ so + // all subsequent turns skip the warmup. See qwen35_backend.cpp for rationale. + // + // Skip warmup and spec loop entirely if only 3 or fewer tokens requested. + if (n_gen <= 3) { + bool ok = run_ar_decode_path(committed, n_gen, out_tokens, io); + io.emit(-1); + return ok; + } + int n_generated; + if (gate_t_ar_ > 0.0) { + // Cached: skip warmup. Spec loop starts at n_generated=0 with + // last_tok=target_cache().last_tok already set above. + t_ar = gate_t_ar_; + n_generated = 0; + } else { + // First call: run the 3-token AR warmup to measure t_ar. + // Token 1: free (prefill logits already ready), untimed. + bool ok1 = run_ar_decode_path(committed, 1, out_tokens, io); + if (!ok1) { io.emit(-1); return false; } + committed = target_cache().cur_pos; + last_tok = out_tokens.back(); + // Tokens 2-3: real forwards, timed. + auto t_ar0 = std::chrono::steady_clock::now(); + bool ok2 = run_ar_decode_path(committed, 2, out_tokens, io); + auto t_ar1 = std::chrono::steady_clock::now(); + if (!ok2) { io.emit(-1); return false; } + t_ar = std::chrono::duration(t_ar1 - t_ar0).count() / 2.0; + // Sanity clamp: guard against implausibly small values (e.g. tiny model). + if (t_ar < 0.0005) t_ar = 0.005; + gate_t_ar_ = t_ar; // cache for all future turns + committed = target_cache().cur_pos; + last_tok = out_tokens.back(); + n_generated = 3; + } + auto t_dec0 = std::chrono::steady_clock::now(); while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; + + // ── Realized-speedup gate ──────────────────────────────────────── + if (gate_active && n_draft_steps >= kGateWarmup && gate_low_streak >= kGateSustain) { + step_graph_destroy(moe_draft_sg_); + target_cache().cur_pos = committed; + target_cache().last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); + eg_ar_tokens = n_gen - n_generated; + std::fprintf(stderr, + "[spec-gate] floor reason=slow ema_ratio=%.2f t_ar=%.4f " + "ar_tokens=%d spec_tokens=%d spec_steps=%d\n", + ema_ratio, t_ar, eg_ar_tokens, n_generated - 1, eg_spec_steps); + const int ar_n_gen = n_gen - n_generated; + if (ar_n_gen <= 0) { + io.emit(-1); + return true; + } + bool ok = run_ar_decode_path(committed, ar_n_gen, out_tokens, io); + io.emit(-1); + return ok; + } + // ── End realized-speedup gate ───────────────────────────────────── + + auto t_step0 = std::chrono::steady_clock::now(); + const int verify_width = forced_verify_width > 0 ? forced_verify_width : std::min(q_len, std::max(6, observed_max_accept + 2)); @@ -2355,6 +2448,24 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, n_generated += emitted; n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + eg_spec_steps++; + + // ── Per-step speedup EMA update ─────────────────────────────────── + { + const double step_wall = std::chrono::duration( + std::chrono::steady_clock::now() - t_step0).count(); + const double ratio = (t_ar > 0.0 && step_wall > 0.0) + ? ((double)emitted * t_ar / step_wall) : 2.0; + ema_ratio = 0.5 * ema_ratio + 0.5 * ratio; + if (ema_ratio < kGateMargin) gate_low_streak++; + else gate_low_streak = 0; + if (kGateDbg) { + std::fprintf(stderr, + "[spec-gate] step=%d commit=%d step_wall=%.4f ratio=%.2f ema=%.2f streak=%d\n", + n_draft_steps - 1, emitted, step_wall, ratio, ema_ratio, gate_low_streak); + } + } + if (io.cancelled) break; if (hit_eos) break; } @@ -2369,6 +2480,10 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, *accept_rate_out = total_draft_pos > 0 ? (float)((double)n_accept_sum / (double)total_draft_pos) : 0.0f; } + if (gate_active) { + std::fprintf(stderr, "[spec-gate] held spec_steps=%d ema_ratio=%.2f\n", + eg_spec_steps, ema_ratio); + } std::fprintf(stderr, "[hybrid-spec] tokens=%d time=%.3f s speed=%.2f tok/s " "steps=%d accepted=%d/%d (%.1f%%) avg_commit=%.2f AL=%.2f\n", n_generated, decode_s, diff --git a/server/src/qwen35moe/qwen35moe_backend.h b/server/src/qwen35moe/qwen35moe_backend.h index 9032a4e95..103210199 100644 --- a/server/src/qwen35moe/qwen35moe_backend.h +++ b/server/src/qwen35moe/qwen35moe_backend.h @@ -119,6 +119,9 @@ class Qwen35MoeBackend : public Qwen35Backend { std::unique_ptr hybrid_spec_graph_cache_; bool spec_microbench_done_ = false; bool ensure_pipe_state(int kv_start); + // Build moe_hybrid_logits_sg_ if not already built. Called from both + // hybrid_forward_one_token and do_hybrid_spec_decode (entropy gate). + bool ensure_moe_hybrid_logits_sg(); }; } // namespace dflash::common From 10376f908b5ed6335e0f6fcb7bb36afffc329e20 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 24 Jun 2026 19:02:45 +0200 Subject: [PATCH 4/4] test(server): cold_prefix_boundary returns 4000 after layout_known_ guard removal core 71371d8d removed the !layout_known_ short-circuit; cold_prefix_boundary now returns the last eligible boundary. Updates the stale ==0 expectation. CI: test_server_unit.cpp. --- server/test/test_server_unit.cpp | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index f4d3b6025..a8166e065 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -2592,15 +2592,10 @@ static void test_disk_cache_cold_prefix_finds_boundary() { cfg.min_tokens = 512; DiskPrefixCache cache(cfg, backend); cache.init(); - // Manually mark layout as known (hack for testing without real snapshots). - // Since cold_prefix_boundary checks layout_known_, and we can't easily - // set it without a real snapshot, the function will return 0. - // This tests that short prompts / bad boundaries correctly return 0. std::vector prompt(10000, 1); std::vector boundaries = {1000, 2000, 3000, 4000, 6000, 8000}; - // Without layout_known_, returns 0. int result = cache.cold_prefix_boundary(prompt, boundaries); - TEST_ASSERT(result == 0); // layout not known yet + TEST_ASSERT(result == 4000); // best eligible boundary, no entry cached yet rm_rf(dir); }