feat(dflash): entropy + self-calibrating spec-decode gate (composes with #436)

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].

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("<f4").reshape(shape)
-        return f32.astype("<f2")
+        # Return as uint16 — the caller will write with raw_dtype=BF16.
+        return u16
     if dtype == "F16":
         return np.frombuffer(raw, dtype="<f2").reshape(shape)
     if dtype == "F32":
@@ -226,8 +244,8 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray:
 SAFETENSORS_DTYPE_TO_GGUF = {
     "F32":  gguf.GGMLQuantizationType.F32,
     "F16":  gguf.GGMLQuantizationType.F16,
-    # BF16 in safetensors -> we narrow to F16 in bytes_to_np above.
-    "BF16": gguf.GGMLQuantizationType.F16,
+    # BF16 preserved as-is for sm80+ native tensor core support.
+    "BF16": gguf.GGMLQuantizationType.BF16,
 }
 
 
@@ -284,6 +302,14 @@ def main():
         print(f"[warn] capture_layer_ids len {len(_cap_ids)} != n_target_layers "
               f"{a['n_target_layers']}; not embedding ids", file=sys.stderr)
 
+    # Sliding-window attention metadata (consumed by draft_gguf_loader.cpp).
+    # Qwen3.6 DFlash pattern: 5 sliding_attention + 1 full_attention, window=4096.
+    if a.get("sliding_window"):
+        writer.add_uint32(f"{ARCH}.attention.sliding_window", a["sliding_window"])
+    if a.get("sliding_window_pattern"):
+        writer.add_array(f"{ARCH}.attention.sliding_window_pattern",
+                         a["sliding_window_pattern"])
+
     # Walk + add tensors. Sort: dflash.* singletons first, then output_*,
     # then per-layer in numeric order — keeps the on-disk layout stable.
     pending = []
@@ -324,7 +350,13 @@ def sort_key(t):
             gguf_name == "dflash.hidden_norm.weight"
         )
         if is_norm:
-            arr = arr.astype("<f4")
+            # Norms must be F32 (ggml-cuda elementwise mul asserts on BF16/F16 src1).
+            # Convert properly from source dtype.
+            if st_dtype == "BF16":
+                u32 = (arr.astype(np.uint32) << 16)
+                arr = u32.view("<f4").reshape(shape).copy()
+            else:
+                arr = arr.astype("<f4")
             raw_dtype = gguf.GGMLQuantizationType.F32
         writer.add_tensor(gguf_name, arr, raw_dtype=raw_dtype)
         print(f"[tensor] {gguf_name:50s} {st_dtype:4s}->{raw_dtype.name:4s} {tuple(shape)}")

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("<Q", f.read(8))[0]
-        header_json = f.read(header_size).decode("utf-8")
-        return header_size, json.loads(header_json)
+Q8_0_BLOCK_SIZE = 32   # elements per Q8_0 block
 
 
 def read_tensor_bytes(path: Path, header_size: int, info: dict) -> bytes:
@@ -109,6 +47,14 @@ def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray:
     return u32.view("<f4").reshape(shape)
 
 
+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"
+    )
+
+
 # ──────────────────────────────────────────────────────────────────────
 # 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 = []