char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Daniel Hanmodel: add Qwen3.8-Flash-Next (qwen4exp) (#27742)6c84c7d5d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.3 KiB195 linesraw
1from __future__ import annotations
2
3from typing import Iterable, cast
4
5import torch
6from torch import Tensor
7
8import gguf
9import numpy as np
10
11from .base import ModelBase
12from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin
13from .qwen3vl import Qwen3VLVisionModel
14
15
16@ModelBase.register("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM")
17@ModelBase.example("Qwen/Qwen3.8-Flash-Next")
18class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
19    """Qwen3.8-Flash-Next.
20
21    Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things:
22    hyper-connections in place of every layer norm, QSA sparse attention on the full
23    attention layers, and PLE n-gram hash embeddings on a single layer.
24    """
25
26    model_arch = gguf.MODEL_ARCH.QWEN4EXP
27
28    # the MTP block is a separate draft head; vLLM drops it too
29    supports_mtp_export = False
30    no_mtp = True
31
32    def __init__(self, *args, **kwargs):
33        super().__init__(*args, **kwargs)
34        # only the shard names, so the table itself is never held
35        self._ple_shards: dict[int, str] = {}
36        self._ple_row_dim: int | None = None
37
38    def _read_hash_constants(self, suffix: str) -> list[int]:
39        """Read an int64 PLE constant straight from the checkpoint.
40
41        prepare_tensors() casts every non-float dtype to float32 before
42        modify_tensors() sees it (base.py), which would silently round these
43        45-bit multipliers. Reading the lazy tensor here bypasses that.
44        """
45        for name, gen in self.model_tensors.items():
46            if name.endswith(suffix):
47                t = gen()
48                if t.dtype != torch.int64:
49                    t = t.to(torch.int64)
50                return [int(x) for x in t.tolist()]
51        raise ValueError(f"PLE constant {suffix!r} missing from the checkpoint")
52
53    def set_gguf_parameters(self):
54        super().set_gguf_parameters()
55        hp = self.hparams
56
57        self.gguf_writer.add_hyper_connection_count(hp["hc_count"])
58        self.gguf_writer.add_hyper_connection_low_rank(hp["hc_lowrank"])
59
60        n_layer = hp["num_hidden_layers"]
61        self.gguf_writer.add_indexer_head_count(hp["indexer_n_heads"])
62        self.gguf_writer.add_indexer_key_length(hp["indexer_head_dim"])
63        self.gguf_writer.add_indexer_top_k(hp["indexer_budget"])
64        ratio = hp["indexer_compress_ratio"]
65        layer_types = hp["layer_types"]
66        self.gguf_writer.add_attention_compress_ratios(
67            [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
68        )
69
70        # ple_layer_ids is 1-based in the HF config; empty means no n-gram table,
71        # so emit no PLE keys rather than optional ones
72        ple_layers = [i - 1 for i in hp["ple_layer_ids"]]
73        if not ple_layers:
74            return
75        self.gguf_writer.add_ple_layers(ple_layers)
76        self.gguf_writer.add_ple_ngram_size(hp["ngram_size"])
77        self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"])
78        self.gguf_writer.add_ple_conv_kernel(hp["ple_conv_kernel_size"])
79        self.gguf_writer.add_ple_eos_token_id(self._eos_token_id())
80        # an image is decoded as an embeddings-only batch, so the graph has no placeholder
81        # ids to hash; carry the id and let it stand in for those positions
82        _img = self._image_token_id()
83        if _img is not None:
84            self.gguf_writer.add_ple_image_token_id(int(_img))
85        if self._ple_row_dim is not None:
86            self.gguf_writer.add_embedding_length_per_layer_input(self._ple_row_dim)
87
88        self.gguf_writer.add_ple_layer_multipliers(
89            self._read_hash_constants("ple_embedding.layer_multipliers"))
90        self.gguf_writer.add_ple_head_offsets(
91            self._read_hash_constants("ple_embedding.ngram_heads_offsets"))
92        self.gguf_writer.add_ple_head_vocab_sizes(
93            self._read_hash_constants("ple_embedding.ngram_heads_vocab_sizes"))
94
95    def _image_token_id(self) -> int | None:
96        img = self.hparams.get("image_token_id")
97        return None if img is None else int(img)
98
99    def _eos_token_id(self) -> int:
100        eos = self.hparams.get("eos_token_id")
101        if isinstance(eos, list):
102            # the PLE hash resets n-grams on the primary EOS
103            return int(eos[-1])
104        if eos is None:
105            raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it")
106        return int(eos)
107
108    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
109        # int64 hash constants must stay exact; 1-D tensors force F32, so use KV
110        if name.endswith("ple_embedding.layer_multipliers"):
111            self._ple_multipliers = [int(x) for x in data_torch.tolist()]
112            return []
113        if name.endswith("ple_embedding.ngram_heads_offsets"):
114            self._ple_head_offsets = [int(x) for x in data_torch.tolist()]
115            return []
116        if name.endswith("ple_embedding.ngram_heads_vocab_sizes"):
117            self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()]
118            return []
119
120        if ".ngram_embedding.shard_" in name:
121            return self._place_ple_shard(data_torch, name)
122
123        # one projection feeds indexer q and k; split it, as minimax-m3 does
124        if ".indexer.index_qk_proj.weight" in name:
125            n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
126            q = data_torch[:n_q]
127            k = data_torch[n_q:]
128            return [
129                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q),
130                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k),
131            ]
132
133        # Gemma zero-centred gammas the inherited norm.weight rule misses
134        if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight",
135                          ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")):
136            return [(self.map_tensor_name(name), data_torch + 1)]
137
138        if name.endswith(".ple.conv1d.weight"):
139            return [(self.map_tensor_name(name), data_torch.squeeze())]
140
141        return super().modify_tensors(data_torch, name, bid)
142
143    # the shards concatenate into a tensor of well over 100 GB
144    # use LazyChunkedTensor here, a single shard resident at a time
145    def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]:
146
147        idx = int(name.rpartition(".shard_")[2].partition(".")[0])
148        n_parts = self.hparams["split_ngram_parts"]
149
150        self._ple_shards[idx] = name
151        self._ple_row_dim = int(data_torch.shape[-1])
152
153        if len(self._ple_shards) < n_parts:
154            return []
155
156        # the checkpoint may yield the shards in any order, the row order is by index
157        shards = [self._ple_shards[i] for i in sorted(self._ple_shards)]
158        rows = 0
159        for shard in shards:
160            shape = self.model_tensors[shard]().shape
161            if int(shape[-1]) != self._ple_row_dim:
162                raise ValueError(
163                    f"PLE shard {shard} has row dim {int(shape[-1])}, expected {self._ple_row_dim}")
164            rows += int(shape[0])
165
166        table = gguf.LazyChunkedTensor(
167            [self._load_ple_shard(shard) for shard in shards],
168            shape=(rows, self._ple_row_dim),
169            dtype=np.float32,
170        )
171        gguf_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD]
172        return [(gguf_name + ".weight", cast(Tensor, table))]
173
174    def _load_ple_shard(self, name: str):
175        def load() -> np.ndarray:
176            from .base import LazyTorchTensor
177
178            # a fresh lazy tensor every call, or to_eager() memoizes every shard
179            eager = LazyTorchTensor.to_eager(self.model_tensors[name]())
180            return eager.to(torch.float32).contiguous().numpy()
181        return load
182
183    def prepare_tensors(self):
184        super().prepare_tensors()
185        n_parts = self.hparams.get("split_ngram_parts", 0)
186        if self._ple_shards and len(self._ple_shards) != n_parts:
187            raise ValueError(
188                f"got {len(self._ple_shards)} PLE embedding shards, expected {n_parts}"
189            )
190
191
192@ModelBase.register("Qwen4ExpForConditionalGeneration")
193@ModelBase.example("Qwen/Qwen3.8-Flash-Next")
194class Qwen4ExpVisionModel(Qwen3VLVisionModel):
195    """The vision tower is an unmodified Qwen3-VL ViT."""