From f99670411f9487530b4fe80234b593ebe2d0da23 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 | 295 +++++++++++------- server/src/qwen35moe/qwen35moe_backend.h | 12 + .../qwen35moe/qwen35moe_pipelined_decode.cpp | 16 +- .../qwen35moe/qwen35moe_pipelined_decode.h | 4 + 18 files changed, 699 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 de26810bb..1a36ae6d1 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, } @@ -351,6 +369,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 = [] @@ -391,7 +417,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 29efc543d..9d1d32958 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -140,6 +140,65 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, } // 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, @@ -241,11 +300,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 f9020c44e..61fc54c3c 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 8e93532fb..33efb52d4 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. @@ -317,6 +319,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 4e376950e..7e03c844b 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++) @@ -311,7 +339,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, @@ -1865,8 +1898,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); @@ -1915,7 +1946,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) { @@ -1938,7 +1969,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; } @@ -1957,52 +1988,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; @@ -2044,7 +2075,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). @@ -2060,7 +2091,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; } @@ -2069,7 +2100,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++; @@ -2102,7 +2133,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. @@ -2121,7 +2152,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; @@ -2137,7 +2168,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; @@ -2175,13 +2206,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++; @@ -2220,7 +2251,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; } @@ -2235,18 +2266,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; } @@ -2300,7 +2331,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; } @@ -2309,7 +2340,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++; @@ -2325,7 +2356,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); @@ -2405,7 +2436,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); @@ -2436,7 +2471,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); @@ -2445,7 +2480,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++; @@ -2460,7 +2495,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) { @@ -2545,7 +2580,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; @@ -2556,7 +2591,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++; @@ -2568,7 +2603,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++; @@ -2616,7 +2651,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 = @@ -2641,7 +2676,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; @@ -2652,14 +2687,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 = @@ -2682,7 +2717,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 742aaa329..79c3ac462 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 48b8c879e..8412d9b94 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -71,8 +71,22 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & ? std::string("hotness:") + hotness_path : std::string("uniform"); - // If all experts fit on GPU, reload with experts included - if (placement.total_hot >= out.n_layer * out.n_expert) { + // If all experts fit on GPU, reload with experts included. + // 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. + // + // 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); // Record the placement result so post_kvflash_init_gate() can disable @@ -81,6 +95,11 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & 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; @@ -1177,6 +1196,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; @@ -1285,6 +1306,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, @@ -1514,54 +1538,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; } @@ -1829,20 +1864,46 @@ 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; - if (!build_lm_head_projection_step(proj_sg, target_weights(), target_backend(), n_tokens)) { + const auto lm_build_t0 = HybridClock::now(); + if (!build_lm_head_projection_step(moe_proj_sg_, target_weights(), target_backend(), n_tokens)) { return false; } - ggml_backend_tensor_set(proj_sg.hidden_input, embed_all.data(), 0, + prof.lm_head_graph_build_us += elapsed_us(lm_build_t0, HybridClock::now()); + ggml_backend_tensor_set(moe_proj_sg_.hidden_input, embed_all.data(), 0, sizeof(float) * (size_t)n_tokens * (size_t)hidden); - auto proj_st = ggml_backend_graph_compute(target_backend(), proj_sg.gf); + const auto lm_compute_t0 = HybridClock::now(); + 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); - step_graph_destroy(proj_sg); + prof.lm_head_compute_us += elapsed_us(lm_compute_t0, HybridClock::now()); + step_graph_destroy(moe_proj_sg_); + prof.total_us = elapsed_us(batch_t0, HybridClock::now()); + + if (hybrid_spec_profile_enabled()) { + std::fprintf(stderr, + "[hybrid-spec-prof][batch] n_tokens=%d total=%.3fms prefn_build=%.3fms " + "prefn_compute=%.3fms ssm_prefn=%.3fms attn_prefn=%.3fms " + "positions=%.3fms masks=%.3fms routing_readback=%.3fms " + "moe_ffn=%.3fms feature_capture=%.3fms lm_head_build=%.3fms " + "lm_head_compute=%.3fms\n", + n_tokens, + prof.total_us / 1000.0, + prof.prefn_graph_build_us / 1000.0, + prof.prefn_compute_us / 1000.0, + prof.prefn_ssm_compute_us / 1000.0, + prof.prefn_attn_compute_us / 1000.0, + prof.position_build_us / 1000.0, + prof.mask_build_us / 1000.0, + prof.routing_readback_us / 1000.0, + prof.moe_ffn_us / 1000.0, + prof.feature_capture_us / 1000.0, + prof.lm_head_graph_build_us / 1000.0, + prof.lm_head_compute_us / 1000.0); + } return true; } @@ -1873,7 +1934,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); @@ -1891,7 +1951,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; } @@ -1908,7 +1968,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; } @@ -1922,80 +1982,97 @@ 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 ALL draft tokens batched + // 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()); target_tok.resize(verify_width); - bool verify_ok = hybrid_forward_batch( - draft_tok.data(), verify_width, committed, - act_cur, target_tok, /*capture_features=*/false); - if (!verify_ok) { - std::fprintf(stderr, "[hybrid-spec] verify failed\n"); - restore_ssm_state(target_cache()); - step_graph_destroy(draft_sg); - return false; + { + ggml_tensor * saved_feat = target_cache().target_feat; + target_cache().target_feat = nullptr; // suppress feature capture during verify + 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(moe_draft_sg_); + return false; + } } // 5. Acceptance: longest matching prefix @@ -2012,7 +2089,9 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, if (commit_n <= accept_n) bonus_tok = -1; } - // 6. Restore and replay accepted tokens + // 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()); std::vector replay_tok((size_t)commit_n); @@ -2020,15 +2099,17 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, replay_tok[i] = (i < accept_n) ? draft_tok[i] : bonus_tok; } - // Replay tokens through batched hybrid forward (captures features for next draft step) - std::vector replay_argmax; - if (!hybrid_forward_batch(replay_tok.data(), commit_n, committed, - act_cur, replay_argmax, /*capture_features=*/true)) { - 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 = replay_argmax[commit_n - 1]; } - last_tok = replay_argmax[commit_n - 1]; // 7. Sync features to mirror for next draft step if (feature_mirror().target_feat && target_cache().target_feat) { @@ -2063,7 +2144,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 a731b4f7a..c784c4135 100644 --- a/server/src/qwen35moe/qwen35moe_backend.h +++ b/server/src/qwen35moe/qwen35moe_backend.h @@ -44,6 +44,17 @@ class Qwen35MoeBackend : public Qwen35Backend { void after_target_compute(StepGraph & sg, int kv_start, int n_tokens) override; private: + 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_; + // All-hot placement signal for post_kvflash_init_gate(): set when // load_target_model takes the all-hot early-return (moe_hybrid null). bool placement_all_hot_ = false; @@ -51,6 +62,7 @@ class Qwen35MoeBackend : public Qwen35Backend { // (KVFlash redundant). When false but placement_all_hot_ is true, the pool // is what kept experts hot — the gate must NOT disable KVFlash. bool placement_all_hot_full_kv_ = false; + 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 917998a244c1d522bdaa511fedf3564933ba5dc2 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:34:05 +0200 Subject: [PATCH 2/4] feat(dflash): rope_parameters-aware Q4 drafter converter (convert_modal_dflash_to_gguf.py) resolve_rope_theta reads rope_parameters.rope_theta -> top-level rope_theta -> sys.exit(1); drops the silent 1M default that baked wrong RoPE for Qwen3.6 drafters whose true 10M lives nested under rope_parameters (missed by the old top-level-only pick(), default-baked 1M -> long-ctx spec-decode acceptance collapse). Q4_K_M output mirrors quantize_draft_q8.py (256 superblock, norm tensors stay F32). 12 pure-python tests cover the 35B-a3b nested case, 27B legacy top-level, gemma4 genuine-1M, and missing-theta fail-loud. Original convert_dflash_to_gguf.py untouched. --- .../scripts/convert_modal_dflash_to_gguf.py | 442 ++++++++++++++++++ server/scripts/test_convert_modal_rope.py | 100 ++++ 2 files changed, 542 insertions(+) create mode 100644 server/scripts/convert_modal_dflash_to_gguf.py create mode 100644 server/scripts/test_convert_modal_rope.py diff --git a/server/scripts/convert_modal_dflash_to_gguf.py b/server/scripts/convert_modal_dflash_to_gguf.py new file mode 100644 index 000000000..3f06c8745 --- /dev/null +++ b/server/scripts/convert_modal_dflash_to_gguf.py @@ -0,0 +1,442 @@ +#!/usr/bin/env python3 +""" +Convert the z-lab DFlash draft (safetensors, bf16) to a Q4_K_M GGUF. + +This is a corrected variant of convert_dflash_to_gguf.py with two fixes: + + FIX 1 — rope_theta resolution: + The original pick("rope_theta") only searched the top-level config.json. + Modern Qwen3.6 drafter repos publish theta nested under "rope_parameters": + { "rope_parameters": { "rope_theta": 10000000 } } + The original code silently fell through to the ROPE_THETA=1_000_000 default, + baking the wrong frequency base into the GGUF → spec-decode acceptance + collapse on 35B-a3b targets. This version checks rope_parameters first, + then top-level rope_theta, then sys.exit(1) — no silent wrong default. + + FIX 2 — Q4_K_M output: + Projection weights are quantized to Q4_K_M (mirrors the Q8_0 mechanism + in quantize_draft_q8.py: gguf.quantize() + raw_dtype override). + Norm weights and the dflash.hidden_norm singleton remain F32. + # ponytail: gguf.quantize(Q4_K_M) is a pure-Python reference path — + # llama-quantize would be faster for large models but is offline here. + +Usage: + PYTHONPATH=../../dflash_ggml/deps/llama.cpp/gguf-py python convert_modal_dflash_to_gguf.py \\ + models/draft/model.safetensors \\ + models/draft/draft-q4_k_m.gguf +""" + +import argparse +import json +import struct +import sys +from pathlib import Path + +import numpy as np +import gguf + + +# ────────────────────────────────────────────────────────────────────── +# Architecture defaults (27B drafter fallback when config.json absent). +# rope_theta intentionally absent — must come from config; see FIX 1. +# ────────────────────────────────────────────────────────────────────── + +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 +RMS_EPS = 1e-6 +MASK_TOKEN_ID = 248070 +BLOCK_SIZE = 16 +CTX_LEN = 32768 + +Q4_K_BLOCK_SIZE = 256 # Q4_K_M: 256 elements per superblock + + +# ────────────────────────────────────────────────────────────────────── +# FIX 1: rope_theta resolution — pure helper, testable without IO. +# ────────────────────────────────────────────────────────────────────── + +def resolve_rope_theta(config_dict: dict) -> float: + """Resolve rope_theta from a parsed config.json dict. + + Priority: + 1. config["rope_parameters"]["rope_theta"] (Qwen3.6 drafter layout) + 2. config["rope_theta"] (legacy top-level layout) + 3. sys.exit(1) — no silent default + + Returns float. Raises SystemExit on missing. + """ + c = config_dict + rp = c.get("rope_parameters", {}) + if rp.get("rope_theta") is not None: + return float(rp["rope_theta"]) + if c.get("rope_theta") is not None: + return float(c["rope_theta"]) + print( + "[error] rope_theta not found in config.json. " + "Checked: config['rope_parameters']['rope_theta'] and config['rope_theta']. " + "Add one of these fields to config.json before converting.", + file=sys.stderr, + ) + sys.exit(1) + + +# ────────────────────────────────────────────────────────────────────── +# Norm tensor predicate — shared with quantize_draft_q8.py consumers. +# ────────────────────────────────────────────────────────────────────── + +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" + ) + + +# ────────────────────────────────────────────────────────────────────── +# Architecture loader +# ────────────────────────────────────────────────────────────────────── + +def load_arch(safetensors: Path, header: dict) -> dict: + """Resolve arch scalars from config.json + tensor shapes. + + rope_theta is resolved via resolve_rope_theta() (FIX 1). + All other fields follow the same pattern as the original converter. + """ + a = dict( + hidden=HIDDEN, n_layer=N_LAYER, n_head=N_HEAD, n_head_kv=N_HEAD_KV, + head_dim=HEAD_DIM, intermediate=INTERMEDIATE, vocab=VOCAB, + n_target_layers=N_TARGET_LAYERS, + rms_eps=RMS_EPS, mask_token_id=MASK_TOKEN_ID, block_size=BLOCK_SIZE, + ctx_len=CTX_LEN, + sliding_window=0, sliding_window_pattern=[], + ) + # rope_theta intentionally absent from defaults — will be set from config + # or sys.exit(1) below. + + cfg_path = safetensors.parent / "config.json" + if cfg_path.exists(): + c = json.loads(cfg_path.read_text()) + dc = c.get("dflash_config", {}) + + def pick(*keys): + for k in keys: + if k in c and c[k] is not None: + return c[k] + return None + + 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")), + ("n_head", pick("num_attention_heads")), + ("n_head_kv", pick("num_key_value_heads")), + ("head_dim", pick("head_dim")), + ("intermediate", pick("intermediate_size")), + ("vocab", pick("vocab_size")), + ("rms_eps", pick("rms_norm_eps")), + ("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 + + # FIX 1: resolve rope_theta through scoped helper (loud on missing). + a["rope_theta"] = resolve_rope_theta(c) + + _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] + + 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("[warn] no config.json next to safetensors; using 27B defaults") + print( + "[error] cannot resolve rope_theta: no config.json found. " + "Provide a config.json with rope_theta or rope_parameters.rope_theta.", + file=sys.stderr, + ) + sys.exit(1) + + # Tensor-shape overrides (ground truth beats config for derived values). + def shape_of(st_name): + e = header.get(st_name) + return e["shape"] if e else None + + k0 = shape_of("layers.0.self_attn.k_proj.weight") + if k0 and a["n_head_kv"]: + derived_hd = k0[0] // a["n_head_kv"] + cfg_text = (cfg_path.read_text() if cfg_path.exists() else "{}") + if "head_dim" not in json.loads(cfg_text): + a["head_dim"] = derived_hd + + g0 = shape_of("layers.0.mlp.gate_proj.weight") + if g0: + a["intermediate"] = g0[0] + + fc = shape_of("fc.weight") + if fc and a["hidden"]: + a["n_target_layers"] = max(fc) // a["hidden"] + + n_blocks = 1 + max( + (int(n.split(".")[1]) for n in header + if n.startswith("layers.") and n.split(".")[1].isdigit()), + default=a["n_layer"] - 1, + ) + a["n_layer"] = n_blocks + + if k0: + exp_kv = a["n_head_kv"] * a["head_dim"] + if exp_kv != k0[0]: + 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']} rope_theta={a['rope_theta']:.0f} " + f"swa={swa_n}/{a['n_layer']} window={a.get('sliding_window', 0)}" + ) + return a + + +# ────────────────────────────────────────────────────────────────────── +# Tensor name mapping — DFlash safetensors -> llama.cpp GGUF +# ────────────────────────────────────────────────────────────────────── + +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 + + +# ────────────────────────────────────────────────────────────────────── +# safetensors I/O helpers +# ────────────────────────────────────────────────────────────────────── + +def load_safetensors_header(path: Path): + with open(path, "rb") as f: + header_size = struct.unpack(" bytes: + start, end = info["data_offsets"] + with open(path, "rb") as f: + f.seek(8 + header_size + start) + return f.read(end - start) + + +def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray: + u16 = np.frombuffer(raw, dtype=np.uint16).reshape(shape) + u32 = u16.astype(np.uint32) << 16 + return u32.view(" np.ndarray: + if dtype == "BF16": + return bf16_bytes_to_f32(raw, shape) + if dtype == "F16": + return np.frombuffer(raw, dtype="F32 {tuple(shape)}") + else: + # FIX 2: projection weights → Q4_K_M. + # Mirrors the Q8_0 mechanism in quantize_draft_q8.py: + # gguf.quantize(arr, type) + add_tensor(..., raw_dtype=type). + # Q4_K_M requires the last dim to be a multiple of 256 (superblock). + last_dim = shape[-1] + if last_dim % Q4_K_BLOCK_SIZE != 0: + print( + f"[error] {gguf_name}: last dim {last_dim} not divisible by " + f"{Q4_K_BLOCK_SIZE} (Q4_K_M superblock); cannot quantize", + file=sys.stderr, + ) + sys.exit(1) + q4_data = gguf.quantize(arr, gguf.GGMLQuantizationType.Q4_K) + writer.add_tensor(gguf_name, q4_data, + raw_dtype=gguf.GGMLQuantizationType.Q4_K) + total_out += q4_data.nbytes + ratio = q4_data.nbytes / len(raw) + print( + f"[tensor] {gguf_name:50s} {info['dtype']:4s}->Q4_K_M {tuple(shape)}" + f" ({q4_data.nbytes:,} bytes, {ratio:.1%} of src)" + ) + + print(f"\n[info] writing {args.out_gguf}") + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + print(f"[done] wrote {args.out_gguf}") + print(f"[size] source: {total_src / 1e9:.2f} GB") + print(f"[size] Q4_K_M out: {total_out / 1e9:.2f} GB") + if total_src: + print(f"[size] compression: {total_out / total_src:.1%}") + + +if __name__ == "__main__": + main() diff --git a/server/scripts/test_convert_modal_rope.py b/server/scripts/test_convert_modal_rope.py new file mode 100644 index 000000000..75c34c033 --- /dev/null +++ b/server/scripts/test_convert_modal_rope.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +""" +Pure-Python unit tests for resolve_rope_theta(). + +No torch, no gguf, no GPU needed. Imports only the pure helper from +convert_modal_dflash_to_gguf, making this safe to run offline. +""" + +import sys +import unittest + +# Import only the pure function — no IO, no gguf, no safetensors needed. +from convert_modal_dflash_to_gguf import resolve_rope_theta + + +class TestResolveRopeTheta(unittest.TestCase): + + # ── happy-path cases ────────────────────────────────────────────── + + def test_rope_parameters_scope_wins(self): + """35B-a3b layout: theta nested under rope_parameters.""" + cfg = {"rope_parameters": {"rope_theta": 10_000_000}} + self.assertEqual(resolve_rope_theta(cfg), 10_000_000.0) + + def test_rope_parameters_beats_top_level(self): + """When both scopes present, rope_parameters takes priority.""" + cfg = { + "rope_theta": 1_000_000, + "rope_parameters": {"rope_theta": 10_000_000}, + } + self.assertEqual(resolve_rope_theta(cfg), 10_000_000.0) + + def test_top_level_fallback(self): + """27B legacy layout: theta at top level, no rope_parameters.""" + cfg = {"rope_theta": 10_000_000} + self.assertEqual(resolve_rope_theta(cfg), 10_000_000.0) + + def test_genuine_1m_top_level(self): + """gemma4 case: real top-level theta=1_000_000 must resolve correctly.""" + cfg = {"rope_theta": 1_000_000} + self.assertEqual(resolve_rope_theta(cfg), 1_000_000.0) + + def test_genuine_1m_in_rope_parameters(self): + """theta=1_000_000 inside rope_parameters must also resolve correctly.""" + cfg = {"rope_parameters": {"rope_theta": 1_000_000}} + self.assertEqual(resolve_rope_theta(cfg), 1_000_000.0) + + def test_returns_float(self): + """Result is always float regardless of source type (int vs float).""" + cfg = {"rope_theta": 500_000} + result = resolve_rope_theta(cfg) + self.assertIsInstance(result, float) + + # ── fail-loud cases ─────────────────────────────────────────────── + + def test_neither_scope_exits(self): + """No rope_theta anywhere → must raise SystemExit (NOT silently return 1M).""" + cfg = {} + with self.assertRaises(SystemExit): + resolve_rope_theta(cfg) + + def test_empty_rope_parameters_exits(self): + """rope_parameters present but empty → must raise SystemExit.""" + cfg = {"rope_parameters": {}} + with self.assertRaises(SystemExit): + resolve_rope_theta(cfg) + + def test_none_top_level_exits(self): + """Explicit null at top level, no rope_parameters → must raise SystemExit.""" + cfg = {"rope_theta": None} + with self.assertRaises(SystemExit): + resolve_rope_theta(cfg) + + def test_none_in_rope_parameters_falls_back_to_top_level(self): + """Null inside rope_parameters → fall through to top-level.""" + cfg = {"rope_parameters": {"rope_theta": None}, "rope_theta": 10_000_000} + self.assertEqual(resolve_rope_theta(cfg), 10_000_000.0) + + def test_none_everywhere_exits(self): + """Null in both scopes → must raise SystemExit.""" + cfg = {"rope_parameters": {"rope_theta": None}, "rope_theta": None} + with self.assertRaises(SystemExit): + resolve_rope_theta(cfg) + + def test_silent_1m_default_is_gone(self): + """Critical regression guard: empty config must NOT return 1_000_000.""" + cfg = {} + try: + result = resolve_rope_theta(cfg) + # If we reach here, the function returned instead of exiting. + self.fail( + f"resolve_rope_theta({{}}) returned {result!r} " + f"instead of raising SystemExit — silent-1M default is back" + ) + except SystemExit: + pass # correct — loud failure, not silent wrong value + + +if __name__ == "__main__": + unittest.main(verbosity=2) From 02f2eb694ea8022b7e09ed1e4d6c5602f13f2022 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 24 Jun 2026 22:36:28 +0200 Subject: [PATCH 3/4] fix(cubic): address production P1/P2 review findings (bounded loop, OOB, UAF, atoi-UB, re-prefill fallback, pin_range, gguf type guards) --- .../scripts/convert_modal_dflash_to_gguf.py | 15 ++++--- server/src/common/kvflash_qk.h | 18 +++++++- server/src/common/moe_hybrid_ffn_eval.cpp | 38 +++++++++++++++++ server/src/draft/draft_gguf_loader.cpp | 22 +++++++--- server/src/qwen35/qwen35_backend.cpp | 41 +++++++++++++++---- server/src/qwen35/qwen35_target_graph.cpp | 12 +++++- server/src/qwen35moe/qwen35moe_backend.cpp | 39 ++++++++++++++++++ server/src/qwen35moe/qwen35moe_backend.h | 1 + 8 files changed, 163 insertions(+), 23 deletions(-) diff --git a/server/scripts/convert_modal_dflash_to_gguf.py b/server/scripts/convert_modal_dflash_to_gguf.py index 3f06c8745..033bfa303 100644 --- a/server/scripts/convert_modal_dflash_to_gguf.py +++ b/server/scripts/convert_modal_dflash_to_gguf.py @@ -33,7 +33,8 @@ from pathlib import Path import numpy as np -import gguf +# gguf is imported lazily inside main() so that resolve_rope_theta() and its +# tests remain importable without gguf installed (offline / pure-unit-test use). # ────────────────────────────────────────────────────────────────────── @@ -74,7 +75,9 @@ def resolve_rope_theta(config_dict: dict) -> float: """ c = config_dict rp = c.get("rope_parameters", {}) - if rp.get("rope_theta") is not None: + # Guard: rope_parameters may be present but null or a non-dict scalar; + # calling .get() on a non-dict would raise AttributeError. + if isinstance(rp, dict) and rp.get("rope_theta") is not None: return float(rp["rope_theta"]) if c.get("rope_theta") is not None: return float(c["rope_theta"]) @@ -176,10 +179,10 @@ def pick_dflash(*keys): print(f"[info] read arch from {cfg_path}") else: - print("[warn] no config.json next to safetensors; using 27B defaults") print( - "[error] cannot resolve rope_theta: no config.json found. " - "Provide a config.json with rope_theta or rope_parameters.rope_theta.", + "[error] no config.json found next to safetensors — cannot resolve " + "rope_theta. Provide a config.json with rope_theta or " + "rope_parameters.rope_theta.", file=sys.stderr, ) sys.exit(1) @@ -302,6 +305,8 @@ def to_f32(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: # ────────────────────────────────────────────────────────────────────── def main(): + import gguf # lazy import so pure helpers (resolve_rope_theta) work offline + ap = argparse.ArgumentParser( description="Convert DFlash draft BF16 safetensors to Q4_K_M GGUF" ) diff --git a/server/src/common/kvflash_qk.h b/server/src/common/kvflash_qk.h index f84bd85b0..a7e3ef445 100644 --- a/server/src/common/kvflash_qk.h +++ b/server/src/common/kvflash_qk.h @@ -47,7 +47,10 @@ inline void kvflash_qk_chunk_scores( const float * query, const KvFlashQkDims & d, std::vector & out, - float missing_score = -2.0f) { + float missing_score = -2.0f, + const float * seeded = nullptr, + float seeded_sentinel = -std::numeric_limits::infinity(), + int seeded_n = -1) { const int group = d.n_q_heads / d.n_kv_heads; const int n_chunks = (int)pooled_keys.size(); out.assign((size_t)n_chunks, missing_score); @@ -83,6 +86,19 @@ inline void kvflash_qk_chunk_scores( } out[(size_t)c] = acc * inv_layers; // layer-MEAN (Phase-0 config) } + // Seeded fallback: for chunks with no pooled key, use the ledger score from + // a prior turn if it is not the sentinel (i.e. it was actually scored). + // seeded_n bounds the valid range of the seeded array; chunks beyond it + // (n_chunks > seeded array length) fall back to missing_score safely. + if (seeded) { + const int seeded_limit = (seeded_n >= 0) ? seeded_n : n_chunks; + for (int c = 0; c < n_chunks; c++) { + if (!pooled_keys[(size_t)c] && c < seeded_limit && + seeded[c] != seeded_sentinel) { + out[(size_t)c] = seeded[c]; + } + } + } } } // namespace dflash::common diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index e8bddebd2..16a2fc547 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -1066,6 +1066,27 @@ static bool eval_moe_hybrid_ffn_batched_core( if (cl >= 0) { cold_sel[i] = cl; cold_wts[i] = selected_weights[i]; fp_has_cold = true; } } } + // Dummy slots (wts==0) may alias a real hot expert's local ID per token → + // ids_to_sorted_host drops entries → ASSERT in slow ggml_mul_mat_id path. + for (int t = 0; t < n_tokens; ++t) { + const int base = t * n_used; + int32_t next = 0; + for (int s = 0; s < n_used; ++s) { + if (hot_wts[base + s] > 0.0f) continue; + // Bounded search: at most n_hot_init probes. If every ID in + // [0, n_hot_init) is already taken by another slot we break and + // keep `next` as-is (duplicate), which is safe — the zero-weight + // slot is ignored by ids_to_sorted_host anyway. + int tries = 0; + while (tries < n_hot_init && + [&]{ for (int k=0; k= n_hot_init) next = 0; + ++tries; + } + hot_sel[base + s] = next++; + if (next >= n_hot_init) next = 0; + } + } CachedHotBatchedGraph & hg = storage.hot_batched_mixed[n_tokens]; const bool hg_ok = (hg.valid() && hg.n_tokens == n_tokens) @@ -1145,6 +1166,23 @@ static bool eval_moe_hybrid_ffn_batched_core( } } } + // Dummy slots (wts==0) may alias a real hot expert's local ID per token → + // ids_to_sorted_host drops entries → ASSERT in slow ggml_mul_mat_id path. + for (int t = 0; t < n_tokens; ++t) { + const int base = t * n_used; + int32_t next = 0; + for (int s = 0; s < n_used; ++s) { + if (hot_wts[base + s] > 0.0f) continue; + int tries = 0; + while (tries < n_hot_init && + [&]{ for (int k=0; k= n_hot_init) next = 0; + ++tries; + } + hot_sel[base + s] = next++; + if (next >= n_hot_init) next = 0; + } + } // ── Step 2: Build and run hot GPU graph (includes shared expert always) ── std::vector hot_partial((size_t)n_embd * (size_t)n_tokens, 0.0f); diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index 9d1d32958..edadd1bec 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -175,17 +175,27 @@ bool read_draft_capture_config(const std::string & path, n_capture = (int)gguf_get_val_u32(gctx, ntl_id); } + // Zero-initialize the output array so callers get safe values even when + // target_layer_ids is absent (P1-G: uninitialized capture IDs). + for (int k = 0; k < max_ids; k++) capture_ids[k] = 0; + // 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]; + // P2-2: verify element type is INT32 before casting (array may be INT64). + if (gguf_get_arr_type(gctx, tli_id) != GGUF_TYPE_INT32) { + std::fprintf(stderr, + "[draft-cfg] target_layer_ids array type is not INT32; skipping\n"); + } else { + 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; } - if (n > 0) n_capture = (int)n; } gguf_free(gctx); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 7e03c844b..3e379f365 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1053,15 +1053,38 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, // identity-style in the first place). const bool kvf_paged = kvflash_active() && kv_offset + prompt_len > kvflash_tokens_ - kvflash_pager_.chunk_tokens(); - if (kvf_paged && kv_offset != 0) { - std::fprintf(stderr, - "[kvflash] restored prefix (%d) + prompt (%d) exceeds pool %d; " - "pooled prefill requires a fresh request\n", - kv_offset, prompt_len, kvflash_tokens_); - set_last_error("kvflash: restore + pooled prefill unsupported"); - return -1; - } - if (kvf_paged) { + // restore-consume: kv_offset != 0 means the caller passed a suffix after + // deserializing the prefix KV; allow it in the pooled path without reset. + const bool kvf_restore_suffix = kvf_paged && kv_offset != 0; + if (kvf_restore_suffix) { + // Suffix pooled prefill: pager was seeded by deserialize(); just set + // ubatch size and skip reset. Log to distinguish from the cold path. + prefill_ubatch = kvflash_pager_.chunk_tokens(); + // Verify kv_offset is chunk-aligned (snapshot boundary contract). + // On misalignment, fall back to a full cold re-prefill rather than + // aborting the request — the KV restored from the snapshot is stale. + if (kv_offset % prefill_ubatch != 0) { + std::fprintf(stderr, + "[kvflash] restore-consume: kv_offset=%d not chunk-aligned " + "(chunk_tokens=%d) — falling back to full re-prefill\n", + kv_offset, prefill_ubatch); + kv_offset = 0; + kvflash_pager_.reset(); + if (kvflash_qk_policy_) { + kvflash_qk_pool_.reset(kvflash_qk_pool_.dims()); + kvflash_qk_pooled_upto_ = 0; + } + std::printf("[kvflash] pooled prefill (fallback): %d tokens through a %d-token pool " + "(%d-token chunks, evicting)\n", + prompt_len, kvflash_tokens_, prefill_ubatch); + std::fflush(stdout); + } else { + std::printf("[kvflash] restore-consume suffix: offset=%d suffix=%d " + "pool=%d chunk=%d\n", + kv_offset, prompt_len, kvflash_tokens_, prefill_ubatch); + std::fflush(stdout); + } + } else if (kvf_paged) { prefill_ubatch = kvflash_pager_.chunk_tokens(); kvflash_pager_.reset(); if (kvflash_qk_policy_) { diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 79c3ac462..fe64b9eae 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -195,7 +195,15 @@ bool create_target_cache_partial(const TargetWeights & w, // 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); + char * endp = nullptr; + long v = std::strtol(e, &endp, 10); + if (endp == e || *endp != '\0' || v <= 0) { + std::fprintf(stderr, + "[dflash] DFLASH_FEAT_RING_CAP='%s' is invalid or non-positive; " + "using default max_ctx=%d\n", e, max_ctx); + } else { + target_feat_cap_default = (int)std::min(v, (long)max_ctx); + } } out.target_feat_cap = std::min(max_ctx, target_feat_cap_default); if (allocate_target_feat) { @@ -1437,7 +1445,7 @@ bool snapshot_target_cache(const TargetWeights & w, return false; } - // Reuse existing buffer if shapes match (same cur_pos); otherwise reallocate. + // Reuse existing buffer if shapes match (same cur_pos AND same pooled mode). // Right-sized KV tensors use [head_dim, cur_pos, n_head_kv] — orders of // magnitude smaller than [head_dim, max_ctx, n_head_kv] for short prefixes. const bool needs_alloc = (snap.ctx == nullptr) || (snap.cur_pos != snap_pos); diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index 8412d9b94..e398ebea1 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -32,8 +32,33 @@ static uint64_t elapsed_us(HybridClock::time_point start, HybridClock::time_poin return (uint64_t) std::chrono::duration_cast(end - start).count(); } +static bool hybrid_spec_profile_enabled() { + static const bool enabled = []() { + const char * v = std::getenv("DFLASH_QWEN35MOE_SPEC_PROFILE"); + return v && std::atoi(v) != 0; + }(); + return enabled; +} + } // namespace +struct Qwen35MoeBackend::HybridSpecBatchProfile { + uint64_t total_us = 0; + uint64_t prefn_graph_build_us = 0; + uint64_t prefn_compute_us = 0; + uint64_t prefn_ssm_compute_us = 0; + uint64_t prefn_attn_compute_us = 0; + uint64_t position_build_us = 0; + uint64_t mask_build_us = 0; + uint64_t routing_readback_us = 0; + uint64_t moe_ffn_us = 0; + uint64_t feature_capture_us = 0; + uint64_t lm_head_graph_build_us = 0; + uint64_t lm_head_compute_us = 0; +}; + +struct Qwen35MoeBackend::HybridSpecGraphCache {}; + Qwen35MoeBackend::Qwen35MoeBackend(const Qwen35Config & cfg) : Qwen35Backend(cfg) {} @@ -332,6 +357,17 @@ bool Qwen35MoeBackend::rebuild_hybrid_from_placement(const MoeHybridPlacement & return true; } +bool Qwen35MoeBackend::park(const std::string & what) { + // Invalidate the persistent hybrid logits step-graph before freeing weights + // so that ensure_moe_hybrid_logits_sg() rebuilds it after unpark. Without + // this, the cached graph retains stale pointers to freed weight tensors. + const bool want_target = (what.empty() || what == "all" || what == "target"); + if (want_target) { + step_graph_destroy(moe_hybrid_logits_sg_); + } + return Qwen35Backend::park(what); +} + bool Qwen35MoeBackend::spark_bootstrap_finalize(const std::string & profile_path) { if (!spark_wants_bootstrap()) return false; std::string err; @@ -1613,6 +1649,9 @@ bool Qwen35MoeBackend::hybrid_forward_batch( std::vector & argmax_out, bool capture_features) { + HybridSpecBatchProfile prof{}; + const auto batch_t0 = HybridClock::now(); + const int hidden = target_weights().n_embd; const int n_layer = target_weights().n_layer; const int n_expert_used = target_weights().n_expert_used; diff --git a/server/src/qwen35moe/qwen35moe_backend.h b/server/src/qwen35moe/qwen35moe_backend.h index c784c4135..18a0d5e23 100644 --- a/server/src/qwen35moe/qwen35moe_backend.h +++ b/server/src/qwen35moe/qwen35moe_backend.h @@ -39,6 +39,7 @@ class Qwen35MoeBackend : public Qwen35Backend { std::vector & out_tokens, const DaemonIO & io) override; bool should_capture_moe_router() const override { return routing_stats_ != nullptr; } + bool park(const std::string & what) override; bool spark_wants_bootstrap() const override; bool spark_bootstrap_finalize(const std::string & profile_path) override; void after_target_compute(StepGraph & sg, int kv_start, int n_tokens) override; From 213fd76df5021acd691f2ba4f7a13a1b71c9d34e Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:11:37 +0200 Subject: [PATCH 4/4] fix(dflash): validate early capture layer ids --- server/scripts/quantize_draft_q8.py | 2 +- server/src/draft/draft_gguf_loader.cpp | 39 ++++++++++++++++++++++---- server/src/qwen35/qwen35_backend.cpp | 3 +- 3 files changed, 36 insertions(+), 8 deletions(-) diff --git a/server/scripts/quantize_draft_q8.py b/server/scripts/quantize_draft_q8.py index 7f147ec53..254f8f479 100644 --- a/server/scripts/quantize_draft_q8.py +++ b/server/scripts/quantize_draft_q8.py @@ -27,7 +27,7 @@ # 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, + ARCH, load_arch, map_name, load_safetensors_header, ) diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index edadd1bec..ec7fa53d1 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -182,24 +182,51 @@ bool read_draft_capture_config(const std::string & path, // 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); + int ids_count = 0; + bool ids_valid = false; if (tli_id >= 0 && gguf_get_kv_type(gctx, tli_id) == GGUF_TYPE_ARRAY) { // P2-2: verify element type is INT32 before casting (array may be INT64). if (gguf_get_arr_type(gctx, tli_id) != GGUF_TYPE_INT32) { std::fprintf(stderr, "[draft-cfg] target_layer_ids array type is not INT32; skipping\n"); } else { - 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]; + const size_t n = (size_t)gguf_get_arr_n(gctx, tli_id); + if (n == 0 || n > (size_t)max_ids) { + std::fprintf(stderr, + "[draft-cfg] target_layer_ids count %zu outside supported range 1..%d; skipping\n", + n, max_ids); + } else { + 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]; + } + ids_count = (int)n; + ids_valid = true; + if (n_capture == 0) n_capture = ids_count; } - if (n > 0) n_capture = (int)n; } } gguf_free(gctx); + if (!ids_valid) { + if (n_capture > 0) { + std::fprintf(stderr, + "[draft-cfg] n_capture=%d but no valid target_layer_ids; skipping early sync\n", + n_capture); + } + n_capture = 0; + return false; + } + if (n_capture <= 0 || n_capture > max_ids || n_capture != ids_count) { + std::fprintf(stderr, + "[draft-cfg] n_capture=%d target_layer_ids=%d mismatch; skipping early sync\n", + n_capture, ids_count); + n_capture = 0; + return false; + } + 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++) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 3e379f365..6f8afea76 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -208,7 +208,8 @@ bool Qwen35Backend::init() { 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)) { + DFLASH_MAX_CAPTURE_LAYERS) && + n_cap > 0 && n_cap <= 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];