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