From 2ddee8d0288606dd4ae8e31016b51dd80c82bd09 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 28 Aug 2026 08:05:50 +0000 Subject: [PATCH] Export the DSpark block drafter for Qwen3.5/3.8 DSpark is a parallel block drafter with a low-rank Markov head. It exposes the same candidate lattice to the runtime as DFlash 2, so it reuses that proposal and verification path and only needs its own graph export. Add the drafter graph builder and wire it into the composite Qwen3.5 builder next to DFlash 2: `dspark_path` produces an auxiliary `dspark.onnx` plus a `dspark` section in genai_config.json, sharing the target's embedding and LM head on disk. The two block drafters are mutually exclusive, since the runtime drives one lattice per model. SpecForge names the layers whose *output* it taps, while the builder's aux_hidden_state_layers names the residual stream *entering* a layer, so the required tap set is each `target_layer_ids` entry plus one. Passing SpecForge's ids through unchanged is an off-by-one that leaves acceptance pinned at exactly 1.0, so a mismatch is rejected at init; the tests cover that case explicitly. `dspark_top_k` controls how many candidates the lattice keeps per block slot. --- src/python/py/models/README.md | 15 + src/python/py/models/builder.py | 9 + src/python/py/models/builders/dspark.py | 801 ++++++++++++++++++ src/python/py/models/builders/qwen.py | 115 ++- .../python/builder/test_qwen_dspark_export.py | 252 ++++++ 5 files changed, 1191 insertions(+), 1 deletion(-) create mode 100644 src/python/py/models/builders/dspark.py create mode 100644 test/python/builder/test_qwen_dspark_export.py diff --git a/src/python/py/models/README.md b/src/python/py/models/README.md index 9c06186b31..73d62b8cae 100644 --- a/src/python/py/models/README.md +++ b/src/python/py/models/README.md @@ -25,6 +25,7 @@ This folder contains the model builder for quickly creating optimized and quanti - [Include Auxiliary Hidden States Output](#include-auxiliary-hidden-states-output) - [Build with Paged Attention](#build-with-paged-attention) - [Build a DFlash 2 Block Drafter](#build-a-dflash-2-block-drafter) + - [Build a DSpark Block Drafter](#build-a-dspark-block-drafter) - [Disable Windowed KV Cache](#disable-windowed-kv-cache) - [Enable Shared Embeddings](#enable-shared-embeddings) - [Enable CUDA Graph Capture](#enable-cuda-graph-capture) @@ -324,6 +325,20 @@ python -m onnxruntime_genai.models.builder -i path_to_target_model -o path_to_ou python builder.py -i path_to_target_model -o path_to_output_folder -p fp16 -e cuda -c cache_dir_for_hf_files --extra_options use_paged_attention=true aux_hidden_state_layers=1,11,21 dflash2_path=path_to_dflash2_checkpoint dflash2_num_draft_tokens=4 ``` +#### Build a DSpark Block Drafter + +Set `dspark_path` to a DSpark checkpoint to export an auxiliary `dspark.onnx` block drafter beside a Qwen3.5 or Qwen3.8 target model. The target must use paged attention. SpecForge identifies the target layers whose outputs are tapped, while `aux_hidden_state_layers` identifies residual streams entering layers, so each configured auxiliary layer must be one greater than the corresponding `target_layer_ids` entry in the DSpark checkpoint. The drafter reuses the target's embedding and LM-head initializers. `dspark_path` and `dflash2_path` are mutually exclusive. + +`dspark_num_draft_tokens` optionally overrides how many tokens the drafter proposes per step and must be a positive integer. `dspark_top_k` controls how many candidates the lattice keeps per block slot; it defaults to `16` and must be a positive integer no greater than the drafter vocabulary size. + +```bash +# From wheel: +python -m onnxruntime_genai.models.builder -i path_to_target_model -o path_to_output_folder -p fp16 -e cuda -c cache_dir_for_hf_files --extra_options use_paged_attention=true aux_hidden_state_layers=1,11,21 dspark_path=path_to_dspark_checkpoint dspark_num_draft_tokens=4 dspark_top_k=16 + +# From source: +python builder.py -i path_to_target_model -o path_to_output_folder -p fp16 -e cuda -c cache_dir_for_hf_files --extra_options use_paged_attention=true aux_hidden_state_layers=1,11,21 dspark_path=path_to_dspark_checkpoint dspark_num_draft_tokens=4 dspark_top_k=16 +``` + #### Disable Windowed KV Cache By default, sliding-window layers use a reduced KV cache on supported execution providers. With paged attention, eligible local layers use a ring of blocks when the exported model also contains at least one full-context layer. Set `windowed_kv_cache=false` to give every layer a full-length KV cache, which is useful for performance comparisons or compatibility testing. The option defaults to `true` and applies to both paged and non-paged models. diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 44e5bc2778..d0e96b0026 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -757,6 +757,15 @@ def get_args(): aux_hidden_state_layers to match the drafter's `target_layer_ids`. Default is unset (disabled). dflash2_num_draft_tokens = Override the number of draft tokens the DFlash 2 block drafter proposes per step. Must be a positive integer. Default is taken from the draft checkpoint. + dspark_path = Path to a DSpark draft checkpoint. Exports an auxiliary `dspark.onnx` + block drafter beside the target model and adds a `dspark` section to + genai_config.json. Mutually exclusive with dflash2_path. Requires + use_paged_attention=true. SpecForge taps each target layer's output, so + aux_hidden_state_layers must be the drafter's `target_layer_ids` each plus one. + dspark_num_draft_tokens = Override the number of draft tokens the DSpark block + drafter proposes per step. Must be a positive integer. Default is taken from the draft checkpoint. + dspark_top_k = Candidates the DSpark lattice keeps per block slot. Must be a positive integer no + greater than the drafter vocabulary size. Default is 16. mtp_quant_config = JSON object/file: Configure MTP I/O, dense weights, MoE, and runtime using the structured QuantConfig schema independently from the main model. linear_attn_op = linear_attention/gated_delta_net: Select the recurrent operator for non-paged diff --git a/src/python/py/models/builders/dspark.py b/src/python/py/models/builders/dspark.py new file mode 100644 index 0000000000..50e4253e0e --- /dev/null +++ b/src/python/py/models/builders/dspark.py @@ -0,0 +1,801 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +"""Builder for a DSpark block-diffusion draft model (``dspark.onnx``). + +DSpark is SpecForge's DFlash backbone plus a low-rank Markov (bigram) head. Like DFlash 2 it +is a *block* drafter: one pass over the target's auxiliary hidden states proposes a whole +block of tokens. It differs from DFlash 2 in four ways that are all visible in the graph: + +1. **No dynamic convolutions.** The layers are plain pre-norm Qwen3 blocks. +2. **Checkpoint-specific full attention.** This Qwen3.8 checkpoint has only + ``full_attention`` layers, so its KV cache covers the whole context + (``local_window_size = -1``). DSpark itself can use sliding-window attention; for a + uniformly windowed checkpoint the builder forwards its positive ``sliding_window`` and + the runtime uses a fixed cache ring. +3. **YaRN RoPE** (factor 32 over an 8192-token pretraining window) rather than default RoPE. +4. **A Markov head instead of a candidate selector.** ``markov_w2[cur] . markov_w1[prev]`` is a + learned bigram bias added to the draft logits. Restricted to the per-position top-k it is + exactly the pairwise edge score of a lattice, so the runtime walks it with the same greedy + path search DFlash 2 uses -- ``draft_candidate_ids`` / ``draft_scores`` mean the same thing. + +Every block row predicts: row ``j`` carries the anchor token's embedding when ``j == 0`` and a +MASK embedding otherwise, and predicts the token at ``anchor_position + j + 1``. (DFlash 2's +row 0 predicts nothing, which is why its block is one row wider for the same draft count.) + +The drafter checkpoint ships no embedding and no LM head; both are the target's and are +emitted with the target's initializer names so ``share_external_initializers`` can fold them +back onto ``text.onnx.data``. +""" + +from __future__ import annotations + +import glob +import json +import math +import os + +import numpy as np +import onnx_ir as ir +import torch +from onnx_ir.tensor_adapters import TorchTensor, to_torch_dtype +from tqdm import tqdm + + +class DSparkBuilder: + """Emits ``dspark.onnx`` from a DSpark draft checkpoint plus the target's embedding and + LM head.""" + + def __init__( + self, + draft_dir, + target_dir, + io_dtype, + paged_block_size, + max_position_embeddings, + filename="dspark.onnx", + num_draft_tokens=None, + top_k=16, + ): + self.draft_dir = draft_dir + self.target_dir = target_dir + # The drafter is a bf16 checkpoint whose activations leave the fp16 range (the fc output + # alone reaches ~1e4), so the body runs in bf16. Only the tensors it shares with the fp16 + # target -- the aux hidden states, the embedding table and the FP8 LM head -- stay at the + # target's dtype. + self.io_dtype = ir.DataType.BFLOAT16 + self.external_dtype = io_dtype + self.filename = filename + self.paged_block_size = paged_block_size + + with open(os.path.join(draft_dir, "config.json")) as f: + cfg = json.load(f) + self.cfg = cfg + dfl = cfg.get("dflash_config") or {} + self.dflash_config = dfl + + self.hidden_size = int(cfg["hidden_size"]) + self.num_layers = int(cfg["num_hidden_layers"]) + self.num_heads = int(cfg["num_attention_heads"]) + self.num_kv_heads = int(cfg["num_key_value_heads"]) + self.head_size = int(cfg["head_dim"]) + self.intermediate_size = int(cfg["intermediate_size"]) + self.vocab_size = int(cfg["vocab_size"]) + self.rms_eps = float(cfg["rms_norm_eps"]) + self.max_position = int(max_position_embeddings or cfg["max_position_embeddings"]) + self.rope_parameters = dict(cfg["rope_parameters"]) + self.sliding_window = int(cfg["sliding_window"]) if cfg.get("use_sliding_window") else -1 + # The dual-source block attends to itself bidirectionally; DFlash's attention module + # hard-codes is_causal=False and the published configs do not override it. + self.is_causal = bool(cfg.get("is_causal", False)) + + self.mask_token_id = int(dfl["mask_token_id"]) + self.target_layer_ids = list(dfl["target_layer_ids"]) + self.markov_rank = int(cfg.get("markov_rank", dfl.get("markov_rank", 0))) + if self.markov_rank <= 0: + raise ValueError("A DSpark drafter needs markov_rank > 0.") + # One query row per draft; a smaller request just takes a prefix of the same block. + try: + self.block_size = int(num_draft_tokens) if num_draft_tokens is not None else int(cfg["block_size"]) + except (TypeError, ValueError) as error: + raise ValueError("num_draft_tokens must be a positive integer.") from error + if self.block_size < 1: + raise ValueError("num_draft_tokens must be a positive integer.") + self.num_draft_tokens = self.block_size + try: + self.top_k = int(top_k) + except (TypeError, ValueError) as error: + raise ValueError("top_k must be a positive integer.") from error + if not 1 <= self.top_k <= self.vocab_size: + raise ValueError(f"top_k must be between 1 and the vocabulary size ({self.vocab_size}).") + self.aux_hidden_size = self.hidden_size * len(self.target_layer_ids) + + self.values: dict[str, ir.Value] = {} + self.node_names: set[str] = set() + self.graph = ir.Graph( + inputs=(), + outputs=(), + nodes=(), + opset_imports={"": 21, "com.microsoft": 1}, + name="dspark_graph", + ) + self.model = ir.Model(self.graph, ir_version=10, producer_name="onnxruntime-genai") + self._const_cache: dict[str, str] = {} + + # ---------------------------------------------------------------- plumbing + + def make_value(self, name, dtype=None, shape=None): + if name == "": + return ir.Value(name="") + value = self.values.setdefault(name, ir.Value(name=name)) + if dtype is not None: + value.dtype = ir.DataType(dtype) + if shape is not None: + value.shape = ir.Shape(shape) + return value + + def make_node(self, op_type, inputs, outputs, *, name, domain="", **attrs): + if name in self.node_names: + raise ValueError(f"duplicate node name {name}") + node = ir.node( + op_type, + inputs=[self.make_value(n) for n in inputs], + attributes=attrs, + domain=domain, + outputs=[self.make_value(n) for n in outputs], + name=name, + ) + self.graph.append(node) + self.node_names.add(name) + + def make_initializer(self, tensor, name, to=None): + if to is not None: + + def tensor_func(t=tensor, dtype=to): + return TorchTensor(t.to(to_torch_dtype(dtype)).contiguous(), name=name) + + ir_tensor = ir.LazyTensor(tensor_func, dtype=to, shape=ir.Shape(tensor.shape), name=name) + elif isinstance(tensor, torch.Tensor): + ir_tensor = TorchTensor(tensor.contiguous(), name=name) + else: + ir_tensor = ir.tensor(tensor, name=name) + value = self.make_value(name, ir_tensor.dtype, ir_tensor.shape) + value.const_value = ir_tensor + self.graph.register_initializer(value) + return name + + def const(self, values, dtype=ir.DataType.INT64): + """Emit (once) a small constant. + + These are ``Constant`` nodes rather than initializers because shape inference has to read + the ones that feed ``Reshape`` / ``Slice`` / ``TopK``, and every initializer is written to + external data. + """ + arr = np.asarray(values) + key = f"{dtype}:{arr.shape}:{arr.tobytes().hex()}" + if key in self._const_cache: + return self._const_cache[key] + name = f"dspark.const.{len(self._const_cache)}" + np_dtype = { + ir.DataType.INT64: np.int64, + ir.DataType.INT32: np.int32, + ir.DataType.FLOAT: np.float32, + }[dtype] + tensor = ir.tensor(arr.astype(np_dtype), name=name) + self.make_node("Constant", [], [name], name=f"{name}/Constant", value=tensor) + self.make_value(name, dtype, tensor.shape) + self._const_cache[key] = name + return name + + # --------------------------------------------------------------- op sugar + + def _out(self, name): + return f"{name}/output_0" + + def unary(self, op, name, x, dtype, shape, domain="", **attrs): + out = self._out(name) + self.make_node(op, [x], [out], name=name, domain=domain, **attrs) + self.make_value(out, dtype, shape) + return out + + def binary(self, op, name, a, b, dtype, shape): + out = self._out(name) + self.make_node(op, [a, b], [out], name=name) + self.make_value(out, dtype, shape) + return out + + def reshape(self, name, x, shape_const, dtype, shape): + return self.binary("Reshape", name, x, self.const(shape_const), dtype, shape) + + def matmul(self, name, x, weight_tensor, in_features, out_features, rows, weight_name=None): + """``x @ weight.T`` for a torch ``[out, in]`` weight.""" + wname = weight_name or (name[1:].replace("/", ".") + ".weight") + if wname not in self.values: + self.make_initializer(weight_tensor.T, wname, to=self.io_dtype) + out = self._out(name) + self.make_node("MatMul", [x, wname], [out], name=name) + self.make_value(out, self.io_dtype, [rows, out_features]) + return out + + def rms_norm(self, name, x, weight_tensor, rows, weight_name=None): + wname = weight_name or (name[1:].replace("/", ".") + ".weight") + if wname not in self.values: + self.make_initializer(weight_tensor, wname, to=self.io_dtype) + out = self._out(name) + self.make_node( + "SimplifiedLayerNormalization", + [x, wname], + [out, "", ""], + name=name, + axis=-1, + epsilon=self.rms_eps, + stash_type=1, + ) + self.make_value(out, self.io_dtype, [rows, self.hidden_size]) + return out + + def skip_rms_norm(self, name, root, skip, weight_tensor, rows, want_sum=True): + wname = name[1:].replace("/", ".") + ".weight" + self.make_initializer(weight_tensor, wname, to=self.io_dtype) + out = self._out(name) + sm = f"{name}/output_3" + self.make_node( + "SkipSimplifiedLayerNormalization", + [root, skip, wname], + [out, "", "", sm] if want_sum else [out], + name=name, + domain="com.microsoft", + epsilon=self.rms_eps, + ) + self.make_value(out, self.io_dtype, [rows, self.hidden_size]) + if want_sum: + self.make_value(sm, self.io_dtype, [rows, self.hidden_size]) + return out, (sm if want_sum else None) + + # ------------------------------------------------------------ rope caches + + def _yarn_inv_freq(self): + """HF ``_compute_yarn_parameters`` for this checkpoint's ``rope_parameters``.""" + params = self.rope_parameters + base = float(params["rope_theta"]) + dim = int(self.head_size * float(params.get("partial_rotary_factor", 1.0))) + factor = float(params["factor"]) + original_max = float(params["original_max_position_embeddings"]) + beta_fast = float(params.get("beta_fast") or 32) + beta_slow = float(params.get("beta_slow") or 1) + + attention_factor = params.get("attention_factor") + if attention_factor is None: + attention_factor = 1.0 if factor <= 1 else 0.1 * math.log(factor) + 1.0 + + def correction_dim(rotations): + return (dim * math.log(original_max / (rotations * 2 * math.pi))) / (2 * math.log(base)) + + low, high = correction_dim(beta_fast), correction_dim(beta_slow) + if params.get("truncate", True): + low, high = math.floor(low), math.ceil(high) + low, high = max(low, 0), min(high, dim - 1) + if low == high: + high += 0.001 + + pos_freqs = base ** (np.arange(0, dim, 2, dtype=np.float32) / dim) + extrapolation = 1.0 / pos_freqs + interpolation = 1.0 / (factor * pos_freqs) + ramp = np.clip((np.arange(dim // 2, dtype=np.float32) - low) / (high - low), 0.0, 1.0) + weight = 1.0 - ramp + inv_freq = interpolation * (1.0 - weight) + extrapolation * weight + return inv_freq.astype(np.float64), float(attention_factor) + + def _make_rope_caches(self): + # PagedAttention type-constrains cos/sin to the query's element type. YaRN's attention + # factor scales cos/sin, exactly as HF's rotary embedding does, so it is baked in here. + if str(self.rope_parameters.get("rope_type", "default")).lower() in ("yarn", "longrope"): + inv_freq, attention_factor = self._yarn_inv_freq() + else: + dim = self.head_size + base = float(self.rope_parameters["rope_theta"]) + inv_freq = 1.0 / (base ** (np.arange(0, dim, 2, dtype=np.float64) / dim)) + attention_factor = 1.0 + pos = np.arange(self.max_position, dtype=np.float64)[:, None] + freqs = pos * inv_freq[None, :] + for name, fn in (("dspark.cos_cache", np.cos), ("dspark.sin_cache", np.sin)): + values = (fn(freqs) * attention_factor).astype(np.float32) + self.make_initializer(torch.from_numpy(values), name, to=self.io_dtype) + + # ------------------------------------------------------------------ graph + + def make_model(self): + w = self._load_weights() + self._declare_io() + self._make_rope_caches() + + rows_q = "num_block" + + # --- context path: the target's aux hidden states become every layer's K/V --- + aux = self.unary( + "Cast", + "/dspark/aux/Cast", + "aux_hidden_states", + self.io_dtype, + ["num_ctx", self.aux_hidden_size], + to=self.io_dtype, + ) + ctx = self.matmul("/dspark/fc/MatMul", aux, w["fc.weight"], self.aux_hidden_size, self.hidden_size, "num_ctx") + ctx_n = self.rms_norm("/dspark/hidden_norm", ctx, w["hidden_norm.weight"], "num_ctx") + ctx_kv = [] + for i in range(self.num_layers): + k = self.matmul( + f"/dspark/layers.{i}/ctx_k/MatMul", + ctx_n, + w[f"layers.{i}.self_attn.k_proj.weight"], + self.hidden_size, + self.num_kv_heads * self.head_size, + "num_ctx", + weight_name=f"dspark.layers.{i}.self_attn.k_proj.weight", + ) + v = self.matmul( + f"/dspark/layers.{i}/ctx_v/MatMul", + ctx_n, + w[f"layers.{i}.self_attn.v_proj.weight"], + self.hidden_size, + self.num_kv_heads * self.head_size, + "num_ctx", + weight_name=f"dspark.layers.{i}.self_attn.v_proj.weight", + ) + ctx_kv.append((k, v)) + + # --- query path --- + self.make_initializer(w["embed_tokens.weight"], "model.embed_tokens.weight", to=self.external_dtype) + emb_ext = self.binary( + "Gather", + "/dspark/embed_tokens/Gather", + "model.embed_tokens.weight", + "input_ids", + self.external_dtype, + [rows_q, self.hidden_size], + ) + hidden = self.unary( + "Cast", "/dspark/embed_tokens/Cast", emb_ext, self.io_dtype, [rows_q, self.hidden_size], to=self.io_dtype + ) + + residual = None + for i in range(self.num_layers): + p = f"/dspark/layers.{i}" + if residual is None: + x = self.rms_norm(f"{p}/input_layernorm", hidden, w[f"layers.{i}.input_layernorm.weight"], rows_q) + residual = hidden + else: + x, residual = self.skip_rms_norm( + f"{p}/input_layernorm", residual, hidden, w[f"layers.{i}.input_layernorm.weight"], rows_q + ) + + attn_out = self._make_attention(i, x, ctx_kv[i], rows_q) + y, residual = self.skip_rms_norm( + f"{p}/post_attention_layernorm", + residual, + attn_out, + w[f"layers.{i}.post_attention_layernorm.weight"], + rows_q, + ) + hidden = self._make_mlp(i, y, w, rows_q) + + final, _ = self.skip_rms_norm("/dspark/norm", residual, hidden, w["norm.weight"], rows_q, want_sum=False) + + self._make_candidates_and_markov(final, w) + self.graph.sort() + return self.model + + def _make_attention(self, i, x, ctx_kv, rows_q): + p = f"/dspark/layers.{i}" + w = self._weights + q = self.matmul( + f"{p}/attn/q_proj/MatMul", + x, + w[f"layers.{i}.self_attn.q_proj.weight"], + self.hidden_size, + self.num_heads * self.head_size, + rows_q, + ) + k = self.matmul( + f"{p}/attn/k_proj/MatMul", + x, + w[f"layers.{i}.self_attn.k_proj.weight"], + self.hidden_size, + self.num_kv_heads * self.head_size, + rows_q, + weight_name=f"dspark.layers.{i}.self_attn.k_proj.weight", + ) + v = self.matmul( + f"{p}/attn/v_proj/MatMul", + x, + w[f"layers.{i}.self_attn.v_proj.weight"], + self.hidden_size, + self.num_kv_heads * self.head_size, + rows_q, + weight_name=f"dspark.layers.{i}.self_attn.v_proj.weight", + ) + + # Interleave the block rows and the context rows into one packed token stream. + # `qkv_row_map` indexes concat(block, context); `q_row_map` indexes the block rows alone + # (context rows point at row 0 and their output is dropped). + kv_dim = self.num_kv_heads * self.head_size + k_cat = self._out(f"{p}/attn/k_concat") + self.make_node("Concat", [k, ctx_kv[0]], [k_cat], name=f"{p}/attn/k_concat", axis=0) + self.make_value(k_cat, self.io_dtype, ["num_rows", kv_dim]) + v_cat = self._out(f"{p}/attn/v_concat") + self.make_node("Concat", [v, ctx_kv[1]], [v_cat], name=f"{p}/attn/v_concat", axis=0) + self.make_value(v_cat, self.io_dtype, ["num_rows", kv_dim]) + + q_all = self.binary( + "Gather", + f"{p}/attn/q_gather", + q, + "q_row_map", + self.io_dtype, + ["num_tokens", self.num_heads * self.head_size], + ) + k_all = self.binary("Gather", f"{p}/attn/k_gather", k_cat, "qkv_row_map", self.io_dtype, ["num_tokens", kv_dim]) + v_all = self.binary("Gather", f"{p}/attn/v_gather", v_cat, "qkv_row_map", self.io_dtype, ["num_tokens", kv_dim]) + + q_norm = self.make_initializer( + w[f"layers.{i}.self_attn.q_norm.weight"], f"dspark.layers.{i}.self_attn.q_norm.weight", to=self.io_dtype + ) + k_norm = self.make_initializer( + w[f"layers.{i}.self_attn.k_norm.weight"], f"dspark.layers.{i}.self_attn.k_norm.weight", to=self.io_dtype + ) + + attn_name = f"{p}/attn/PagedAttention" + attn_out = self._out(attn_name) + self.make_node( + "PagedAttention", + [ + q_all, + k_all, + v_all, + f"past_key_values.{i}.key", + f"past_key_values.{i}.value", + "cumulative_sequence_lengths", + "past_sequence_lengths", + "block_table", + "dspark.cos_cache", + "dspark.sin_cache", + "", # slot_mapping: derived from block_table + past_sequence_lengths + "", # head_sink + q_norm, + k_norm, + "", + "", # k_scale / v_scale + "attention_metadata", + ], + [attn_out, f"present.{i}.key", f"present.{i}.value"], + name=attn_name, + domain="com.microsoft", + num_heads=self.num_heads, + kv_num_heads=self.num_kv_heads, + local_window_size=self.sliding_window, + is_causal=1 if self.is_causal else 0, + do_rotary=1, + rotary_interleaved=0, + qk_norm_epsilon=self.rms_eps, + ) + self.make_value(attn_out, self.io_dtype, ["num_tokens", self.num_heads * self.head_size]) + for suffix in ("key", "value"): + self.make_value( + f"present.{i}.{suffix}", + self.io_dtype, + ["num_blocks", self.paged_block_size, self.num_kv_heads, self.head_size], + ) + + block_out = self.binary( + "Gather", + f"{p}/attn/out_gather", + attn_out, + "block_row_index", + self.io_dtype, + [rows_q, self.num_heads * self.head_size], + ) + return self.matmul( + f"{p}/attn/o_proj/MatMul", + block_out, + w[f"layers.{i}.self_attn.o_proj.weight"], + self.num_heads * self.head_size, + self.hidden_size, + rows_q, + ) + + def _make_mlp(self, i, x, w, rows_q): + p = f"/dspark/layers.{i}/mlp" + gate = self.matmul( + f"{p}/gate_proj/MatMul", + x, + w[f"layers.{i}.mlp.gate_proj.weight"], + self.hidden_size, + self.intermediate_size, + rows_q, + ) + up = self.matmul( + f"{p}/up_proj/MatMul", + x, + w[f"layers.{i}.mlp.up_proj.weight"], + self.hidden_size, + self.intermediate_size, + rows_q, + ) + sig = self.unary("Sigmoid", f"{p}/act/Sigmoid", gate, self.io_dtype, [rows_q, self.intermediate_size]) + silu = self.binary("Mul", f"{p}/act/Mul", gate, sig, self.io_dtype, [rows_q, self.intermediate_size]) + prod = self.binary("Mul", f"{p}/act/MulUp", silu, up, self.io_dtype, [rows_q, self.intermediate_size]) + return self.matmul( + f"{p}/down_proj/MatMul", + prod, + w[f"layers.{i}.mlp.down_proj.weight"], + self.intermediate_size, + self.hidden_size, + rows_q, + ) + + def _make_candidates_and_markov(self, final, w): + """Top-k per position plus the Markov bigram bias, packaged as a candidate lattice.""" + n_spec, top_k, rank = self.num_draft_tokens, self.top_k, self.markov_rank + + logits = self._make_lm_head(final) + logits32 = self.unary( + "Cast", "/dspark/topk/Cast", logits, ir.DataType.FLOAT, ["num_block", self.vocab_size], to=ir.DataType.FLOAT + ) + vals, idx = "/dspark/topk/values", "/dspark/topk/indices" + self.make_node( + "TopK", [logits32, self.const([top_k])], [vals, idx], name="/dspark/topk/TopK", axis=-1, largest=1, sorted=1 + ) + self.make_value(vals, ir.DataType.FLOAT, ["num_block", top_k]) + self.make_value(idx, ir.DataType.INT64, ["num_block", top_k]) + + cand = self.reshape( + "/dspark/topk/cand", idx, [-1, n_spec, top_k], ir.DataType.INT64, ["batch_size", n_spec, top_k] + ) + unary_logits = self.reshape( + "/dspark/topk/unary", vals, [-1, n_spec, top_k], ir.DataType.FLOAT, ["batch_size", n_spec, top_k] + ) + + # The bias at slot l is conditioned on the token at slot l-1: the anchor for l == 0 and + # slot l-1's own candidate afterwards. That is a pairwise edge, so it walks like a lattice. + ids2 = self.reshape( + "/dspark/markov/ids", "input_ids", [-1, self.block_size], ir.DataType.INT64, ["batch_size", self.block_size] + ) + anchor = self._out("/dspark/markov/anchor") + self.make_node( + "Slice", [ids2, self.const([0]), self.const([1]), self.const([1])], [anchor], name="/dspark/markov/anchor" + ) + self.make_value(anchor, ir.DataType.INT64, ["batch_size", 1]) + anchor3 = self.reshape("/dspark/markov/anchor3", anchor, [-1, 1, 1], ir.DataType.INT64, ["batch_size", 1, 1]) + anchor_tiled = self.binary( + "Tile", + "/dspark/markov/anchor_tile", + anchor3, + self.const([1, 1, top_k]), + ir.DataType.INT64, + ["batch_size", 1, top_k], + ) + prev_tail = self._out("/dspark/markov/prev_tail") + self.make_node( + "Slice", + [cand, self.const([0]), self.const([n_spec - 1]), self.const([1])], + [prev_tail], + name="/dspark/markov/prev_tail", + ) + self.make_value(prev_tail, ir.DataType.INT64, ["batch_size", n_spec - 1, top_k]) + prev = self._out("/dspark/markov/prev") + self.make_node("Concat", [anchor_tiled, prev_tail], [prev], name="/dspark/markov/prev", axis=1) + self.make_value(prev, ir.DataType.INT64, ["batch_size", n_spec, top_k]) + + self.make_initializer(w["markov_head.markov_w1.weight"], "dspark.markov_w1", to=self.io_dtype) + self.make_initializer(w["markov_head.markov_w2.weight"], "dspark.markov_w2", to=self.io_dtype) + pred = self.binary( + "Gather", + "/dspark/markov/pred_gather", + "dspark.markov_w1", + prev, + self.io_dtype, + ["batch_size", n_spec, top_k, rank], + ) + succ = self.binary( + "Gather", + "/dspark/markov/succ_gather", + "dspark.markov_w2", + cand, + self.io_dtype, + ["batch_size", n_spec, top_k, rank], + ) + # The bias is added to a float32 logit, and a bf16 rank-256 dot product carries enough + # relative error to reorder the lattice, so the (tiny) gathered slices go to float32 first. + pred32 = self.unary( + "Cast", + "/dspark/markov/pred_cast", + pred, + ir.DataType.FLOAT, + ["batch_size", n_spec, top_k, rank], + to=ir.DataType.FLOAT, + ) + succ32 = self.unary( + "Cast", + "/dspark/markov/succ_cast", + succ, + ir.DataType.FLOAT, + ["batch_size", n_spec, top_k, rank], + to=ir.DataType.FLOAT, + ) + succ_t = self.unary( + "Transpose", + "/dspark/markov/succ_t", + succ32, + ir.DataType.FLOAT, + ["batch_size", n_spec, rank, top_k], + perm=[0, 1, 3, 2], + ) + pair32 = self.binary( + "MatMul", "/dspark/markov/pair", pred32, succ_t, ir.DataType.FLOAT, ["batch_size", n_spec, top_k, top_k] + ) + unary4 = self.reshape( + "/dspark/markov/unary4", + unary_logits, + [-1, n_spec, 1, top_k], + ir.DataType.FLOAT, + ["batch_size", n_spec, 1, top_k], + ) + self.make_node("Add", [pair32, unary4], ["draft_scores"], name="/dspark/markov/scores") + self.make_value("draft_scores", ir.DataType.FLOAT, ["batch_size", n_spec, top_k, top_k]) + self.make_node("Cast", [cand], ["draft_candidate_ids"], name="/dspark/markov/cand_cast", to=ir.DataType.INT32) + self.make_value("draft_candidate_ids", ir.DataType.INT32, ["batch_size", n_spec, top_k]) + + self.graph.outputs.append(self.values["draft_candidate_ids"]) + self.graph.outputs.append(self.values["draft_scores"]) + for i in range(self.num_layers): + self.graph.outputs.append(self.values[f"present.{i}.key"]) + self.graph.outputs.append(self.values[f"present.{i}.value"]) + + def _make_lm_head(self, root): + w = self._weights + weight, scale = w["lm_head.weight"], w["lm_head.weight_scale"] + if weight.dtype != torch.float8_e4m3fn: + name = "/lm_head/MatMul" + return self.matmul( + name, root, weight, self.hidden_size, self.vocab_size, "num_block", weight_name="lm_head.MatMul.weight" + ) + # The quantized head is the target's and only runs in the target's dtype. Its input is the + # drafter's final normed hidden state, which is back inside the fp16 range. + root = self.unary( + "Cast", "/lm_head/Cast", root, self.external_dtype, ["num_block", self.hidden_size], to=self.external_dtype + ) + self.make_initializer(weight.contiguous(), "lm_head.MatMul.fp8_weight") + self.make_initializer( + scale.reshape(self.vocab_size, 1), "lm_head.MatMul.fp8_weight_scale", to=ir.DataType.FLOAT + ) + out = "/lm_head/MatMul/output_0" + self.make_node( + "MatMulBlockQuantizedFp8Weight", + [root, "lm_head.MatMul.fp8_weight", "lm_head.MatMul.fp8_weight_scale"], + [out], + name="/lm_head/MatMul", + domain="com.microsoft", + block_size=int(weight.shape[1]), + ) + self.make_value(out, self.external_dtype, ["num_block", self.vocab_size]) + return out + + # ------------------------------------------------------------------- I/O + + def _declare_io(self): + decls = [ + ("aux_hidden_states", self.external_dtype, ["num_ctx", self.aux_hidden_size]), + ("input_ids", ir.DataType.INT64, ["num_block"]), + ("q_row_map", ir.DataType.INT32, ["num_tokens"]), + ("qkv_row_map", ir.DataType.INT32, ["num_tokens"]), + ("block_row_index", ir.DataType.INT32, ["num_block"]), + ("cumulative_sequence_lengths", ir.DataType.INT32, ["batch_size + 1"]), + ("past_sequence_lengths", ir.DataType.INT32, ["batch_size"]), + ("block_table", ir.DataType.INT32, ["batch_size", "max_num_blocks"]), + ("attention_metadata", ir.DataType.INT32, [3]), + ] + for name, dtype, shape in decls: + self.graph.inputs.append(self.make_value(name, dtype, shape)) + for i in range(self.num_layers): + for suffix in ("key", "value"): + self.graph.inputs.append( + self.make_value( + f"past_key_values.{i}.{suffix}", + self.io_dtype, + ["num_blocks", self.paged_block_size, self.num_kv_heads, self.head_size], + ) + ) + + # --------------------------------------------------------------- weights + + def _load_weights(self): + import safetensors.torch as safetensors_torch # noqa: PLC0415 + + weights = {} + for shard in sorted(glob.glob(os.path.join(self.draft_dir, "*.safetensors"))): + with safetensors_torch.safe_open(shard, framework="pt") as f: + for key in f.keys(): # noqa: SIM118 -- safetensors handles are not iterable + weights[key] = f.get_tensor(key) + for required in ("fc.weight", "markov_head.markov_w1.weight", "markov_head.markov_w2.weight"): + if required not in weights: + raise ValueError(f"'{self.draft_dir}' does not look like a DSpark checkpoint (no {required}).") + + embed_keys = {"model.embed_tokens.weight", "model.language_model.embed_tokens.weight"} + for shard in sorted(glob.glob(os.path.join(self.target_dir, "*.safetensors"))): + if os.path.basename(shard).startswith("model_mtp"): + continue + with safetensors_torch.safe_open(shard, framework="pt") as f: + for key in f.keys(): # noqa: SIM118 -- safetensors handles are not iterable + if key in embed_keys: + weights["embed_tokens.weight"] = f.get_tensor(key) + elif key == "lm_head.weight": + weights["lm_head.weight"] = f.get_tensor(key) + elif key in ("lm_head.weight_scale", "lm_head.weight_global_scale"): + weights["lm_head.weight_scale"] = f.get_tensor(key) + for required in ("embed_tokens.weight", "lm_head.weight"): + if required not in weights: + raise ValueError(f"Could not find '{required}' in the target checkpoint '{self.target_dir}'.") + self._weights = weights + return weights + + # ------------------------------------------------------------------ save + + def save_model(self, out_dir): + out_path = os.path.join(out_dir, self.filename) + data_path = out_path + ".data" + for path in (out_path, data_path): + if os.path.exists(path): + os.remove(path) + with tqdm() as pbar: + total_set = False + + def callback(tensor, metadata): + nonlocal total_set + if not total_set: + pbar.total = metadata.total + total_set = True + pbar.update() + pbar.set_description(f"Saving {tensor.name} ({tensor.dtype.short_name()}, {tensor.shape})") + + ir.save( + self.model, + out_path, + external_data=os.path.basename(data_path), + size_threshold_bytes=0, + callback=callback, + ) + + def genai_config_section(self): + return { + "filename": self.filename, + "num_hidden_layers": self.num_layers, + "num_key_value_heads": self.num_kv_heads, + "head_size": self.head_size, + "block_size": self.block_size, + "num_draft_tokens": self.num_draft_tokens, + "selector_top_k": self.top_k, + "mask_token_id": self.mask_token_id, + "sliding_window": self.sliding_window, + "main_aux_hidden_states": "aux_hidden_states", + "inputs": { + "aux_hidden_states": "aux_hidden_states", + "input_ids": "input_ids", + "q_row_map": "q_row_map", + "qkv_row_map": "qkv_row_map", + "block_row_index": "block_row_index", + "cumulative_sequence_lengths": "cumulative_sequence_lengths", + "past_sequence_lengths": "past_sequence_lengths", + "block_table": "block_table", + "attention_metadata": "attention_metadata", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + }, + "outputs": { + "candidate_ids": "draft_candidate_ids", + "scores": "draft_scores", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + }, + } diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index 194e83ef57..c6e097c741 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -833,7 +833,7 @@ class Qwen35MoEModel(MTPModel): # Extra options naming a block drafter. The Engine drives one drafter per model, so any of # these supersedes the MTP head rather than shipping beside it. - block_drafter_options = ("dflash2_path",) + block_drafter_options = ("dflash2_path", "dspark_path") def requested_block_drafter(self, extra_options): return next((name for name in self.block_drafter_options if extra_options.get(name)), None) @@ -858,6 +858,10 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): self.dflash2_shared_initializers = [] self.make_dflash2_init(io_dtype, extra_options) + self.dspark = None + self.dspark_shared_initializers = [] + self.make_dspark_init(io_dtype, extra_options) + self.vocab_size = self.decoder.vocab_size self.hf_token = self.decoder.hf_token self.hf_remote = self.decoder.hf_remote @@ -909,6 +913,9 @@ def make_mtp_model(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_opti # A block drafter reads the target's aux hidden states, so it is never nested in the head. mtp_options.pop("dflash2_path", None) mtp_options.pop("dflash2_num_draft_tokens", None) + mtp_options.pop("dspark_path", None) + mtp_options.pop("dspark_num_draft_tokens", None) + mtp_options.pop("dspark_top_k", None) self.mtp = self.get_mtp_model_class()( copy.deepcopy(config), self.mtp_attrs["io_dtype"], @@ -937,6 +944,7 @@ def make_model(self, input_path): print("Building MTP (multi-token prediction) head -> mtp.onnx") self.mtp.make_model(input_path) self.make_dflash2_model(input_path) + self.make_dspark_model(input_path) def save_model(self, output_dir): self.decoder.save_model(output_dir) @@ -946,6 +954,7 @@ def save_model(self, output_dir): output_dir, self.decoder.filename, self.mtp.filename ) self.save_dflash2_model(output_dir) + self.save_dspark_model(output_dir) def make_genai_config(self, config, extra_kwargs, out_dir): self.decoder.model_type = self.model_type @@ -954,6 +963,8 @@ def make_genai_config(self, config, extra_kwargs, out_dir): self.add_mtp_to_genai_config(out_dir) if self.dflash2 is not None: self.add_dflash2_to_genai_config(out_dir) + if self.dspark is not None: + self.add_dspark_to_genai_config(out_dir) def add_mtp_to_genai_config(self, out_dir): config_path = os.path.join(out_dir, "genai_config.json") @@ -1076,6 +1087,108 @@ def add_dflash2_to_genai_config(self, out_dir): json.dump(genai_config, config_file, indent=4) print("Added 'dflash2' section to genai_config.json") + def make_dspark_init(self, io_dtype, extra_options): + """DSpark block drafter, exported as an auxiliary ``dspark.onnx``. + + ``dspark_path`` points at the draft checkpoint. SpecForge taps the *output* of each + ``target_layer_ids`` entry (``hidden_states[layer_id + 1]``), which is the residual stream + entering layer ``layer_id + 1`` -- the tensor aux_hidden_state_layers names. Getting that + off by one leaves acceptance at exactly 1.0. + """ + self.dspark_path = extra_options.get("dspark_path") + if not self.dspark_path: + return + if self.dflash2_path: + raise ValueError("dspark_path and dflash2_path are mutually exclusive.") + if not self.decoder.use_paged_attention: + raise ValueError("dspark_path requires use_paged_attention=true.") + + num_draft_tokens = None + if "dspark_num_draft_tokens" in extra_options: + try: + num_draft_tokens = int(extra_options["dspark_num_draft_tokens"]) + except (TypeError, ValueError) as error: + raise ValueError("dspark_num_draft_tokens must be a positive integer.") from error + if num_draft_tokens < 1: + raise ValueError("dspark_num_draft_tokens must be a positive integer.") + + try: + top_k = int(extra_options.get("dspark_top_k", 16)) + except (TypeError, ValueError) as error: + raise ValueError("dspark_top_k must be a positive integer.") from error + if top_k < 1: + raise ValueError("dspark_top_k must be a positive integer.") + + with open(os.path.join(self.dspark_path, "config.json"), encoding="utf-8") as handle: + draft_config = json.load(handle) + vocab_size = int(draft_config["vocab_size"]) + if top_k > vocab_size: + raise ValueError(f"dspark_top_k must not exceed the drafter vocabulary size ({vocab_size}).") + + self.dspark_attrs = { + "io_dtype": io_dtype, + "num_draft_tokens": num_draft_tokens, + "top_k": top_k, + } + + target_layer_ids = draft_config["dflash_config"]["target_layer_ids"] + expected = ",".join(str(i + 1) for i in target_layer_ids) + actual = ",".join(str(i) for i in self.decoder.aux_hidden_state_layers) + if actual != expected: + raise ValueError( + f"The DSpark drafter needs aux_hidden_state_layers={expected} on the main model, got '{actual}'." + ) + + def make_dspark_model(self, input_path): + if not self.dspark_path: + return + from .dspark import DSparkBuilder # noqa: PLC0415 + + print("Building DSpark draft model -> dspark.onnx") + target_dir = input_path if input_path and os.path.isdir(input_path) else self.decoder.model_name_or_path + self.dspark = DSparkBuilder( + self.dspark_path, + target_dir, + self.dspark_attrs["io_dtype"], + self.decoder.attention_attrs["paged_block_size"], + self.decoder.original_context_length or self.decoder.context_length, + num_draft_tokens=self.dspark_attrs["num_draft_tokens"], + top_k=self.dspark_attrs["top_k"], + ) + self.dspark.make_model() + + def save_dspark_model(self, output_dir): + if self.dspark is None: + return + self.dspark.save_model(output_dir) + self.dspark_shared_initializers = self.share_initializers( + output_dir, self.decoder.filename, self.dspark.filename + ) + + def add_dspark_to_genai_config(self, out_dir): + config_path = os.path.join(out_dir, "genai_config.json") + with open(config_path) as config_file: + genai_config = json.load(config_file) + + decoder = genai_config["model"]["decoder"] + decoder.setdefault("outputs", {}).setdefault("aux_hidden_states", "aux_hidden_states") + + section = self.dspark.genai_config_section() + section["aux_hidden_state_layers"] = list(self.decoder.aux_hidden_state_layers) + if self.dspark_shared_initializers: + existing = decoder.get("shared_initializers", []) + known = {json.dumps(entry, sort_keys=True) for entry in existing} + for entry in self.dspark_shared_initializers: + if json.dumps(entry, sort_keys=True) not in known: + existing.append(entry) + decoder["shared_initializers"] = existing + section["shared_initializers"] = self.dspark_shared_initializers + genai_config["model"]["dspark"] = section + + with open(config_path, "w") as config_file: + json.dump(genai_config, config_file, indent=4) + print("Added 'dspark' section to genai_config.json") + class Qwen35MTPModel(Qwen35MoETextModel): """Qwen3.6 multi-token-prediction self-speculative head builder.""" diff --git a/test/python/builder/test_qwen_dspark_export.py b/test/python/builder/test_qwen_dspark_export.py new file mode 100644 index 0000000000..e2f3cb7d6e --- /dev/null +++ b/test/python/builder/test_qwen_dspark_export.py @@ -0,0 +1,252 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License + +import importlib +import json +import types + +import onnx_ir as ir +import pytest + +from models.builders.dspark import DSparkBuilder +from models.builders.mtp import MTPModel +from models.builders.qwen import Qwen35MoEModel + +# SpecForge names the layers whose *output* it taps, and the builder names the residual stream +# *entering* a layer, so the tap set is each target layer id plus one. +TARGET_LAYER_IDS = [0, 10, 20] +AUX_LAYERS = [1, 11, 21] + + +def _draft_checkpoint(tmp_path, name="dspark_draft", target_layer_ids=TARGET_LAYER_IDS): + draft_dir = tmp_path / name + draft_dir.mkdir() + config = { + "hidden_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 16, + "vocab_size": 32, + "rms_norm_eps": 1e-6, + "max_position_embeddings": 128, + "rope_parameters": {"rope_theta": 10000.0}, + "block_size": 5, + "markov_rank": 4, + "dflash_config": { + "mask_token_id": 31, + "target_layer_ids": target_layer_ids, + }, + } + (draft_dir / "config.json").write_text(json.dumps(config)) + return str(draft_dir) + + +def _composite(aux_layers=AUX_LAYERS, use_paged_attention=True, dflash2_path=None): + model = object.__new__(Qwen35MoEModel) + model.dspark = None + model.dspark_shared_initializers = [] + model.dflash2_path = dflash2_path + model.decoder = types.SimpleNamespace( + use_paged_attention=use_paged_attention, + aux_hidden_state_layers=list(aux_layers), + filename="model.onnx", + attention_attrs={"paged_block_size": 256}, + context_length=32768, + original_context_length=131072, + ) + return model + + +def test_absent_option_builds_no_drafter(tmp_path): + model = _composite() + + model.make_dspark_init(io_dtype=None, extra_options={}) + model.make_dspark_model(str(tmp_path)) + + assert model.dspark is None + + +def test_drafter_requires_paged_attention(tmp_path): + model = _composite(use_paged_attention=False) + + with pytest.raises(ValueError, match="use_paged_attention"): + model.make_dspark_init(io_dtype=None, extra_options={"dspark_path": _draft_checkpoint(tmp_path)}) + + +def test_two_block_drafters_cannot_be_exported_together(tmp_path): + model = _composite(dflash2_path=_draft_checkpoint(tmp_path, name="dflash2_draft")) + + with pytest.raises(ValueError, match="mutually exclusive"): + model.make_dspark_init(io_dtype=None, extra_options={"dspark_path": _draft_checkpoint(tmp_path)}) + + +# Passing SpecForge's own target_layer_ids straight through is the off-by-one that leaves +# acceptance pinned at exactly 1.0, so it has to be rejected rather than tolerated. +@pytest.mark.parametrize("aux_layers", [TARGET_LAYER_IDS, [1, 11], [2, 11, 21], [21, 11, 1], []]) +def test_mismatched_tap_layers_are_rejected(tmp_path, aux_layers): + model = _composite(aux_layers=aux_layers) + + with pytest.raises(ValueError, match="aux_hidden_state_layers"): + model.make_dspark_init(io_dtype=None, extra_options={"dspark_path": _draft_checkpoint(tmp_path)}) + + +def test_tap_layers_one_past_each_target_layer_are_accepted(tmp_path): + model = _composite() + + model.make_dspark_init(io_dtype=None, extra_options={"dspark_path": _draft_checkpoint(tmp_path)}) + + assert model.dspark_attrs["num_draft_tokens"] is None + assert model.dspark_attrs["top_k"] == 16 + + +def test_lattice_width_and_draft_count_can_be_overridden(tmp_path): + model = _composite() + + model.make_dspark_init( + io_dtype=None, + extra_options={ + "dspark_path": _draft_checkpoint(tmp_path), + "dspark_num_draft_tokens": "6", + "dspark_top_k": "8", + }, + ) + + assert model.dspark_attrs["num_draft_tokens"] == 6 + assert model.dspark_attrs["top_k"] == 8 + + +@pytest.mark.parametrize("num_draft_tokens", ["0", "-1", "invalid"]) +def test_draft_token_count_must_be_positive(tmp_path, num_draft_tokens): + model = _composite() + + with pytest.raises(ValueError, match="dspark_num_draft_tokens must be a positive integer"): + model.make_dspark_init( + io_dtype=None, + extra_options={ + "dspark_path": _draft_checkpoint(tmp_path), + "dspark_num_draft_tokens": num_draft_tokens, + }, + ) + + +@pytest.mark.parametrize("top_k", ["0", "-1", "invalid"]) +def test_lattice_width_must_be_positive(tmp_path, top_k): + model = _composite() + + with pytest.raises(ValueError, match="dspark_top_k must be a positive integer"): + model.make_dspark_init( + io_dtype=None, + extra_options={"dspark_path": _draft_checkpoint(tmp_path), "dspark_top_k": top_k}, + ) + + +def test_lattice_width_cannot_exceed_vocabulary(tmp_path): + model = _composite() + + with pytest.raises(ValueError, match="dspark_top_k must not exceed the drafter vocabulary size"): + model.make_dspark_init( + io_dtype=None, + extra_options={"dspark_path": _draft_checkpoint(tmp_path), "dspark_top_k": "33"}, + ) + + +def test_duplicate_node_names_are_rejected(): + builder = object.__new__(DSparkBuilder) + builder.node_names = {"duplicate"} + + with pytest.raises(ValueError, match="duplicate node name duplicate"): + builder.make_node("Identity", [], [], name="duplicate") + + +def test_kv_cache_uses_configured_paged_block_size(tmp_path): + builder = DSparkBuilder( + _draft_checkpoint(tmp_path), + str(tmp_path), + ir.DataType.FLOAT16, + paged_block_size=512, + max_position_embeddings=128, + ) + + builder._declare_io() + + assert builder.values["past_key_values.0.key"].shape[1] == 512 + + +def test_drafter_uses_original_context_length(tmp_path, monkeypatch): + captured = {} + + class StubDSparkBuilder: + def __init__(self, _draft_dir, _target_dir, _io_dtype, _paged_block_size, max_position, **_kwargs): + captured["max_position"] = max_position + + def make_model(self): + pass + + dspark_module = importlib.import_module("models.builders.dspark") + monkeypatch.setattr(dspark_module, "DSparkBuilder", StubDSparkBuilder) + model = _composite() + model.dspark_path = _draft_checkpoint(tmp_path) + model.dspark_attrs = {"io_dtype": None, "num_draft_tokens": None, "top_k": 16} + + model.make_dspark_model(str(tmp_path)) + + assert captured["max_position"] == model.decoder.original_context_length + + +def test_genai_config_gains_the_drafter_and_the_target_tap(tmp_path): + config_path = tmp_path / "genai_config.json" + config_path.write_text(json.dumps({"model": {"decoder": {}}})) + model = _composite() + model.dspark = types.SimpleNamespace(genai_config_section=lambda: {"filename": "dspark.onnx"}) + + model.add_dspark_to_genai_config(str(tmp_path)) + + config = json.loads(config_path.read_text()) + assert config["model"]["decoder"]["outputs"]["aux_hidden_states"] == "aux_hidden_states" + assert config["model"]["dspark"]["filename"] == "dspark.onnx" + assert config["model"]["dspark"]["aux_hidden_state_layers"] == AUX_LAYERS + + +def test_shared_initializers_are_recorded_once_on_both_sides(tmp_path): + config_path = tmp_path / "genai_config.json" + shared = {"name": "model.embed_tokens.weight", "filename": "model.onnx.data"} + config_path.write_text(json.dumps({"model": {"decoder": {"shared_initializers": [shared]}}})) + model = _composite() + model.dspark = types.SimpleNamespace(genai_config_section=lambda: {"filename": "dspark.onnx"}) + model.dspark_shared_initializers = [shared] + + model.add_dspark_to_genai_config(str(tmp_path)) + + config = json.loads(config_path.read_text()) + assert config["model"]["decoder"]["shared_initializers"] == [shared] + assert config["model"]["dspark"]["shared_initializers"] == [shared] + + +def test_builder_exposes_the_api_the_composite_drives(): + assert all(hasattr(DSparkBuilder, name) for name in ("make_model", "save_model", "genai_config_section")) + + +@pytest.fixture +def mtp_init(monkeypatch): + """Drive Qwen35MoEModel.make_mtp_init with the base seeding stubbed out.""" + + def run(extra_options, num_mtp_layers=1): + model = object.__new__(Qwen35MoEModel) + model.mtp_attrs = {} + monkeypatch.setattr(MTPModel, "make_mtp_init", lambda self, _c, opts: dict(opts)) + config = types.SimpleNamespace(mtp_num_hidden_layers=num_mtp_layers) + Qwen35MoEModel.make_mtp_init(model, config, extra_options) + return model.mtp_attrs["build"] + + return run + + +# The Engine drives one drafter per model, so DSpark replaces the MTP head just as DFlash 2 does. +def test_a_dspark_drafter_suppresses_the_mtp_head(tmp_path, mtp_init): + assert mtp_init({"dspark_path": _draft_checkpoint(tmp_path)}) is False + + +def test_an_mtp_head_is_still_built_without_a_block_drafter(mtp_init): + assert mtp_init({}) is True