from __future__ import annotations
import json
from typing import Any, Callable, Iterable, TYPE_CHECKING
import numpy as np
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger
@ModelBase.register("QWenLMHeadModel")
@ModelBase.example("Qwen/Qwen-7B")
class QwenModel(TextModel):
model_arch = gguf.MODEL_ARCH.QWEN
@staticmethod
def token_bytes_to_string(b):
from transformers.convert_slow_tokenizer import bytes_to_unicode
byte_encoder = bytes_to_unicode()
return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')])
@staticmethod
def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]:
parts = [bytes([b]) for b in token]
while True:
min_idx = None
min_rank = None
for i, pair in enumerate(zip(parts[:-1], parts[1:])):
rank = mergeable_ranks.get(pair[0] + pair[1])
if rank is not None and (min_rank is None or rank < min_rank):
min_idx = i
min_rank = rank
if min_rank is None or (max_rank is not None and min_rank >= max_rank):
break
assert min_idx is not None
parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:]
return parts
def set_vocab(self):
self._set_vocab_qwen()
@ModelBase.register(
"Qwen2Model",
"Qwen2ForCausalLM",
"Qwen2AudioForConditionalGeneration",
"KORMoForCausalLM",
"AudioFlamingo3ForConditionalGeneration",
"DotsOCRForCausalLM",
)
@ModelBase.example("Qwen/Qwen2.5-7B-Instruct")
class Qwen2Model(TextModel):
model_arch = gguf.MODEL_ARCH.QWEN2
def set_vocab(self):
try:
self._set_vocab_sentencepiece()
except FileNotFoundError:
self._set_vocab_gpt2()
def set_gguf_parameters(self):
super().set_gguf_parameters()
self._try_set_pooling_type()
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if self.hf_arch == "Qwen2Model":
name = f"model.{name}" # map to Qwen2ForCausalLM tensors
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Qwen2MoeForCausalLM")
@ModelBase.example("Qwen/Qwen1.5-MoE-A2.7B")
class Qwen2MoeModel(TextModel):
model_arch = gguf.MODEL_ARCH.QWEN2MOE
def set_gguf_parameters(self):
super().set_gguf_parameters()
if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:
self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)
logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}")
if (shared_expert_intermediate_size := self.hparams.get('shared_expert_intermediate_size')) is not None:
self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size)
logger.info(f"gguf: expert shared feed forward length = {shared_expert_intermediate_size}")
_experts: list[dict[str, Tensor]] | None = None
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
# handle aggregated expert tensors
# GGUF stores dimensions reversed from PyTorch, so:
# PyTorch (A,B,C) -> GGUF writes [C,B,A] -> GGML reads ne={C,B,A}
# Input shapes from HF: (n_expert, n_ff_exp, n_embd) or (n_expert, n_embd, n_ff_exp)
# Expected GGML ne: {n_embd, n_ff_exp, n_expert} for gate/up, {n_ff_exp, n_embd, n_expert} for down
if name.endswith("mlp.experts.down_proj") or name.endswith("mlp.experts.down_proj.weight"):
mapped = f"{name}.weight" if not name.endswith(".weight") else name
# HF: [n_expert, n_embd, n_ff] -> GGML: {n_ff, n_embd, n_expert}
yield from super().modify_tensors(data_torch, mapped, bid)
return
if name.endswith("mlp.experts.gate_up_proj") or name.endswith("mlp.experts.gate_up_proj.weight"):
if data_torch.ndim < 3 or data_torch.shape[-2] % 2 != 0:
raise ValueError(f"Unexpected gate_up_proj shape for {name}: {tuple(data_torch.shape)}")
# HF: [n_expert, 2*n_ff, n_embd] -> split on dim=-2
n_ff = data_torch.shape[-2] // 2
gate = data_torch[..., :n_ff, :].contiguous()
up = data_torch[..., n_ff:, :].contiguous()
# gate/up: [n_expert, n_ff, n_embd] -> GGML: {n_embd, n_ff, n_expert}
base_name = name.removesuffix(".weight").removesuffix(".gate_up_proj")
mapped_gate = f"{base_name}.gate_proj.weight"
mapped_up = f"{base_name}.up_proj.weight"
yield from super().modify_tensors(gate, mapped_gate, bid)
yield from super().modify_tensors(up, mapped_up, bid)
return
if name.find("experts") != -1:
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
assert bid is not None
if self._experts is None:
self._experts = [{} for _ in range(self.block_count)]
self._experts[bid][name] = data_torch
if len(self._experts[bid]) >= n_experts * 3:
# merge the experts into a single 3d tensor
for w_name in ["down_proj", "gate_proj", "up_proj"]:
datas: list[Tensor] = []
for xid in range(n_experts):
ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
datas.append(self._experts[bid][ename])
del self._experts[bid][ename]
data_torch = torch.stack(datas, dim=0)
merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
yield from super().modify_tensors(data_torch, merged_name, bid)
return
else:
return
yield from super().modify_tensors(data_torch, name, bid)
def prepare_tensors(self):
super().prepare_tensors()
if self._experts is not None:
# flatten `list[dict[str, Tensor]]` into `list[str]`
experts = [k for d in self._experts for k in d.keys()]
if len(experts) > 0:
raise ValueError(f"Unprocessed experts: {experts}")
@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model")
@ModelBase.example("Qwen/Qwen3-8B")
class Qwen3Model(Qwen2Model):
model_arch = gguf.MODEL_ARCH.QWEN3
# extra logic for rerank models
is_rerank: bool = False
is_tied_embeddings: bool = False
token_false_id: int | None = None
token_true_id: int | None = None
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# track for intern-s1-mini
hparams = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)
self.origin_hf_arch = hparams.get('architectures', [None])[0]
if self._is_qwen3_reranker():
self._find_rerank_config()
def _is_qwen3_reranker(self) -> bool:
readme_path = self.dir_model / "README.md"
readme_text = ""
if readme_path.exists():
with readme_path.open("r", encoding="utf-8") as f:
readme_text = f.read()
name_hints = [
str(self.dir_model.name),
str(self.hparams.get("_name_or_path", "")),
str(self.hparams.get("model_type", "")),
str(self.origin_hf_arch or ""),
]
name_hints = [hint.lower() for hint in name_hints if hint]
if "# qwen3-reranker" in readme_text.lower() or "# qwen3-vl-reranker" in readme_text.lower():
return True
if any("qwen3-reranker" in hint or "qwen3-vl-reranker" in hint for hint in name_hints):
return True
return "sequenceclassification" in (self.origin_hf_arch or "").lower()
def set_vocab(self):
# deal with intern-s1-mini
if self.origin_hf_arch == 'InternS1ForConditionalGeneration':
self._set_vocab_interns1()
return
super().set_vocab()
def _find_rerank_config(self):
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
self.is_rerank = True
self.is_tied_embeddings = self.hparams.get("tie_word_embeddings", False)
self.token_false_id = tokenizer.convert_tokens_to_ids("no") # ty: ignore[unresolved-attribute, invalid-assignment]
self.token_true_id = tokenizer.convert_tokens_to_ids("yes") # ty: ignore[unresolved-attribute, invalid-assignment]
self.sep_token_id = tokenizer.convert_tokens_to_ids("|") # ty: ignore[unresolved-attribute]
assert self.token_false_id is not None and self.token_true_id is not None
def set_gguf_parameters(self):
super().set_gguf_parameters()
if self.is_rerank:
self.gguf_writer.add_pooling_type(gguf.PoolingType.RANK)
self.gguf_writer.add_classifier_output_labels(["yes", "no"])
self.gguf_writer.add_chat_template([{
"name": "rerank",
"template": "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n"
"<|im_start|>user\n<Instruct>: Given a web search query, retrieve relevant passages that answer the query\n<Query>: {query}\n<Document>: {document}<|im_end|>\n"
"<|im_start|>assistant\n<think>\n\n</think>\n\n"
}])
def _get_cls_out_tensor(self, data_torch: Tensor) -> Tensor:
# extract "yes" and "no" tokens from the output lm_head tensor
false_row = data_torch[self.token_false_id]
true_row = data_torch[self.token_true_id]
return torch.stack([true_row, false_row], dim=0)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if self.is_rerank:
is_tied_head = self.is_tied_embeddings and "embed_tokens" in name
is_real_head = not self.is_tied_embeddings and "lm_head" in name
if is_tied_head or is_real_head:
cls_out_head = (
gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight",
self._get_cls_out_tensor(data_torch),
)
yield cls_out_head
if is_tied_head:
yield from super().modify_tensors(data_torch, name, bid)
return
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Qwen3MoeForCausalLM")
@ModelBase.example("Qwen/Qwen3-30B-A3B")
class Qwen3MoeModel(Qwen2MoeModel):
model_arch = gguf.MODEL_ARCH.QWEN3MOE
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
hparams = ModelBase.load_hparams(self.dir_model, False)
self.origin_hf_arch = hparams.get('architectures', [None])[0]
def set_vocab(self):
# deal with intern-s1
if self.origin_hf_arch == 'InternS1ForConditionalGeneration':
self._set_vocab_interns1()
return
super().set_vocab()
class _QwenMtpMixin:
"""Shared MTP wiring for Qwen3-Next and Qwen3.5/3.6 text variants. The HF
config carries the MTP block under `mtp_num_hidden_layers` (computed from
the checkpoint when absent, e.g. Qwen3-Next) and the tensors under
`mtp.*`; we extend block_count, emit the nextn metadata key, and remap
`mtp.*` to the standard layer-indexed nextn naming so the existing
tensor_map handles them."""
supports_mtp_export = True
hparams: dict[str, Any]
model_arch: gguf.MODEL_ARCH
gguf_writer: gguf.GGUFWriter
block_count: int
tensor_map: gguf.TensorNameMap
no_mtp: bool
mtp_only: bool
_original_block_count: int | None = None
opt_num_mtp_layers: int = 0
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.block_count = self.hparams["num_hidden_layers"]
if not self.no_mtp:
n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)
# Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config.
if n_mtp == 0:
assert self.opt_num_mtp_layers != 0
n_mtp = self.opt_num_mtp_layers
self.block_count += n_mtp
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
hparams = {**self.hparams, **self.hparams.get("text_config", {})}
key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
type(self)._original_block_count = hparams.get(key)
type(self).opt_num_mtp_layers = 0
return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute]
@classmethod
def filter_tensors(cls, item):
assert cls._original_block_count is not None
# TODO: change TextModel to super()
if (titem := TextModel.filter_tensors(item)) is None:
return None
name, gen = titem
if name.startswith("model.mtp."):
name = name.replace("model.", "", 1)
if name.startswith("mtp."):
if cls.no_mtp:
return None
remapper = {
"fc": "eh_proj",
"pre_fc_norm_embedding": "enorm",
"pre_fc_norm_hidden": "hnorm",
"norm": "shared_head.norm",
}
parts = name.split(".", 3)
if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():
mtp_idx = int(parts[2])
name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"
cls.opt_num_mtp_layers = max(cls.opt_num_mtp_layers, mtp_idx + 1)
elif len(parts) == 3 and parts[1] in remapper:
name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}"
elif cls.mtp_only:
keep = name in (
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
"embed_tokens.weight", "norm.weight",
)
if not keep:
return None
return name, gen
def set_gguf_parameters(self):
super().set_gguf_parameters() # ty: ignore[unresolved-attribute]
if self.no_mtp:
return
if (n := self.block_count - self.hparams["num_hidden_layers"]) > 0:
self.gguf_writer.add_nextn_predict_layers(n)
def prepare_metadata(self, vocab_only: bool):
from_dir = self.fname_out.is_dir()
super().prepare_metadata(vocab_only=vocab_only) # ty: ignore[unresolved-attribute]
if not self.mtp_only or not from_dir:
return
output_type: str = self.ftype.name.partition("_")[2] # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
fname_default: str = gguf.naming_convention(
self.metadata.name, self.metadata.basename, self.metadata.finetune, # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
self.metadata.version, size_label=None, output_type=output_type, model_type=None) # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
@ModelBase.register("Qwen3NextForCausalLM")
@ModelBase.example("Qwen/Qwen3-Next-80B-A3B-Instruct")
class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel):
model_arch = gguf.MODEL_ARCH.QWEN3NEXT
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_ssm_conv_kernel(self.hparams["linear_conv_kernel_dim"])
self.gguf_writer.add_ssm_state_size(self.hparams["linear_key_head_dim"])
self.gguf_writer.add_ssm_group_count(self.hparams["linear_num_key_heads"])
self.gguf_writer.add_ssm_time_step_rank(self.hparams["linear_num_value_heads"])
self.gguf_writer.add_ssm_inner_size(self.hparams["linear_value_head_dim"] * self.hparams["linear_num_value_heads"])
self.gguf_writer.add_full_attention_interval(self.hparams.get("full_attention_interval", 4))
if (rope_dim := self.hparams.get("head_dim")) is None:
rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.25)))
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name.endswith(".A_log"):
data_torch = -torch.exp(data_torch)
elif name.endswith(".dt_bias"):
name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
elif "conv1d" in name:
data_torch = data_torch.squeeze()
elif name.endswith("norm.weight") and not name.endswith("linear_attn.norm.weight"):
data_torch = data_torch + 1
if "in_proj_qkvz.weight" in name:
# original order: [q, k, v, z] * head_count
# corrected order: [q * head_count, k * head_count, v * head_count, z * head_count]
head_k_dim = self.hparams["linear_key_head_dim"]
head_v_dim = self.hparams["linear_value_head_dim"]
num_v_heads = self.hparams["linear_num_value_heads"]
num_k_heads = self.hparams["linear_num_key_heads"]
hidden_size = self.hparams["hidden_size"]
split_arg_list_qkvz = [
head_k_dim, # q partition
head_k_dim, # k partition
(num_v_heads // num_k_heads * head_v_dim), # v partition
(num_v_heads // num_k_heads * head_v_dim), # z partition
]
# view as (n_embd, head_count, [q+k+v+z])
data_torch = data_torch.permute(1, 0).contiguous()
data_torch = data_torch.view(-1, num_k_heads, sum(split_arg_list_qkvz))
# split into q, k, v, z
q, k, v, z = torch.split(data_torch, split_arg_list_qkvz, dim=-1)
# flatten dim + head_count
q = q.contiguous().view(hidden_size, -1)
k = k.contiguous().view(hidden_size, -1)
v = v.contiguous().view(hidden_size, -1)
z = z.contiguous().view(hidden_size, -1)
# stack back
qkv = torch.cat([q, k, v], dim=-1).permute(1, 0).contiguous()
z = z.permute(1, 0).contiguous()
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, bid, ".weight"), qkv)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_GATE, bid, ".weight"), z)
else:
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("RND1")
@ModelBase.example("radicalnumerics/RND1-Base-0910")
class RND1Model(Qwen2MoeModel):
model_arch = gguf.MODEL_ARCH.RND1
def set_gguf_parameters(self):
super().set_gguf_parameters()
# RND1 specific parameters
# RND1 uses bidirectional attention
self.gguf_writer.add_causal_attention(False)
if (mask_token_id := self.hparams.get("mask_token_id")) is not None:
self.gguf_writer.add_mask_token_id(mask_token_id)
class _LinearAttentionVReorderBase(Qwen3NextModel):
model_arch = gguf.MODEL_ARCH.QWEN3NEXT # overridden by subclasses
"""reorders V heads from grouped to tiled order for ggml broadcast
see https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306
Linear attention may has num_k_heads < num_v_heads. The HF weights store
V heads grouped by K head: [G0_v0..v{r-1}, G1_v0..v{r-1}, ...].
ggml binary ops use tiled broadcast: [K0, K1, ..., K0, K1, ...].
We reorder V heads to tiled order so ggml_repeat can replace the expensive
interleaved repeat: [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...].
"""
@staticmethod
def _reorder_v_heads(tensor: Tensor, dim: int, num_k_heads: int, num_v_per_k: int, head_dim: int) -> Tensor:
"""Reorder V heads from grouped (by K head) to tiled order along the given dimension."""
shape = list(tensor.shape)
if dim < 0:
dim += len(shape)
new_shape = shape[:dim] + [num_k_heads, num_v_per_k, head_dim] + shape[dim + 1:]
tensor = tensor.reshape(*new_shape)
perm = list(range(len(new_shape)))
perm[dim], perm[dim + 1] = perm[dim + 1], perm[dim]
return tensor.permute(*perm).contiguous().reshape(*shape)
def _transform_nvfp4_weight(self, name: str, weight: Tensor, scale: Tensor) -> tuple[Tensor, Tensor]:
if not name.endswith((
".linear_attn.in_proj_qkv.weight",
".linear_attn.in_proj_z.weight",
".linear_attn.in_proj_a.weight",
".linear_attn.in_proj_b.weight",
".linear_attn.out_proj.weight",
)):
return weight, scale
num_k_heads = self.hparams["linear_num_key_heads"]
num_v_heads = self.hparams["linear_num_value_heads"]
head_k_dim = self.hparams["linear_key_head_dim"]
head_v_dim = self.hparams["linear_value_head_dim"]
num_v_per_k = num_v_heads // num_k_heads
def unpack_nibbles(qs: Tensor) -> Tensor:
lo = torch.bitwise_and(qs, 0x0F)
hi = torch.bitwise_right_shift(qs, 4)
return torch.stack((lo, hi), dim=-1).reshape(*qs.shape[:-1], qs.shape[-1] * 2)
def pack_nibbles(codes: Tensor) -> Tensor:
codes = codes.reshape(*codes.shape[:-1], codes.shape[-1] // 2, 2)
lo = torch.bitwise_and(codes[..., 0], 0x0F)
hi = torch.bitwise_left_shift(torch.bitwise_and(codes[..., 1], 0x0F), 4)
return torch.bitwise_or(lo, hi).contiguous()
def apply_col_perm(qs: Tensor, scales: Tensor, col_perm: Tensor) -> tuple[Tensor, Tensor]:
assert qs.ndim >= 2
assert scales.ndim >= 2
k = qs.shape[-1] * 2
assert col_perm.numel() == k
assert k % 16 == 0
group_cols = col_perm.reshape(-1, 16)
group_starts = group_cols[:, 0]
expected = group_starts.unsqueeze(1) + torch.arange(16, dtype=col_perm.dtype)
assert torch.equal(group_cols, expected)
assert torch.all(group_starts % 16 == 0)
group_perm = (group_starts // 16).to(dtype=torch.long)
expected_groups = torch.arange(scales.shape[-1], dtype=torch.long)
assert group_perm.numel() == scales.shape[-1]
assert torch.equal(torch.sort(group_perm).values, expected_groups)
codes = unpack_nibbles(qs)
codes = codes.index_select(-1, col_perm.to(device=qs.device, dtype=torch.long))
qs = pack_nibbles(codes)
scales = scales.index_select(-1, group_perm.to(device=scales.device))
return qs, scales
def reorder_rows(qs: Tensor, scales: Tensor, head_dim: int) -> tuple[Tensor, Tensor]:
row_perm = self._reorder_v_heads(
torch.arange(num_v_heads * head_dim, dtype=torch.long).unsqueeze(-1),
0, num_k_heads, num_v_per_k, head_dim,
).squeeze(-1)
return (
qs.index_select(0, row_perm.to(device=qs.device)),
scales.index_select(0, row_perm.to(device=scales.device)),
)
if name.endswith(".linear_attn.in_proj_qkv.weight"):
q_dim = head_k_dim * num_k_heads
k_dim = head_k_dim * num_k_heads
q = weight[:q_dim]
k = weight[q_dim:q_dim + k_dim]
v = weight[q_dim + k_dim:]
q_scale = scale[:q_dim]
k_scale = scale[q_dim:q_dim + k_dim]
v_scale = scale[q_dim + k_dim:]
v, v_scale = reorder_rows(v, v_scale, head_v_dim)
return torch.cat([q, k, v], dim=0), torch.cat([q_scale, k_scale, v_scale], dim=0)
if name.endswith(".linear_attn.in_proj_z.weight"):
weight, scale = reorder_rows(weight, scale, head_v_dim)
elif name.endswith((".linear_attn.in_proj_a.weight", ".linear_attn.in_proj_b.weight")):
weight, scale = reorder_rows(weight, scale, 1)
elif name.endswith(".linear_attn.out_proj.weight"):
if self._hadamard_folds_tensor(name):
# folded latent: a column permutation on the rotation axis cannot be
# refolded, so keep the training (grouped) order and let the runtime
# permute the activation instead
self._hadamard_gdn_v_grouped = True
else:
col_perm = self._reorder_v_heads(
torch.arange(num_v_heads * head_v_dim, dtype=torch.long).unsqueeze(0),
1, num_k_heads, num_v_per_k, head_v_dim,
).squeeze(0)
weight, scale = apply_col_perm(weight, scale, col_perm)
return weight, scale
def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):
weight, scale = self._transform_nvfp4_weight(name, weight, scale)
super()._repack_nvfp4(name, weight, scale, scale2, input_scale)
def _hadamard_folds_tensor(self, name: str) -> bool:
# a manifest entry and `name` may differ only by leading wrapper prefixes
return any(name.endswith(n) or n.endswith(name) for n in self.hadamard_folded_names())
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
num_k_heads = self.hparams.get("linear_num_key_heads", 0)
num_v_heads = self.hparams.get("linear_num_value_heads", 0)
if num_k_heads > 0 and num_v_heads > 0 and num_k_heads != num_v_heads and "linear_attn." in name:
head_k_dim = self.hparams["linear_key_head_dim"]
head_v_dim = self.hparams["linear_value_head_dim"]
num_v_per_k = num_v_heads // num_k_heads
if ".in_proj_qkv." in name:
# QKV weight: reorder only the V rows
q_dim = head_k_dim * num_k_heads
k_dim = head_k_dim * num_k_heads
q = data_torch[:q_dim]
k = data_torch[q_dim:q_dim + k_dim]
v = data_torch[q_dim + k_dim:]
v = self._reorder_v_heads(v, 0, num_k_heads, num_v_per_k, head_v_dim)
data_torch = torch.cat([q, k, v], dim=0)
elif ".in_proj_z." in name:
# Z gate weight: reorder rows (num_v_heads * head_v_dim)
data_torch = self._reorder_v_heads(data_torch, 0, num_k_heads, num_v_per_k, head_v_dim)
elif ".in_proj_b." in name or ".in_proj_a." in name:
# Beta/Alpha weight: reorder rows (num_v_heads, head_dim=1)
data_torch = self._reorder_v_heads(data_torch, 0, num_k_heads, num_v_per_k, 1)
elif ".A_log" in name or ".dt_bias" in name or ".dt_proj" in name:
# A_log / dt_bias: 1D parameters with num_v_heads elements
if data_torch.ndim == 1:
data_torch = self._reorder_v_heads(
data_torch.unsqueeze(-1), 0, num_k_heads, num_v_per_k, 1
).squeeze(-1)
else:
data_torch = self._reorder_v_heads(data_torch, -1, num_k_heads, num_v_per_k, 1)
elif ".conv1d" in name:
# Conv1d kernel: reorder only the V channel portion
data = data_torch.squeeze()
qk_channels = head_k_dim * num_k_heads * 2
qk_part = data[:qk_channels]
v_part = data[qk_channels:]
v_part = self._reorder_v_heads(v_part, 0, num_k_heads, num_v_per_k, head_v_dim)
data_torch = torch.cat([qk_part, v_part], dim=0)
elif ".out_proj." in name:
# wrapper prefixes may already be stripped by the time the name reaches
# modify_tensors, so match on the full name modulo those prefixes. Matching
# on the tensor-kind suffix alone would treat every layer as folded as soon
# as one layer is.
if self._hadamard_folds_tensor(name):
# Hadamard-folded latent: the rotation axis must keep the
# training (grouped) V order; the runtime permutes the
# activation tiled->grouped before the transform instead.
self._hadamard_gdn_v_grouped = True
else:
# Out projection weight: reorder columns (input dimension)
data_torch = self._reorder_v_heads(data_torch, 1, num_k_heads, num_v_per_k, head_v_dim)
yield from super().modify_tensors(data_torch, name, bid)
class _Qwen35MRopeMixin:
# Qwen3.5 always applies interleaved MRoPE (see Qwen3_5RotaryEmbedding in transformers);
# the upstream default mrope_section is [11, 11, 10] and llama.cpp's QWEN35 / QWEN35MOE
# loaders treat qwen35.rope.dimension_sections as required, so make sure it is always
# written even when a particular checkpoint omits the field in `rope_parameters`.
_QWEN35_DEFAULT_MROPE_SECTION = [11, 11, 10, 0]
gguf_writer: gguf.GGUFWriter
rope_parameters: dict
def set_gguf_parameters(self):
super().set_gguf_parameters() # ty: ignore[unresolved-attribute]
if "mrope_section" not in self.rope_parameters:
self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)
@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")
@ModelBase.example("Qwen/Qwen3.5-9B")
class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
model_arch = gguf.MODEL_ARCH.QWEN35
@ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM")
@ModelBase.example("Qwen/Qwen3.5-35B-A3B")
class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
model_arch = gguf.MODEL_ARCH.QWEN35MOE
@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel")
@ModelBase.example("z-lab/Qwen3.5-9B-DFlash")
class DFlashModel(Qwen3Model):
model_arch = gguf.MODEL_ARCH.DFLASH
def set_vocab(self):
if self.target_model_dir is None:
raise ValueError(
"DFlash draft model requires --target-model-dir to be specified. "
"Please provide the path to the target model directory containing the tokenizer."
)
logger.info(f"DFlash: Using tokenizer from target model: {self.target_model_dir}")
original_dir = self.dir_model
self.dir_model = self.target_model_dir
# Reuse the target model's own vocab handler (e.g. Gemma-4 needs its
# own tokenizer logic, not the Qwen default).
from . import get_model_class
with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
target_hparams = json.load(f)
target_arch = target_hparams["architectures"][0]
target_cls = get_model_class(target_arch)
if target_cls is not type(self):
if target_arch == "NemotronHForCausalLM":
setattr(self, "is_moe", "num_experts_per_tok" in target_hparams)
target_cls.set_vocab(self) # ty: ignore[unresolved-attribute]
else:
super().set_vocab()
self.dir_model = original_dir
mask_token_id = self.hparams.get("dflash_config", {}).get("mask_token_id")
if mask_token_id is not None:
self.gguf_writer.add_mask_token_id(mask_token_id)
def set_gguf_parameters(self):
super().set_gguf_parameters()
dflash_config = self.hparams.get("dflash_config", {})
block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16))
self.gguf_writer.add_block_size(block_size)
if "conv_kernel_size" in dflash_config:
self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"]))
self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"]))
self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"]))
self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"]))
output_multiplier = dflash_config.get(
"output_multiplier", self.hparams.get("output_multiplier")
)
if output_multiplier is not None:
self.gguf_writer.add_logit_scale(float(output_multiplier))
softcap = dflash_config.get(
"final_logit_softcapping", self.hparams.get("final_logit_softcapping")
)
if softcap is not None and float(softcap) > 0:
self.gguf_writer.add_final_logit_softcapping(float(softcap))
embedding_scale = dflash_config.get(
"input_embedding_scale", self.hparams.get("input_embedding_scale")
)
if embedding_scale is not None:
self.gguf_writer.add_embedding_scale(float(embedding_scale))
target_layer_ids = dflash_config.get("target_layer_ids", [])
if target_layer_ids:
extract_layer_ids = [i + 1 for i in target_layer_ids]
self.gguf_writer.add_target_layers(extract_layer_ids)
use_sliding_window = self.hparams.get("use_sliding_window", False)
sliding_window = self.hparams.get("sliding_window")
layer_types = self.hparams.get("layer_types")
if use_sliding_window and sliding_window and layer_types:
is_swa = [lt == "sliding_attention" for lt in layer_types]
self.gguf_writer.add_sliding_window(sliding_window)
self.gguf_writer.add_sliding_window_pattern(is_swa)
# M-RoPE target: the draft ropes on the temporal dim only, so write
# degenerate sections [n_rot/2, 0, 0, 0]
if self._target_uses_mrope():
head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0])
def _target_uses_mrope(self) -> bool:
if self.target_model_dir is None:
return False
with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
cfg = json.load(f)
cfg = cfg.get("text_config", cfg)
rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {}
return "mrope_section" in rope
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
name, gen = item
if not name.startswith("model."):
name = "model." + name
return super().filter_tensors((name, gen))
_ROPE_PERMUTE_SUFFIXES = (
"self_attn.q_proj.weight",
"self_attn.k_proj.weight",
"self_attn.q_norm.weight",
"self_attn.k_norm.weight",
)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True):
return
# interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd
if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES):
head_dim = self.hparams["head_dim"]
shape = data_torch.shape
data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape)
if name in (
"model.candidate_selector.predecessor_codebook",
"model.candidate_selector.successor_codebook",
):
name += ".weight"
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Qwen3DFlyModel", "Qwen3DSparkDFlareV2Model")
@ModelBase.example("AngelSlim/Qwen3-8B-DFly-Block8")
class DFlyModel(DFlashModel):
# AngelSpec DFly = DFlash + (a) a per-DRAFT-layer fusion of the raw target-layer features
# on top of the shared context projection, and (b) a TreeFlash predecessor correction
# applied before the target head. There is no Markov/confidence head, so the runtime reads
# it with the DFlash draft reader (block rows 1..n-1), not the DSpark one.
model_arch = gguf.MODEL_ARCH.DFLASH
def __init__(self, dir_model, *args, **kwargs):
hparams = kwargs.pop("hparams", None)
if hparams is None:
hparams = ModelBase.load_hparams(dir_model, False)
# DFly carries target_layer_ids/mask_token_id flat; normalize to DFlash's nested schema
hparams.setdefault("dflash_config", {
k: hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in hparams
})
super().__init__(dir_model, *args, hparams=hparams, **kwargs)
hp = self.hparams
# The published main-branch config of the reference checkpoint states a target depth and
# vocab that do not match the weights (80/120832 against the real 36/151936), which loads
# clean and then mis-projects. Refuse instead of converting a checkpoint that lies.
target_layers = hp.get("target_num_hidden_layers")
layer_ids = hp.get("target_layer_ids") or []
# A declared depth cannot police itself: the reference config's 80 is self-consistent with
# capture ids up to 33 and still wrong. The target model is the only authoritative shape,
# so compare the declared target_* against it and bound the ids by the real depth.
real = self._target_shapes()
for cfg_key, hp_key in (("num_hidden_layers", "target_num_hidden_layers"),
("vocab_size", "target_vocab_size"),
("hidden_size", "target_hidden_size")):
declared = hp.get(hp_key)
if declared is not None and cfg_key in real and int(declared) != real[cfg_key]:
raise ValueError(
f"DFly {hp_key} is {declared} but the target model reports "
f"{cfg_key}={real[cfg_key]}. The config metadata does not describe the target "
"-- pin a known-good revision (the reference checkpoint's is 5712926)."
)
if "num_hidden_layers" in real:
target_layers = real["num_hidden_layers"]
if target_layers and layer_ids and max(layer_ids) >= int(target_layers):
raise ValueError(
f"DFly target_layer_ids {layer_ids} exceed target_num_hidden_layers {target_layers}. "
"The config metadata does not describe the weights -- pin a known-good revision "
"(the reference checkpoint's is 5712926)."
)
# an omitted target_hidden_size would default to the draft width and compare equal to itself,
# so prefer the target model's own hidden_size whenever it is readable
target_hidden = real.get("hidden_size", hp.get("target_hidden_size", hp["hidden_size"]))
if int(target_hidden) != int(hp["hidden_size"]):
raise ValueError(
"DFly residual fusion requires the target hidden_size to equal the drafter's, got "
f"{target_hidden} vs {hp['hidden_size']}."
)
if hp.get("markov_rank") or hp.get("enable_confidence_head"):
raise ValueError(
"DFly does not use the DSpark Markov/confidence head, but this config declares one. "
"A drafter reporting a Markov head is read one block row late by the runtime."
)
self._has_correction = bool(hp.get("enable_hidden_correction", True))
if self._has_correction:
correction_type = hp.get("hidden_correction_type", "swiglu")
if correction_type != "swiglu":
raise ValueError(f"unsupported hidden_correction_type {correction_type!r} (only 'swiglu')")
if not bool(hp.get("dspark_bonus_anchor", True)):
raise ValueError("DFly requires dspark_bonus_anchor=true; slot 0 is the committed anchor")
self._sample_from_anchor = False
def _target_shapes(self) -> dict[str, int]:
"""Shapes read from --target-model-dir, for the keys it declares. Empty when unavailable."""
if self.target_model_dir is None:
return {} # set_vocab raises on this later; nothing authoritative to compare against
try:
with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
cfg = json.load(f)
except (OSError, ValueError):
return {} # unreadable; set_vocab reads the same file and raises
if not isinstance(cfg, dict):
return {}
cfg = {**cfg, **(cfg.get("text_config") or {})}
shapes = {}
for k in ("num_hidden_layers", "vocab_size", "hidden_size"):
try:
shapes[k] = int(cfg[k])
except (KeyError, ValueError, TypeError):
continue # one unusable value must not disable the other comparisons
return shapes
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor)
def prepare_tensors(self):
super().prepare_tensors()
if self._has_correction and not self._seen_correction:
raise ValueError(
"config sets enable_hidden_correction but no hidden_correction.* weights were "
"found; the export is incomplete and would draft without the correction."
)
_seen_correction = False
_dropped_correction = False
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if "hidden_correction." in name:
# the runtime turns correction on from the presence of hidden_correction.down.weight,
# so shipping these while the config disables it silently re-enables the feature
if not self._has_correction:
if not self._dropped_correction:
logger.info("DFly: enable_hidden_correction is false, dropping hidden_correction.* weights")
self._dropped_correction = True
return
self._seen_correction = True
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register(
"Qwen3DSparkModel",
"DSparkDraftModel",
"DSparkSpeculator",
"Lfm2DSparkDraftModel",
"LingDSparkModel",
)
@ModelBase.example("satgeze/Qwen3.6-27B-DSpark")
class DSparkModel(DFlashModel):
# DSpark = DFlash + a semi-autoregressive Markov head.
model_arch = gguf.MODEL_ARCH.DFLASH
def __init__(self, dir_model, *args, **kwargs):
hparams = kwargs.pop("hparams", None)
if hparams is None:
hparams = ModelBase.load_hparams(dir_model, False)
# EAGLE3-style exports use the 1+N bonus-anchor block, DFlash-lineage exports sample from the anchor
self._sample_from_anchor = hparams.get(
"sample_from_anchor",
"transformer_layer_config" not in hparams and "aux_hidden_state_layer_ids" not in hparams)
if "transformer_layer_config" in hparams:
hparams = {**hparams, **hparams["transformer_layer_config"]}
super().__init__(dir_model, *args, hparams=hparams, **kwargs)
# normalize both schemas to DFlash's nested dflash_config
if "aux_hidden_state_layer_ids" in self.hparams:
self.hparams.setdefault("dflash_config", {
"mask_token_id": self.hparams.get("mask_token_id"),
"target_layer_ids": [i - 1 for i in self.hparams["aux_hidden_state_layer_ids"]],
})
else:
self.hparams.setdefault("dflash_config", {
k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams
})
if (markov_head_type := self.hparams.get("markov_head_type", "vanilla")) != "vanilla":
raise ValueError(f"unsupported markov_head_type {markov_head_type!r} (only 'vanilla' is supported)")
n_vocab = self.hparams["vocab_size"]
self._n_vocab_draft = self.hparams.get("draft_vocab_size") or n_vocab
if self._n_vocab_draft > n_vocab:
raise ValueError(f"draft_vocab_size {self._n_vocab_draft} exceeds vocab_size {n_vocab}")
self._d2t: Tensor | None = None
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor)
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
if item[0] == "t2d": # not used at runtime
return None
return super().filter_tensors(item)
_ROPE_PERMUTE_SUFFIXES = (
"self_attn.q_proj.weight",
"self_attn.k_proj.weight",
"self_attn.q_norm.weight",
"self_attn.k_norm.weight",
)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name == "model.d2t":
self._d2t = data_torch
return
if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")):
return
# interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd
if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES):
head_dim = self.hparams["head_dim"]
shape = data_torch.shape
data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape)
yield from super().modify_tensors(data_torch, name, bid)
def prepare_tensors(self):
super().prepare_tensors()
n_vocab = self.hparams["vocab_size"]
if self._n_vocab_draft < n_vocab and self._d2t is None:
raise ValueError(f"draft_vocab_size {self._n_vocab_draft} < vocab_size {n_vocab} but no d2t table found")
# write d2t as absolute target token ids
if self._d2t is not None:
data = LazyTorchTensor.to_eager(self._d2t).to(torch.int64).cpu().numpy().reshape(-1)
if data.size != self._n_vocab_draft:
raise ValueError(f"d2t size {data.size} does not match draft_vocab_size {self._n_vocab_draft}")
data = data + np.arange(data.size, dtype=np.int64)
if np.any((data < 0) | (data >= n_vocab)):
raise ValueError(f"d2t target ids out of range for target vocab size {n_vocab}")
if np.unique(data).size != data.size:
raise ValueError("d2t contains duplicate target ids")
logger.info(f"{'d2t,':<30} --> I64, shape = {{{data.size}}}")
self.gguf_writer.add_tensor("d2t", data, raw_dtype=gguf.GGMLQuantizationType.I64)