diff --git a/pyproject.toml b/pyproject.toml index abbacc60..d8efeb34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,8 @@ extend-exclude = [ "vllm_gguf_plugin/csrc/gguf/dequantize.cuh", "vllm_gguf_plugin/csrc/gguf/vecdotq.cuh", "vllm_gguf_plugin/csrc/gguf/mmq.cuh", + "vllm_gguf_plugin/csrc/gguf/mma_v2.cuh", + "vllm_gguf_plugin/csrc/gguf/mmq_v2.cuh", "vllm_gguf_plugin/csrc/gguf/mmvq.cuh", ] ignore-hidden = false diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 70fd12d6..21f2cb41 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -1,10 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 +import json +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest import torch +import vllm.config.model as model_config_module import vllm.engine.arg_utils as arg_utils_module +import vllm.model_executor.layers.quantization as quantization_module import vllm.model_executor.layers.vocab_parallel_embedding as vocab_embedding_module import vllm.model_executor.parameter as parameter_module import vllm.transformers_utils.config as config_module +from gguf import GGMLQuantizationType as WeightType +from gguf.constants import Keys, VisionProjectorType from transformers import PretrainedConfig from vllm.config.load import LoadConfig from vllm.engine.arg_utils import EngineArgs @@ -12,27 +23,48 @@ WEIGHT_LOADER_V2_SUPPORTED, MergedColumnParallelLinear, QKVParallelLinear, + RowParallelLinear, ) -from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader import get_model_loader from vllm.transformers_utils.config import get_config_parser import vllm_gguf_plugin.config_parser as gguf_config_parser_module +import vllm_gguf_plugin.gguf_tokenizer_builder as gguf_tokenizer_builder_module +import vllm_gguf_plugin.gguf_utils as gguf_utils_module +import vllm_gguf_plugin.plugin as gguf_plugin_module import vllm_gguf_plugin.quantization as gguf_quantization +import vllm_gguf_plugin.quantization.params as gguf_params_module +import vllm_gguf_plugin.weights_adapter.default as default_adapter_module +import vllm_gguf_plugin.weights_adapter.qwen3_5 as qwen3_5_adapter_module from vllm_gguf_plugin import OOTGGUFConfig, OOTGGUFModelLoader, register from vllm_gguf_plugin.config_parser import GGUFConfigParser +from vllm_gguf_plugin.gguf_tokenizer_builder import build_tokenizer_from_gguf +from vllm_gguf_plugin.gguf_utils import ( + _gguf_sequence_edge, + extract_vision_config_from_gguf, + maybe_patch_hf_config_from_gguf, + resolve_gguf_config_source, +) from vllm_gguf_plugin.quantization import ( GGUFUninitializedParameter, GGUFWeightParameter, GGUFWeightTypeParameter, ) +from vllm_gguf_plugin.weights_adapter.default import ( + GGUFWeightsAdapter, + _add_gemma4_gguf_mappings, + _add_gemma4_mtp_gguf_mappings, + _add_qwen3_5_mtp_gguf_mappings, +) +from vllm_gguf_plugin.weights_adapter.gemma4 import Gemma4GGUFAdapter +from vllm_gguf_plugin.weights_adapter.qwen3_5 import Qwen3_5GGUFAdapter def test_register_overrides_gguf_config(): register() - quant_config = get_quantization_config("gguf") + quant_config = quantization_module.get_quantization_config("gguf") assert quant_config is OOTGGUFConfig @@ -49,13 +81,47 @@ def test_register_is_idempotent(): register() register() - assert get_quantization_config("gguf") is OOTGGUFConfig + assert quantization_module.get_quantization_config("gguf") is OOTGGUFConfig assert isinstance( get_model_loader(LoadConfig(load_format="gguf")), OOTGGUFModelLoader ) assert isinstance(get_config_parser("gguf"), GGUFConfigParser) +@pytest.mark.parametrize( + "script", + [ + """ +import torch +import vllm_gguf_plugin +import vllm.model_executor.layers.quantization.gguf +assert hasattr(torch.ops.vllm_gguf_plugin, "_fused_mul_mat_gguf") +assert hasattr(torch.ops.vllm_gguf_plugin, "_fused_moe_gguf") +assert hasattr(torch.ops.vllm_gguf_plugin, "_apply_gguf_embedding") +""", + """ +import torch +import vllm.model_executor.layers.quantization.gguf +import vllm_gguf_plugin +assert hasattr(torch.ops.vllm_gguf_plugin, "_fused_mul_mat_gguf") +assert hasattr(torch.ops.vllm_gguf_plugin, "_fused_moe_gguf") +assert hasattr(torch.ops.vllm_gguf_plugin, "_apply_gguf_embedding") +""", + ], +) +def test_plugin_custom_ops_do_not_conflict_with_core_gguf_import(script): + subprocess.run([sys.executable, "-c", script], check=True) + + +def test_register_patches_model_config_gguf_helper(): + register() + + assert ( + model_config_module.maybe_patch_hf_config_from_gguf + is gguf_utils_module.maybe_patch_hf_config_from_gguf + ) + + def test_oot_config_reuses_in_tree_behavior(): quant_config = OOTGGUFConfig.from_config({}) @@ -64,6 +130,27 @@ def test_oot_config_reuses_in_tree_behavior(): assert repr(quant_config) == "GGUFConfig()" +def test_gguf_override_quantization_method_accepts_hf_config_keyword(): + register() + + # hf_config keyword matches core QuantizationConfig.override_quantization_method + # signature. + # This was added to fix a TypeError when ModelConfig._verify_quantization() calls + # override_quantization_method(hf_quant_cfg, user_quant, hf_config=hf_config). + result_explicit = OOTGGUFConfig.override_quantization_method( + {}, "gguf", hf_config=object() + ) + assert result_explicit == "gguf" + + result_non_gguf = OOTGGUFConfig.override_quantization_method( + {}, "awq", hf_config=object() + ) + assert result_non_gguf is None + + result_none = OOTGGUFConfig.override_quantization_method({}, None) + assert result_none is None + + def test_gguf_linear_uses_weight_loader_v2(monkeypatch): register() monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) @@ -101,6 +188,40 @@ def test_gguf_linear_uses_weight_loader_v2(monkeypatch): assert layer.qweight_type.shard_weight_type == {0: 3, 1: 4} +def test_gguf_linear_keeps_multi_shards_separate(monkeypatch): + register() + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + layer = MergedColumnParallelLinear( + input_size=4, + output_sizes=[4, 4], + bias=False, + quant_config=OOTGGUFConfig.from_config({}), + disable_tp=True, + ) + layer.weight_loader_v2(layer.qweight, torch.ones((4, 4), dtype=torch.uint8), 0) + layer.weight_loader_v2(layer.qweight, 2 * torch.ones((4, 4), dtype=torch.uint8), 1) + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), 0) + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), 1) + + layer.quant_method.process_weights_after_loading(layer) + + assert layer.qweight.numel() == 0 + assert layer.qweight.shard_id == [0, 1] + assert layer.qweight.shard_id_map == {0: 0, 1: 1} + assert layer.qweight.shard_offset_map == {0: (0, 4, 4), 1: (4, 8, 4)} + assert len(layer.qweight.data_container) == 2 + assert torch.equal( + layer.qweight.data_container[0], torch.ones((4, 4), dtype=torch.uint8) + ) + assert torch.equal( + layer.qweight.data_container[1], 2 * torch.ones((4, 4), dtype=torch.uint8) + ) + + def test_gguf_embedding_uses_plugin_weight_loader(monkeypatch): monkeypatch.setattr( vocab_embedding_module, "get_tensor_model_parallel_rank", lambda: 0 @@ -172,145 +293,2184 @@ def fake_fused_mul_mat_gguf(x, qweight, qweight_type): ) out = layer.quant_method.apply(layer, torch.ones((2, 4), dtype=torch.float32)) - assert calls == [((8, 4), 3)] + assert calls == [((4, 4), 3), ((4, 4), 3)] assert out.shape == (2, 8) -def test_gguf_config_parser_uses_parent_dir_for_local_file(tmp_path, monkeypatch): - gguf_path = tmp_path / "model.gguf" - gguf_path.write_bytes(b"GGUF") - calls = {} +def test_gguf_tuple_shard_loader_splits_fused_qweight(monkeypatch): + register() + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + layer = MergedColumnParallelLinear( + input_size=4, + output_sizes=[4, 2, 6, 8], + bias=False, + quant_config=OOTGGUFConfig.from_config({}), + disable_tp=True, + ) + + fused_qkv = torch.cat( + [ + torch.full((4, 4), 1, dtype=torch.uint8), + torch.full((2, 4), 2, dtype=torch.uint8), + torch.full((6, 4), 3, dtype=torch.uint8), + ], + dim=0, + ) + layer.qweight.weight_loader(layer.qweight, fused_qkv, (0, 1, 2)) + layer.qweight.weight_loader( + layer.qweight, torch.full((8, 4), 4, dtype=torch.uint8), 3 + ) + layer.qweight_type.weight_loader( + layer.qweight_type, + torch.tensor(WeightType.Q4_0, dtype=torch.uint8), + (0, 1, 2), + ) + layer.qweight_type.weight_loader( + layer.qweight_type, torch.tensor(WeightType.Q4_1, dtype=torch.uint8), 3 + ) + + layer.quant_method.process_weights_after_loading(layer) + + assert layer.qweight.shard_id == [0, 1, 2, 3] + assert layer.qweight.shard_offset_map == { + 0: (0, 4, 4), + 1: (4, 6, 4), + 2: (6, 12, 4), + 3: (12, 20, 4), + } + assert layer.qweight.numel() == 0 + assert torch.equal( + layer.qweight.data_container[0], torch.full((4, 4), 1, dtype=torch.uint8) + ) + assert torch.equal( + layer.qweight.data_container[1], torch.full((2, 4), 2, dtype=torch.uint8) + ) + assert torch.equal( + layer.qweight.data_container[2], torch.full((6, 4), 3, dtype=torch.uint8) + ) + assert torch.equal( + layer.qweight.data_container[3], torch.full((8, 4), 4, dtype=torch.uint8) + ) + assert layer.qweight_type.shard_weight_type == { + 0: WeightType.Q4_0, + 1: WeightType.Q4_0, + 2: WeightType.Q4_0, + 3: WeightType.Q4_1, + } - def fake_parse( - self, model, trust_remote_code, revision=None, code_revision=None, **kwargs - ): - calls["model"] = model - calls["trust_remote_code"] = trust_remote_code - return {}, PretrainedConfig(model_type="qwen3_moe") +def test_gguf_row_parallel_weight_loader_v2_omits_empty_shard_id(monkeypatch): + register() + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr( - gguf_config_parser_module.HFConfigParser, - "parse", - fake_parse, + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 ) + monkeypatch.setattr(gguf_params_module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr( - gguf_config_parser_module, - "maybe_patch_hf_config_from_gguf", - lambda model, config: config, + gguf_params_module, "get_tensor_model_parallel_world_size", lambda: 1 ) - config_dict, config = GGUFConfigParser().parse(gguf_path, trust_remote_code=False) + layer = RowParallelLinear( + input_size=4, + output_size=8, + bias=False, + quant_config=OOTGGUFConfig.from_config({}), + disable_tp=True, + ) - assert calls["model"] == gguf_path.parent - assert calls["trust_remote_code"] is False - assert config_dict["norm_topk_prob"] is True - assert config.architectures == ["Qwen3MoeForCausalLM"] + layer.qweight.weight_loader(layer.qweight, torch.ones((8, 4), dtype=torch.uint8)) + assert torch.equal(layer.qweight.data, torch.ones((8, 4), dtype=torch.uint8)) -def test_register_sets_engine_args_for_gguf_model(monkeypatch): - register() - captured = {} - def fake_model_config(**kwargs): - captured.update(kwargs) - return kwargs +def test_gemma4_adapter_transforms_quantized_moe_names(): + adapter = Gemma4GGUFAdapter(PretrainedConfig(model_type="gemma4")) + weight = torch.empty((2, 3), dtype=torch.uint8) - monkeypatch.setattr(arg_utils_module, "ModelConfig", fake_model_config) - engine_args = EngineArgs(model="/tmp/model.gguf", tokenizer="/tmp/tokenizer") + transformed = dict( + adapter.map_weights( + [ + ( + "model.language_model.layers.0.experts.gate_up_proj.qweight", + weight, + ), + ( + "model.language_model.layers.0.experts.gate_up_proj.qweight_type", + torch.tensor(1, dtype=torch.uint8), + ), + ( + "model.language_model.layers.0.experts.down_proj.qweight", + weight, + ), + ( + "model.language_model.layers.0.experts.down_proj.qweight_type", + torch.tensor(1, dtype=torch.uint8), + ), + ] + ) + ) - engine_args.create_model_config() + assert ( + "model.language_model.layers.0.moe.experts.routed_experts.w13_qweight" + in transformed + ) + assert ( + "model.language_model.layers.0.moe.experts.routed_experts.w13_qweight_type" + in transformed + ) + assert ( + "model.language_model.layers.0.moe.experts.routed_experts.w2_qweight" + in transformed + ) + assert ( + "model.language_model.layers.0.moe.experts.routed_experts.w2_qweight_type" + in transformed + ) - assert captured["config_format"] == "gguf" - assert captured["model"] == "/tmp/tokenizer" - assert captured["model_weights"] == "/tmp/model.gguf" - assert captured["quantization"] == "gguf" - assert engine_args.load_format == "gguf" +def test_gemma4_gguf_mappings_match_current_hf_names(): + config = PretrainedConfig(model_type="gemma4", num_hidden_layers=2) + config.vision_config = PretrainedConfig(num_hidden_layers=2) + mapping: dict[str, str] = {} + sideload_params = [] -def test_register_skips_speculator_probe_for_gguf(): - register() + _add_gemma4_gguf_mappings(config, mapping, sideload_params) - model, tokenizer, speculative_config = ( - config_module.maybe_override_with_speculators( - model="/tmp/model.gguf", - tokenizer="/tmp/tokenizer", - trust_remote_code=False, - revision=None, - vllm_speculative_config={"foo": "bar"}, - hf_token=None, - ) + assert mapping["blk.1.ffn_gate_inp.scale"] == ( + "model.language_model.layers.1.router.scale" + ) + assert mapping["blk.1.ffn_gate_inp.weight"] == ( + "model.language_model.layers.1.router.proj.weight" + ) + assert mapping["blk.1.ffn_down_exps.scale"] == ( + "model.language_model.layers.1.router.per_expert_scale" + ) + assert mapping["blk.1.ffn_gate_up_exps.weight"] == ( + "model.language_model.layers.1.experts.gate_up_proj.weight" + ) + assert mapping["blk.1.ffn_down_exps.weight"] == ( + "model.language_model.layers.1.experts.down_proj.weight" + ) + assert mapping["v.blk.1.ln1.weight"] == ( + "model.vision_tower.encoder.layers.1.input_layernorm.weight" + ) + assert mapping["v.blk.1.ln2.weight"] == ( + "model.vision_tower.encoder.layers.1.pre_feedforward_layernorm.weight" ) - assert model == "/tmp/model.gguf" - assert tokenizer == "/tmp/tokenizer" - assert speculative_config == {"foo": "bar"} +def test_gemma4_text_only_does_not_add_vision_projector_mappings(): + config = PretrainedConfig( + architectures=["Gemma4ForCausalLM"], + model_type="gemma4", + num_hidden_layers=2, + ) + mapping: dict[str, str] = {} + sideload_params = [] -def test_gguf_qkv_shards_are_padded_in_qkv_order(monkeypatch): - register() - monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + _add_gemma4_gguf_mappings(config, mapping, sideload_params) + + assert mapping["blk.1.ffn_gate_inp.scale"] == ( + "model.layers.1.router.scale" + ) + assert mapping["blk.1.ffn_gate_up_exps.weight"] == ( + "model.layers.1.experts.gate_up_proj.weight" + ) + assert "v.std_bias" not in mapping + assert "v.patch_embd.weight" not in mapping + assert "mm.input_projection.weight" not in mapping + assert "v.blk.1.ln1.weight" not in mapping + + +def test_gemma4_causal_lm_with_vision_config_uses_text_layout(): + config = PretrainedConfig( + architectures=["Gemma4ForCausalLM"], + model_type="gemma4", + num_hidden_layers=2, + ) + config.vision_config = PretrainedConfig(num_hidden_layers=2) + mapping: dict[str, str] = {} + sideload_params = [] + + _add_gemma4_gguf_mappings(config, mapping, sideload_params) + + assert mapping["blk.1.ffn_gate_inp.scale"] == ( + "model.layers.1.router.scale" + ) + assert "v.std_bias" not in mapping + assert "v.patch_embd.weight" not in mapping + + +def test_gemma4_adapter_flattens_patch_embed_weight(): + adapter = Gemma4GGUFAdapter(PretrainedConfig(model_type="gemma4")) + weight = torch.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5) + + transformed = adapter.transform_weight( + "model.vision_tower.patch_embedder.input_proj.weight", + weight, + ) + + assert transformed.shape == (2, 60) + assert torch.equal(transformed, weight.flatten(1)) + + +def test_gemma4_mtp_gguf_mappings(): + config = PretrainedConfig(model_type="gemma4_assistant", num_hidden_layers=2) + mapping: dict[str, str] = {} + + _add_gemma4_mtp_gguf_mappings(config, mapping) + + assert mapping["token_embd.weight"] == "model.embed_tokens.weight" + assert mapping["nextn.pre_projection.weight"] == "model.pre_projection.weight" + assert mapping["nextn.post_projection.weight"] == "model.post_projection.weight" + assert mapping["blk.1.attn_q.weight"] == ("model.layers.1.self_attn.q_proj.weight") + assert "blk.1.attn_k.weight" not in mapping + assert "blk.1.attn_v.weight" not in mapping + assert mapping["blk.1.ffn_gate.weight"] == ("model.layers.1.mlp.gate_proj.weight") + assert mapping["blk.1.layer_output_scale.weight"] == "model.layers.1.layer_scalar" + + +def test_gguf_sequence_edge_accepts_scalar_and_sequence_values(): + assert _gguf_sequence_edge(None, first=True) is None + assert _gguf_sequence_edge(8, first=True) == 8 + assert _gguf_sequence_edge(8, first=False) == 8 + assert _gguf_sequence_edge([8, 8, 8, 2], first=True) == 8 + assert _gguf_sequence_edge([8, 8, 8, 2], first=False) == 2 + + +class _FakeGGUFField: + def __init__(self, value): + self.value = value + self.parts = [value] + + def contents(self): + return self.value + + +class _FakeGGUFReader: + def __init__(self, fields): + self.fields = {key: _FakeGGUFField(value) for key, value in fields.items()} + self.tensors = [] + + def get_field(self, key): + return self.fields.get(key) + + +def test_build_tokenizer_from_gguf_metadata_uses_arch_alias_and_cache( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + mmproj_path = tmp_path / "mmproj.gguf" + gguf_path.write_bytes(b"GGUF") + mmproj_path.write_bytes(b"GGUF") + qwen_mm_tokens = [ + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>", + ] + qwen_control_tokens = [ + "", + "", + "", + "", + ] + main_reader = _FakeGGUFReader( + { + "general.architecture": "qwen35moe", + "tokenizer.ggml.tokens": [ + "", + "", + "", + "hello", + *qwen_mm_tokens, + *qwen_control_tokens, + "[PAD000]", + ], + "tokenizer.ggml.token_type": [ + 3, + 3, + 3, + 1, + *([3] * len(qwen_mm_tokens)), + *([4] * len(qwen_control_tokens)), + 5, + ], + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.merges": ["h ello"], + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 2, + "tokenizer.ggml.padding_token_id": 0, + "tokenizer.chat_template": "{{ messages }}", + } + ) + mmproj_reader = _FakeGGUFReader( + { + "general.architecture": "clip", + "general.type": "mmproj", + "clip.vision.patch_size": 16, + "clip.vision.spatial_merge_size": 2, + "clip.vision.temporal_patch_size": 2, + } + ) + calls = [] + + class FakeTokenizer: + def __init__(self, *args, **kwargs): + calls.append(("fast", kwargs)) + self.chat_template = None + + def save_pretrained(self, path): + (path / "tokenizer.json").write_text("{}", encoding="utf-8") + (path / "tokenizer_config.json").write_text("{}", encoding="utf-8") + + def fake_convert(architecture, tokenizer_dict): + calls.append(("convert", architecture, tokenizer_dict)) + return object(), {} + + def fake_gguf_reader(path): + return mmproj_reader if Path(path) == mmproj_path else main_reader + + monkeypatch.setenv( + "VLLM_GGUF_TOKENIZER_CACHE", + str(tmp_path / "tokenizer-cache"), + ) monkeypatch.setattr( - parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + gguf_tokenizer_builder_module.gguf, + "GGUFReader", + fake_gguf_reader, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "convert_gguf_tokenizer", + fake_convert, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "PreTrainedTokenizerFast", + FakeTokenizer, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "detect_gguf_multimodal", + lambda model: mmproj_path, ) - layer = QKVParallelLinear( - hidden_size=4, - head_size=2, - total_num_heads=2, - total_num_kv_heads=1, - bias=False, - quant_config=OOTGGUFConfig.from_config({}), - disable_tp=True, + tokenizer_path = build_tokenizer_from_gguf(gguf_path) + + assert tokenizer_path is not None + tokenizer_cache = Path(tokenizer_path) + assert (tmp_path / "tokenizer-cache").is_dir() + processor_config = json.loads( + (tokenizer_cache / "processor_config.json").read_text(encoding="utf-8") + ) + preprocessor_config = json.loads( + (tokenizer_cache / "preprocessor_config.json").read_text(encoding="utf-8") + ) + video_config = json.loads( + (tokenizer_cache / "video_preprocessor_config.json").read_text(encoding="utf-8") ) + assert processor_config["processor_class"] == "Qwen3VLProcessor" + assert processor_config["image_processor"]["patch_size"] == 16 + assert preprocessor_config["image_processor_type"] == "Qwen2VLImageProcessor" + assert preprocessor_config["merge_size"] == 2 + assert video_config["video_processor_type"] == "Qwen3VLVideoProcessor" + tokenizer_config = json.loads( + (tokenizer_cache / "tokenizer_config.json").read_text(encoding="utf-8") + ) + assert tokenizer_config["additional_special_tokens"] == [ + *qwen_mm_tokens, + *qwen_control_tokens, + ] + assert tokenizer_config["image_token"] == "<|image_pad|>" + assert tokenizer_config["video_token"] == "<|video_pad|>" + assert calls[0][0] == "convert" + assert calls[0][1] == "qwen3_moe" + assert calls[0][2]["tokens"] == [ + "", + "", + "", + "hello", + *qwen_mm_tokens, + *qwen_control_tokens, + "[PAD000]", + ] + assert calls[0][2]["token_type"] == [ + 3, + 3, + 3, + 1, + *([3] * len(qwen_mm_tokens)), + *([4] * len(qwen_control_tokens)), + 5, + ] + assert calls[1][1]["bos_token"] == "" + assert calls[1][1]["eos_token"] == "" + assert calls[1][1]["pad_token"] == "" - q = torch.full((4, 4), 1, dtype=torch.uint8) - k = torch.full((2, 4), 2, dtype=torch.uint8) - v = torch.full((2, 4), 3, dtype=torch.uint8) - # Load out of canonical order to match GGUF tensor iteration order. - layer.weight_loader_v2(layer.qweight, k, "k") - layer.weight_loader_v2(layer.qweight, q, "q") - layer.weight_loader_v2(layer.qweight, v, "v") - layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), "k") - layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), "q") - layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), "v") + def fail_convert(*args, **kwargs): + raise AssertionError("cached tokenizer should not call converter again") - layer.quant_method.process_weights_after_loading(layer) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "convert_gguf_tokenizer", + fail_convert, + ) + (tokenizer_cache / "tokenizer_config.json").write_text("{}", encoding="utf-8") + (tokenizer_cache / "preprocessor_config.json").unlink() + assert build_tokenizer_from_gguf(gguf_path) == tokenizer_path + assert (tokenizer_cache / "preprocessor_config.json").is_file() + cached_tokenizer_config = json.loads( + (tokenizer_cache / "tokenizer_config.json").read_text(encoding="utf-8") + ) + assert cached_tokenizer_config["additional_special_tokens"] == [ + *qwen_mm_tokens, + *qwen_control_tokens, + ] - assert layer.qweight.shard_id == ["q", "k", "v"] - assert layer.qweight.shard_offset_map == { - "q": (0, 4, 4), - "k": (4, 6, 4), - "v": (6, 8, 4), - } - assert torch.equal(layer.qweight[:4], q) - assert torch.equal(layer.qweight[4:6], k) - assert torch.equal(layer.qweight[6:8], v) +def test_build_tokenizer_from_qwen35_gguf_uses_dense_arch_alias( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + main_reader = _FakeGGUFReader( + { + "general.architecture": "qwen35", + "tokenizer.ggml.tokens": ["", "", "", "hello"], + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.merges": ["h ello"], + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 2, + "tokenizer.ggml.padding_token_id": 0, + } + ) + calls = [] -def test_gguf_linear_preserves_cuda_weight_device(monkeypatch): - if not torch.cuda.is_available(): - return + class FakeTokenizer: + def __init__(self, *args, **kwargs): + calls.append(("fast", kwargs)) + self.chat_template = None - register() - monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + def save_pretrained(self, path): + (path / "tokenizer.json").write_text("{}", encoding="utf-8") + (path / "tokenizer_config.json").write_text("{}", encoding="utf-8") + + def fake_convert(architecture, tokenizer_dict): + calls.append(("convert", architecture, tokenizer_dict)) + return object(), {} + + monkeypatch.setenv( + "VLLM_GGUF_TOKENIZER_CACHE", + str(tmp_path / "tokenizer-cache"), + ) monkeypatch.setattr( - parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + gguf_tokenizer_builder_module.gguf, + "GGUFReader", + lambda path: main_reader, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "convert_gguf_tokenizer", + fake_convert, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "PreTrainedTokenizerFast", + FakeTokenizer, ) - with torch.device("cuda"): - layer = MergedColumnParallelLinear( - input_size=4, - output_sizes=[4, 4], - bias=False, - quant_config=OOTGGUFConfig.from_config({}), - disable_tp=True, - ) + tokenizer_path = build_tokenizer_from_gguf(gguf_path) - layer.weight_loader_v2(layer.qweight, torch.ones((4, 4), dtype=torch.uint8), 0) - layer.weight_loader_v2(layer.qweight, 2 * torch.ones((4, 4), dtype=torch.uint8), 1) - layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), 0) - layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), 1) - layer.quant_method.process_weights_after_loading(layer) + assert tokenizer_path is not None + assert calls[0][0] == "convert" + assert calls[0][1] == "qwen3" - assert layer.qweight.device.type == "cuda" - assert layer.qweight_type.device.type == "cuda" + +def test_build_tokenizer_from_gguf_returns_none_when_cache_key_stat_fails( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "missing.gguf" + + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "check_gguf_file", + lambda model: True, + ) + + def fail_reader(*args, **kwargs): + raise AssertionError("stat failure must return before reading GGUF") + + monkeypatch.setattr( + gguf_tokenizer_builder_module.gguf, + "GGUFReader", + fail_reader, + ) + + assert build_tokenizer_from_gguf(gguf_path) is None + + +def test_build_tokenizer_from_gguf_copies_local_sidecars_first( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + local_preprocessor = {"processor_class": "LocalProcessor"} + (tmp_path / "preprocessor_config.json").write_text( + json.dumps(local_preprocessor), + encoding="utf-8", + ) + fake_reader = _FakeGGUFReader( + { + "general.architecture": "gemma4", + "tokenizer.ggml.tokens": ["", "", "", "hello"], + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.merges": ["h ello"], + } + ) + + class FakeTokenizer: + def __init__(self, *args, **kwargs): + pass + + def save_pretrained(self, path): + (path / "tokenizer.json").write_text("{}", encoding="utf-8") + + monkeypatch.setenv( + "VLLM_GGUF_TOKENIZER_CACHE", + str(tmp_path / "tokenizer-cache"), + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module.gguf, + "GGUFReader", + lambda path: fake_reader, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "convert_gguf_tokenizer", + lambda architecture, tokenizer_dict: (object(), {}), + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "PreTrainedTokenizerFast", + FakeTokenizer, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "detect_gguf_multimodal", + lambda model: None, + ) + + tokenizer_path = build_tokenizer_from_gguf(gguf_path) + + assert tokenizer_path is not None + copied = json.loads( + (Path(tokenizer_path) / "preprocessor_config.json").read_text(encoding="utf-8") + ) + assert copied == local_preprocessor + + +def test_build_tokenizer_from_gguf_patches_gemma4_special_tokens( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + fake_reader = _FakeGGUFReader( + { + "general.architecture": "gemma4", + "tokenizer.ggml.tokens": [ + "", + "", + "", + "<|image>", + "<|image|>", + "", + "<|audio>", + "<|audio|>", + "", + "<|video|>", + ], + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.merges": ["h ello"], + } + ) + + class FakeTokenizer: + def __init__(self, *args, **kwargs): + pass + + def save_pretrained(self, path): + (path / "tokenizer.json").write_text("{}", encoding="utf-8") + (path / "tokenizer_config.json").write_text("{}", encoding="utf-8") + + monkeypatch.setenv( + "VLLM_GGUF_TOKENIZER_CACHE", + str(tmp_path / "tokenizer-cache"), + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module.gguf, + "GGUFReader", + lambda path: fake_reader, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "convert_gguf_tokenizer", + lambda architecture, tokenizer_dict: (object(), {}), + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "PreTrainedTokenizerFast", + FakeTokenizer, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "detect_gguf_multimodal", + lambda model: None, + ) + + tokenizer_path = build_tokenizer_from_gguf(gguf_path) + + assert tokenizer_path is not None + tokenizer_config = json.loads( + (Path(tokenizer_path) / "tokenizer_config.json").read_text(encoding="utf-8") + ) + assert tokenizer_config["processor_class"] == "Gemma4Processor" + assert tokenizer_config["model_specific_special_tokens"]["image_token"] == ( + "<|image|>" + ) + assert tokenizer_config["model_specific_special_tokens"]["boi_token"] == ( + "<|image>" + ) + assert tokenizer_config["model_specific_special_tokens"]["eoi_token"] == ( + "" + ) + assert tokenizer_config["extra_special_tokens"]["image_token"] == "<|image|>" + assert tokenizer_config["extra_special_tokens"]["boi_token"] == "<|image>" + assert tokenizer_config["extra_special_tokens"]["eoi_token"] == "" + assert tokenizer_config["extra_special_tokens"]["audio_token"] == "<|audio|>" + assert tokenizer_config["extra_special_tokens"]["video_token"] == "<|video|>" + + tokenizer_config_path = Path(tokenizer_path) / "tokenizer_config.json" + tokenizer_config_path.write_text("{}", encoding="utf-8") + assert build_tokenizer_from_gguf(gguf_path) == tokenizer_path + tokenizer_config = json.loads(tokenizer_config_path.read_text(encoding="utf-8")) + assert tokenizer_config["extra_special_tokens"]["image_token"] == "<|image|>" + assert tokenizer_config["extra_special_tokens"]["boi_token"] == "<|image>" + assert tokenizer_config["extra_special_tokens"]["eoi_token"] == "" + + +def test_build_tokenizer_from_gguf_prefers_local_config_special_token_ids( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + (tmp_path / "config.json").write_text( + json.dumps( + { + "eos_token_id": 4, + "pad_token_id": 0, + "text_config": { + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + }, + } + ), + encoding="utf-8", + ) + fake_reader = _FakeGGUFReader( + { + "general.architecture": "gemma4", + "tokenizer.ggml.tokens": [ + "", + "", + "", + "hello", + "", + ], + "tokenizer.ggml.token_type": [3, 3, 3, 1, 3], + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.merges": ["h ello"], + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 2, + "tokenizer.ggml.padding_token_id": 0, + } + ) + calls = [] + + class FakeTokenizer: + def __init__(self, *args, **kwargs): + calls.append(kwargs) + + def save_pretrained(self, path): + (path / "tokenizer.json").write_text("{}", encoding="utf-8") + (path / "tokenizer_config.json").write_text("{}", encoding="utf-8") + + monkeypatch.setenv( + "VLLM_GGUF_TOKENIZER_CACHE", + str(tmp_path / "tokenizer-cache"), + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module.gguf, + "GGUFReader", + lambda path: fake_reader, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "convert_gguf_tokenizer", + lambda architecture, tokenizer_dict: (object(), {}), + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "PreTrainedTokenizerFast", + FakeTokenizer, + ) + monkeypatch.setattr( + gguf_tokenizer_builder_module, + "detect_gguf_multimodal", + lambda model: None, + ) + + tokenizer_path = build_tokenizer_from_gguf(gguf_path) + + assert tokenizer_path is not None + assert calls[0]["bos_token"] == "" + assert calls[0]["eos_token"] == "" + assert calls[0]["pad_token"] == "" + tokenizer_config_path = Path(tokenizer_path) / "tokenizer_config.json" + tokenizer_config = json.loads(tokenizer_config_path.read_text(encoding="utf-8")) + assert tokenizer_config["eos_token"] == "" + assert "" not in tokenizer_config.get("additional_special_tokens", []) + + tokenizer_config_path.write_text( + json.dumps({"additional_special_tokens": ["", "<|think|>"]}), + encoding="utf-8", + ) + assert build_tokenizer_from_gguf(gguf_path) == tokenizer_path + cached_tokenizer_config = json.loads( + tokenizer_config_path.read_text(encoding="utf-8") + ) + assert cached_tokenizer_config["eos_token"] == "" + assert "" not in cached_tokenizer_config.get( + "additional_special_tokens", + [], + ) + assert "<|think|>" in cached_tokenizer_config.get( + "additional_special_tokens", + [], + ) + + +def test_extract_vision_config_accepts_single_value_metadata(monkeypatch): + fake_reader = _FakeGGUFReader( + { + Keys.Clip.PROJECTOR_TYPE: VisionProjectorType.GEMMA3, + Keys.ClipVision.EMBEDDING_LENGTH: [1152], + Keys.ClipVision.FEED_FORWARD_LENGTH: [4304], + Keys.ClipVision.BLOCK_COUNT: [27], + Keys.ClipVision.Attention.HEAD_COUNT: [16], + Keys.ClipVision.IMAGE_SIZE: [896], + Keys.ClipVision.PATCH_SIZE: [14], + Keys.ClipVision.Attention.LAYERNORM_EPS: [1e-6], + } + ) + monkeypatch.setattr( + gguf_utils_module.gguf, + "GGUFReader", + lambda path: fake_reader, + ) + + config = extract_vision_config_from_gguf("mmproj.gguf") + + assert config is not None + assert config.hidden_size == 1152 + assert config.intermediate_size == 4304 + assert config.num_hidden_layers == 27 + assert config.vision_use_head is False + + +def test_qwen35moe_gguf_config_is_normalized_for_mm(monkeypatch): + fake_reader = _FakeGGUFReader( + { + "general.architecture": "qwen35moe", + "qwen35moe.attention.key_length": 256, + "qwen35moe.full_attention_interval": 4, + "qwen35moe.nextn_predict_layers": 1, + "qwen35moe.block_count": 41, + "qwen35moe.rope.dimension_count": 64, + "qwen35moe.rope.dimension_sections": [11, 11, 10, 0], + "qwen35moe.rope.freq_base": 10000000.0, + } + ) + monkeypatch.setattr(gguf_utils_module, "check_gguf_file", lambda model: True) + monkeypatch.setattr( + gguf_utils_module, + "extract_vocab_size_from_gguf", + lambda model: None, + ) + monkeypatch.setattr( + gguf_utils_module, + "extract_lm_head_from_gguf", + lambda model: None, + ) + monkeypatch.setattr( + gguf_utils_module, + "detect_gguf_multimodal", + lambda model: "mmproj.gguf", + ) + monkeypatch.setattr( + gguf_utils_module.gguf, + "GGUFReader", + lambda path: fake_reader, + ) + + config = maybe_patch_hf_config_from_gguf( + "model.gguf", + PretrainedConfig(model_type="qwen35moe"), + ) + + assert config.model_type == "qwen3_5_moe" + assert config.architectures == ["Qwen3_5MoeForConditionalGeneration"] + assert config.mtp_num_hidden_layers == 1 + assert config.num_nextn_predict_layers == 1 + assert config.num_hidden_layers == 40 + assert config.full_attention_interval == 4 + assert config.rope_parameters["mrope_section"] == [11, 11, 10] + assert config.rope_parameters["mrope_interleaved"] is True + + +def test_qwen35_gguf_config_subtracts_nextn_layers(monkeypatch): + fake_reader = _FakeGGUFReader( + { + "general.architecture": "qwen35", + "qwen35.attention.key_length": 256, + "qwen35.full_attention_interval": 4, + "qwen35.nextn_predict_layers": 1, + "qwen35.block_count": 65, + "qwen35.rope.dimension_count": 64, + "qwen35.rope.dimension_sections": [11, 11, 10, 0], + "qwen35.rope.freq_base": 10000000.0, + } + ) + monkeypatch.setattr(gguf_utils_module, "check_gguf_file", lambda model: True) + monkeypatch.setattr( + gguf_utils_module, + "extract_vocab_size_from_gguf", + lambda model: None, + ) + monkeypatch.setattr( + gguf_utils_module, + "extract_lm_head_from_gguf", + lambda model: None, + ) + monkeypatch.setattr( + gguf_utils_module, + "detect_gguf_multimodal", + lambda model: None, + ) + monkeypatch.setattr( + gguf_utils_module.gguf, + "GGUFReader", + lambda path: fake_reader, + ) + + config = maybe_patch_hf_config_from_gguf( + "model.gguf", + PretrainedConfig(model_type="qwen35"), + ) + + assert config.model_type == "qwen3_5" + assert config.architectures == ["Qwen3_5ForCausalLM"] + assert config.mtp_num_hidden_layers == 1 + assert config.num_nextn_predict_layers == 1 + assert config.num_hidden_layers == 64 + assert config.full_attention_interval == 4 + assert config.rope_parameters["mrope_section"] == [11, 11, 10] + assert config.rope_parameters["mrope_interleaved"] is True + + +def test_default_adapter_adds_mmproj_for_multimodal_config(tmp_path, monkeypatch): + main_path = tmp_path / "model.gguf" + mmproj_path = tmp_path / "mmproj-BF16.gguf" + config = PretrainedConfig(model_type="qwen3_5_moe") + config.vision_config = PretrainedConfig() + adapter = GGUFWeightsAdapter(config) + + monkeypatch.setattr( + default_adapter_module, + "detect_gguf_multimodal", + lambda model: mmproj_path, + ) + + assert adapter._get_weight_sources(str(main_path), config) == [ + str(main_path), + str(mmproj_path), + ] + + +def test_default_adapter_keeps_text_only_sources_without_mmproj(tmp_path, monkeypatch): + main_path = tmp_path / "model.gguf" + config = PretrainedConfig(model_type="qwen3_5_moe") + adapter = GGUFWeightsAdapter(config) + + monkeypatch.setattr( + default_adapter_module, + "detect_gguf_multimodal", + lambda model: tmp_path / "mmproj-BF16.gguf", + ) + + assert adapter._get_weight_sources(str(main_path), config) == [str(main_path)] + + +def test_default_adapter_ignores_mmproj_for_causal_lm_architecture( + tmp_path, monkeypatch +): + main_path = tmp_path / "model.gguf" + config = PretrainedConfig( + model_type="qwen3_5", + architectures=["Qwen3_5ForCausalLM"], + ) + config.vision_config = PretrainedConfig() + adapter = GGUFWeightsAdapter(config) + + monkeypatch.setattr( + default_adapter_module, + "detect_gguf_multimodal", + lambda model: tmp_path / "mmproj-BF16.gguf", + ) + + assert adapter._get_weight_sources(str(main_path), config) == [str(main_path)] + + +def _build_qwen3_5_test_name_map( + monkeypatch, + config, + state_names, + tensor_name_map=None, +): + tensor_name_map = tensor_name_map or {} + + class FakeNameMap: + def get_name(self, name): + return tensor_name_map.get(name) + + class FakeAutoModel: + @staticmethod + def from_config(config, trust_remote_code=False): + return SimpleNamespace( + state_dict=lambda: { + name: torch.empty((), device="meta") for name in state_names + }, + ) + + monkeypatch.setattr( + default_adapter_module.gguf, + "MODEL_ARCH_NAMES", + {object(): "qwen35", object(): "qwen35moe"}, + ) + monkeypatch.setattr( + default_adapter_module.gguf, + "get_tensor_name_map", + lambda *args, **kwargs: FakeNameMap(), + ) + monkeypatch.setattr( + default_adapter_module, + "AutoModelForImageTextToText", + FakeAutoModel, + ) + monkeypatch.setattr( + default_adapter_module, + "AutoModelForCausalLM", + FakeAutoModel, + ) + + model_config = SimpleNamespace(hf_config=config, trust_remote_code=False) + return GGUFWeightsAdapter(config).build_name_map(model_config) + + +def test_qwen3_5_dense_multimodal_maps_visual_merger(monkeypatch): + config = PretrainedConfig( + model_type="qwen3_5", + architectures=["Qwen3_5ForConditionalGeneration"], + num_hidden_layers=1, + layer_types=["linear_attention"], + ) + config.vision_config = PretrainedConfig(num_hidden_layers=1) + state_names = [ + "model.language_model.embed_tokens.weight", + "model.visual.patch_embed.proj.weight.1", + "model.visual.merger.linear_fc1.weight", + "model.visual.merger.linear_fc1.bias", + "model.visual.merger.linear_fc2.weight", + "model.visual.merger.linear_fc2.bias", + "model.visual.merger.norm.weight", + "model.visual.merger.norm.bias", + ] + + mapping = _build_qwen3_5_test_name_map(monkeypatch, config, state_names) + + assert mapping["token_embd.weight"] == ( + "model.language_model.embed_tokens.weight" + ) + assert mapping["v.patch_embd.weight.1"] == ( + "model.visual.patch_embed.proj.weight.1" + ) + assert mapping["v.post_ln.weight"] == "model.visual.merger.norm.weight" + assert mapping["v.post_ln.bias"] == "model.visual.merger.norm.bias" + assert mapping["mm.0.weight"] == "model.visual.merger.linear_fc1.weight" + assert mapping["mm.0.bias"] == "model.visual.merger.linear_fc1.bias" + assert mapping["mm.2.weight"] == "model.visual.merger.linear_fc2.weight" + assert mapping["mm.2.bias"] == "model.visual.merger.linear_fc2.bias" + + +def test_qwen3_5_moe_multimodal_maps_token_embd_to_language_model(monkeypatch): + config = PretrainedConfig( + model_type="qwen3_5_moe", + architectures=["Qwen3_5MoeForConditionalGeneration"], + num_hidden_layers=1, + layer_types=["linear_attention"], + ) + config.vision_config = PretrainedConfig(num_hidden_layers=1) + + mapping = _build_qwen3_5_test_name_map( + monkeypatch, + config, + ["model.language_model.embed_tokens.weight"], + ) + + assert mapping["token_embd.weight"] == ( + "model.language_model.embed_tokens.weight" + ) + + +def test_qwen3_5_text_only_does_not_add_visual_merger_mappings(monkeypatch): + config = PretrainedConfig( + model_type="qwen3_5", + num_hidden_layers=1, + layer_types=["linear_attention"], + ) + + mapping = _build_qwen3_5_test_name_map(monkeypatch, config, []) + + assert "v.patch_embd.weight.1" not in mapping + assert "v.post_ln.weight" not in mapping + assert "mm.0.weight" not in mapping + assert "mm.2.weight" not in mapping + + +def test_qwen3_5_causal_lm_uses_text_weight_layout(monkeypatch): + config = PretrainedConfig( + model_type="qwen3_5", + architectures=["Qwen3_5ForCausalLM"], + num_hidden_layers=1, + layer_types=["linear_attention"], + ) + config.vision_config = PretrainedConfig(num_hidden_layers=1) + + mapping = _build_qwen3_5_test_name_map( + monkeypatch, + config, + [ + "model.embed_tokens.weight", + "model.layers.0.linear_attn.dt_bias", + ], + { + "model.embed_tokens": "token_embd", + "model.layers.0.linear_attn.dt_bias": "blk.0.ssm_dt", + }, + ) + + assert mapping["token_embd.weight"] == "model.embed_tokens.weight" + assert mapping["blk.0.ssm_dt.bias"] == "model.layers.0.linear_attn.dt_bias" + assert "v.patch_embd.weight.1" not in mapping + assert "mm.0.weight" not in mapping + assert "mm.2.weight" not in mapping + + +def test_qwen3_5_adapter_reshapes_gguf_weights(): + adapter = Qwen3_5GGUFAdapter(PretrainedConfig(model_type="qwen3_5_moe")) + shared_gate = torch.arange(4) + conv1d = torch.arange(6).reshape(2, 3) + + assert adapter.transform_weight( + "model.layers.0.mlp.shared_expert_gate", + shared_gate, + ).shape == (1, 4) + transformed_conv = adapter.transform_weight( + "model.layers.0.linear_attn.conv1d.weight", + conv1d, + ) + assert transformed_conv.shape == (2, 1, 3) + assert torch.equal(transformed_conv[:, 0, :], conv1d) + + +def test_qwen3_5_adapter_combines_split_patch_embed_weight(): + adapter = Qwen3_5GGUFAdapter(PretrainedConfig(model_type="qwen3_5_moe")) + patch_weight = torch.zeros((4, 3, 16, 16)) + patch_weight_1 = torch.ones((4, 3, 16, 16)) + other_weight = torch.full((2, 2), 2.0) + + mapped = list( + adapter.map_weights( + [ + ("model.visual.patch_embed.proj.weight.1", patch_weight_1), + ("model.layers.0.self_attn.q_proj.weight", other_weight), + ("model.visual.patch_embed.proj.weight", patch_weight), + ] + ) + ) + + assert mapped[0][0] == "model.layers.0.self_attn.q_proj.weight" + assert torch.equal(mapped[0][1], other_weight) + assert mapped[1][0] == "model.visual.patch_embed.proj.weight" + assert mapped[1][1].shape == (4, 3, 2, 16, 16) + assert torch.equal(mapped[1][1][:, :, 0], patch_weight) + assert torch.equal(mapped[1][1][:, :, 1], patch_weight_1) + + +def test_qwen3_5_mtp_gguf_mappings(): + config = PretrainedConfig( + model_type="qwen3_5_moe", + num_hidden_layers=40, + mtp_num_hidden_layers=2, + ) + mapping: dict[str, str] = {} + sideload_params = [] + + _add_qwen3_5_mtp_gguf_mappings(config, mapping, sideload_params) + + assert mapping["blk.40.attn_q.weight"] == ("mtp.layers.0.self_attn.q_proj.weight") + assert mapping["blk.40.attn_k.weight"] == ("mtp.layers.0.self_attn.k_proj.weight") + assert mapping["blk.41.attn_q.weight"] == ("mtp.layers.1.self_attn.q_proj.weight") + assert mapping["blk.40.ffn_gate_inp.weight"] == "mtp.layers.0.mlp.gate.weight" + assert mapping["blk.40.ffn_gate_inp_shexp.weight"] == ( + "mtp.layers.0.mlp.shared_expert_gate.weight" + ) + assert mapping["blk.40.ffn_gate_exps.weight"] == ( + "mtp.layers.0.mlp.experts.0.gate_proj.weight" + ) + assert mapping["blk.40.nextn.eh_proj.weight"] == "mtp.fc.weight" + assert mapping["blk.40.nextn.shared_head_norm.weight"] == "mtp.norm.weight" + assert "blk.41.nextn.eh_proj.weight" not in mapping + assert sideload_params[0].fullmatch("mtp.layers.0.mlp.experts.15.gate_proj.weight") + + +def test_qwen3_5_dense_mtp_gguf_mappings_use_trunk_layer_count(): + config = PretrainedConfig( + model_type="qwen3_5", + num_hidden_layers=64, + mtp_num_hidden_layers=1, + ) + mapping: dict[str, str] = {} + sideload_params = [] + + _add_qwen3_5_mtp_gguf_mappings(config, mapping, sideload_params) + + assert mapping["blk.64.attn_q.weight"] == ("mtp.layers.0.self_attn.q_proj.weight") + assert mapping["blk.64.attn_k.weight"] == ("mtp.layers.0.self_attn.k_proj.weight") + assert mapping["blk.64.attn_v.weight"] == ("mtp.layers.0.self_attn.v_proj.weight") + assert mapping["blk.64.ffn_gate.weight"] == "mtp.layers.0.mlp.gate_proj.weight" + assert mapping["blk.64.ffn_up.weight"] == "mtp.layers.0.mlp.up_proj.weight" + assert mapping["blk.64.ffn_down.weight"] == "mtp.layers.0.mlp.down_proj.weight" + assert mapping["blk.64.nextn.eh_proj.weight"] == "mtp.fc.weight" + assert mapping["blk.64.nextn.shared_head_norm.weight"] == "mtp.norm.weight" + assert "blk.65.attn_q.weight" not in mapping + + +def _qwen3_5_linear_attn_test_config(): + return PretrainedConfig( + model_type="qwen3_5_moe", + linear_num_key_heads=2, + linear_num_value_heads=6, + linear_key_head_dim=2, + linear_value_head_dim=2, + ) + + +def _grouped_to_tiled_v_heads_for_test( + tensor: torch.Tensor, + dim: int, + *, + num_k_heads: int = 2, + num_v_per_k: int = 3, + head_dim: int = 2, +) -> torch.Tensor: + shape = list(tensor.shape) + if dim < 0: + dim += len(shape) + tensor = tensor.reshape( + *shape[:dim], num_k_heads, num_v_per_k, head_dim, *shape[dim + 1 :] + ) + perm = list(range(tensor.dim())) + perm[dim], perm[dim + 1] = perm[dim + 1], perm[dim] + return tensor.permute(*perm).contiguous().reshape(*shape) + + +def test_qwen3_5_adapter_restores_gdn_layout(): + torch.manual_seed(2) + adapter = Qwen3_5GGUFAdapter(_qwen3_5_linear_attn_test_config()) + key_dim = 4 + value_dim = 12 + hidden = 5 + + qkv = torch.randn(2 * key_dim + value_dim, hidden) + stored_qkv = torch.cat( + ( + qkv[: 2 * key_dim], + _grouped_to_tiled_v_heads_for_test(qkv[2 * key_dim :], 0), + ), + dim=0, + ) + assert torch.allclose( + adapter.transform_weight( + "model.layers.0.linear_attn.in_proj_qkv.weight", stored_qkv + ), + qkv, + ) + + z = torch.randn(value_dim, hidden) + stored_z = _grouped_to_tiled_v_heads_for_test(z, 0) + assert torch.allclose( + adapter.transform_weight( + "model.layers.0.linear_attn.in_proj_z.weight", stored_z + ), + z, + ) + + beta = torch.randn(6, hidden) + stored_beta = _grouped_to_tiled_v_heads_for_test(beta, 0, head_dim=1) + assert torch.allclose( + adapter.transform_weight( + "model.layers.0.linear_attn.in_proj_b.weight", stored_beta + ), + beta, + ) + + dt_bias = torch.randn(6) + stored_dt_bias = _grouped_to_tiled_v_heads_for_test( + dt_bias.unsqueeze(-1), 0, head_dim=1 + ).squeeze(-1) + assert torch.allclose( + adapter.transform_weight("model.layers.0.linear_attn.dt_bias", stored_dt_bias), + dt_bias, + ) + + a_log = torch.rand(6) + 0.1 + stored_a_log = _grouped_to_tiled_v_heads_for_test( + (-torch.exp(a_log)).unsqueeze(-1), 0, head_dim=1 + ).squeeze(-1) + assert torch.allclose( + adapter.transform_weight("model.layers.0.linear_attn.A_log", stored_a_log), + a_log, + atol=1e-6, + ) + + conv = torch.randn(2 * key_dim + value_dim, 3) + stored_conv = torch.cat( + ( + conv[: 2 * key_dim], + _grouped_to_tiled_v_heads_for_test(conv[2 * key_dim :], 0), + ), + dim=0, + ) + restored_conv = adapter.transform_weight( + "model.layers.0.linear_attn.conv1d.weight", stored_conv + ) + assert restored_conv.shape == (2 * key_dim + value_dim, 1, 3) + assert torch.allclose(restored_conv[:, 0, :], conv) + + out_proj = torch.randn(hidden, value_dim) + stored_out_proj = _grouped_to_tiled_v_heads_for_test(out_proj, 1) + assert torch.allclose( + adapter.transform_weight( + "model.layers.0.linear_attn.out_proj.weight", stored_out_proj + ), + out_proj, + ) + + +def test_qwen3_5_adapter_restores_text_norm_weights(): + adapter = Qwen3_5GGUFAdapter(_qwen3_5_linear_attn_test_config()) + weight = torch.tensor([1.25, 2.5, 0.75]) + + assert torch.allclose( + adapter.transform_weight("model.layers.0.input_layernorm.weight", weight), + weight - 1, + ) + assert torch.allclose( + adapter.transform_weight( + "model.language_model.layers.0.post_attention_layernorm.weight", + weight, + ), + weight - 1, + ) + assert torch.allclose( + adapter.transform_weight("model.norm.weight", weight), + weight - 1, + ) + + assert torch.equal( + adapter.transform_weight("model.layers.0.linear_attn.norm.weight", weight), + weight, + ) + assert torch.equal( + adapter.transform_weight("model.visual.merger.norm.weight", weight), + weight, + ) + + +def test_qwen3_5_adapter_forwards_quantized_gdn_weights(): + adapter = Qwen3_5GGUFAdapter(PretrainedConfig(model_type="qwen3_5_moe")) + qweight_type = torch.tensor(int(WeightType.Q5_K)) + qweight = torch.randn(2, 3) + + mapped = list( + adapter.map_weights( + [ + ("model.layers.0.linear_attn.in_proj_a.qweight_type", qweight_type), + ("model.layers.0.linear_attn.in_proj_a.qweight", qweight), + ] + ) + ) + + assert mapped == [ + ("model.layers.0.linear_attn.in_proj_a.qweight_type", qweight_type), + ("model.layers.0.linear_attn.in_proj_a.qweight", qweight), + ] + assert adapter._qweight_types["model.layers.0.linear_attn.in_proj_a"] == ( + WeightType.Q5_K + ) + + +def test_qwen3_5_adapter_dequantizes_forced_out_proj(monkeypatch): + adapter = Qwen3_5GGUFAdapter(_qwen3_5_linear_attn_test_config()) + module_name = "model.layers.0.linear_attn.out_proj" + adapter._forced_dequantized_modules.add(module_name) + dense = torch.randn(5, 12) + stored_dense = _grouped_to_tiled_v_heads_for_test(dense, 1) + + def fake_dequantize(weight, qweight_type): + assert qweight_type == WeightType.Q5_K + return stored_dense.numpy() + + monkeypatch.setattr( + qwen3_5_adapter_module.gguf.quants, "dequantize", fake_dequantize + ) + + mapped = list( + adapter.map_weights( + [ + (f"{module_name}.qweight_type", torch.tensor(int(WeightType.Q5_K))), + (f"{module_name}.qweight", torch.zeros(1)), + ] + ) + ) + + assert len(mapped) == 1 + assert mapped[0][0] == f"{module_name}.weight" + assert torch.allclose(mapped[0][1], dense) + + +def test_gguf_config_parser_uses_parent_dir_for_local_file(tmp_path, monkeypatch): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + calls = {} + + def fake_parse( + self, model, trust_remote_code, revision=None, code_revision=None, **kwargs + ): + calls["model"] = model + calls["trust_remote_code"] = trust_remote_code + calls["gguf_file"] = kwargs.get("gguf_file") + return {}, PretrainedConfig(model_type="qwen3_moe") + + monkeypatch.setattr( + gguf_config_parser_module.HFConfigParser, + "parse", + fake_parse, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "file_or_path_exists", + lambda *args, **kwargs: True, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "maybe_patch_hf_config_from_gguf", + lambda model, config: config, + ) + + config_dict, config = GGUFConfigParser().parse(gguf_path, trust_remote_code=False) + + assert calls["model"] == gguf_path.parent + assert calls["trust_remote_code"] is False + assert calls["gguf_file"] is None + assert config_dict["norm_topk_prob"] is True + assert config.architectures == ["Qwen3MoeForCausalLM"] + + +def test_gguf_config_parser_prefers_sidecar_config( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + calls = {} + + def fake_parse( + self, model, trust_remote_code, revision=None, code_revision=None, **kwargs + ): + calls["model"] = model + calls["gguf_file"] = kwargs.get("gguf_file") + return {}, PretrainedConfig( + model_type="qwen3_5_moe", + architectures=["Qwen3_5MoeForCausalLM"], + ) + + monkeypatch.setattr( + gguf_config_parser_module.HFConfigParser, + "parse", + fake_parse, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "maybe_patch_hf_config_from_gguf", + lambda model, config: config, + ) + + config_dict, config = GGUFConfigParser().parse(gguf_path, trust_remote_code=False) + + assert calls["model"] == gguf_path.parent + assert calls["gguf_file"] is None + assert config.model_type == "qwen3_5_moe" + assert config_dict["architectures"] == ["Qwen3_5MoeForCausalLM"] + + +def test_gguf_config_parser_uses_gguf_file_when_parent_has_no_config( + tmp_path, monkeypatch +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + calls = {} + + def fake_parse( + self, model, trust_remote_code, revision=None, code_revision=None, **kwargs + ): + calls["model"] = model + calls["trust_remote_code"] = trust_remote_code + calls["gguf_file"] = kwargs.get("gguf_file") + return {}, PretrainedConfig(model_type="qwen3_moe") + + monkeypatch.setattr( + gguf_config_parser_module.HFConfigParser, + "parse", + fake_parse, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "file_or_path_exists", + lambda *args, **kwargs: False, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "maybe_patch_hf_config_from_gguf", + lambda model, config: config, + ) + + GGUFConfigParser().parse(gguf_path, trust_remote_code=False) + + assert calls["model"] == gguf_path.parent + assert calls["trust_remote_code"] is False + assert calls["gguf_file"] == gguf_path.name + + +def test_gguf_config_parser_resolves_presplit_local_gguf( + tmp_path, + monkeypatch, +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + calls = {} + + def fake_parse( + self, model, trust_remote_code, revision=None, code_revision=None, **kwargs + ): + calls["model"] = model + calls["trust_remote_code"] = trust_remote_code + calls["revision"] = revision + calls["gguf_file"] = kwargs.get("gguf_file") + return {}, PretrainedConfig(model_type="qwen3_moe") + + monkeypatch.setattr( + gguf_config_parser_module, + "resolve_gguf_config_source", + lambda model, revision=None: "base/repo", + ) + monkeypatch.setattr( + gguf_config_parser_module.HFConfigParser, + "parse", + fake_parse, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "maybe_patch_hf_config_from_gguf", + lambda model, config: config, + ) + + GGUFConfigParser().parse( + gguf_path.parent, + trust_remote_code=True, + revision="gguf-revision", + gguf_file=gguf_path.name, + ) + + assert calls["model"] == "base/repo" + assert calls["trust_remote_code"] is False + assert calls["revision"] is None + assert calls["gguf_file"] is None + + +def test_gguf_config_parser_preserves_patched_architecture(tmp_path, monkeypatch): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + + def fake_parse( + self, model, trust_remote_code, revision=None, code_revision=None, **kwargs + ): + return {}, PretrainedConfig(model_type="qwen35moe") + + def fake_patch(model, config): + config.update( + { + "model_type": "qwen3_5_moe", + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + } + ) + return config + + monkeypatch.setattr( + gguf_config_parser_module.HFConfigParser, + "parse", + fake_parse, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "file_or_path_exists", + lambda *args, **kwargs: False, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "maybe_patch_hf_config_from_gguf", + fake_patch, + ) + + config_dict, config = GGUFConfigParser().parse(gguf_path, trust_remote_code=False) + + assert config.model_type == "qwen3_5_moe" + assert config.architectures == ["Qwen3_5MoeForConditionalGeneration"] + assert config_dict["architectures"] == ["Qwen3_5MoeForConditionalGeneration"] + + +def test_gguf_config_source_uses_nearest_parent_config(tmp_path): + model_dir = tmp_path / "model" + mtp_dir = model_dir / "MTP" + mtp_dir.mkdir(parents=True) + (model_dir / "config.json").write_text("{}", encoding="utf-8") + gguf_path = mtp_dir / "draft.gguf" + gguf_path.write_bytes(b"GGUF") + + assert resolve_gguf_config_source(gguf_path) == model_dir + + +def test_gguf_config_parser_disables_trust_for_base_model_redirect( + tmp_path, monkeypatch +): + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + calls = {} + + def fake_parse( + self, model, trust_remote_code, revision=None, code_revision=None, **kwargs + ): + calls["model"] = model + calls["trust_remote_code"] = trust_remote_code + calls["revision"] = revision + calls["gguf_file"] = kwargs.get("gguf_file") + return {}, PretrainedConfig(model_type="qwen3_moe") + + monkeypatch.setattr( + gguf_config_parser_module, + "resolve_gguf_config_source", + lambda model, revision=None: "base/repo", + ) + monkeypatch.setattr( + gguf_config_parser_module.HFConfigParser, + "parse", + fake_parse, + ) + monkeypatch.setattr( + gguf_config_parser_module, + "maybe_patch_hf_config_from_gguf", + lambda model, config: config, + ) + + GGUFConfigParser().parse( + gguf_path, + trust_remote_code=True, + revision="gguf-revision", + ) + + assert calls["model"] == "base/repo" + assert calls["trust_remote_code"] is False + assert calls["revision"] is None + assert calls["gguf_file"] is None + + +def test_register_sets_engine_args_for_gguf_model(monkeypatch): + register() + captured = {} + + def fake_model_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr(arg_utils_module, "ModelConfig", fake_model_config) + engine_args = EngineArgs(model="/tmp/model.gguf", tokenizer="/tmp/tokenizer") + + engine_args.create_model_config() + + assert captured["config_format"] == "gguf" + assert captured["model"] == "/tmp/model.gguf" + assert captured["model_weights"] == "/tmp/model.gguf" + assert captured["quantization"] == "gguf" + assert engine_args.load_format == "gguf" + + +def test_register_sets_embedded_tokenizer_for_local_gguf(tmp_path, monkeypatch): + register() + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + captured = {} + + monkeypatch.setattr( + gguf_plugin_module, + "build_tokenizer_from_gguf", + lambda model: "/tmp/gguf-tokenizer-cache", + ) + + def fake_model_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr(arg_utils_module, "ModelConfig", fake_model_config) + engine_args = EngineArgs(model=str(gguf_path)) + + engine_args.create_model_config() + + assert captured["tokenizer"] == "/tmp/gguf-tokenizer-cache" + + +def test_register_preserves_explicit_tokenizer_for_local_gguf(tmp_path, monkeypatch): + register() + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + captured = {} + + def fail_build_tokenizer(model): + raise AssertionError("explicit tokenizer must not be replaced") + + monkeypatch.setattr( + gguf_plugin_module, + "build_tokenizer_from_gguf", + fail_build_tokenizer, + ) + + def fake_model_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr(arg_utils_module, "ModelConfig", fake_model_config) + engine_args = EngineArgs(model=str(gguf_path), tokenizer="/tmp/tokenizer") + + engine_args.create_model_config() + + assert captured["tokenizer"] == "/tmp/tokenizer" + + +def test_register_skips_speculator_probe_for_gguf(): + register() + + model, tokenizer, speculative_config = ( + config_module.maybe_override_with_speculators( + model="/tmp/model.gguf", + tokenizer="/tmp/tokenizer", + trust_remote_code=False, + revision=None, + vllm_speculative_config={"foo": "bar"}, + hf_token=None, + ) + ) + + assert model == "/tmp/model.gguf" + assert tokenizer == "/tmp/tokenizer" + assert speculative_config == {"foo": "bar"} + + +def test_register_speculator_probe_prefers_sidecar_config( + tmp_path, + monkeypatch, +): + register() + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + calls = {} + + def fake_get_config_dict(config_source, **kwargs): + calls["config_source"] = config_source + calls["gguf_file"] = kwargs.get("gguf_file") + return {"model_type": "qwen3_5"}, {} + + monkeypatch.setattr( + gguf_plugin_module.PretrainedConfig, + "get_config_dict", + fake_get_config_dict, + ) + + model, tokenizer, speculative_config = ( + config_module.maybe_override_with_speculators( + model=str(gguf_path), + tokenizer="/tmp/tokenizer", + trust_remote_code=False, + revision=None, + vllm_speculative_config=None, + hf_token=None, + ) + ) + + assert calls["config_source"] == gguf_path.parent + assert calls["gguf_file"] is None + assert model == str(gguf_path) + assert tokenizer == "/tmp/tokenizer" + assert speculative_config is None + + +def test_register_disables_trust_for_gguf_config_redirect(tmp_path, monkeypatch): + register() + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + captured = {} + + monkeypatch.setattr( + gguf_plugin_module, + "resolve_gguf_config_source", + lambda model, revision=None: "base/repo", + ) + + def fake_model_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr( + arg_utils_module, + "ModelConfig", + fake_model_config, + ) + + engine_args = EngineArgs( + model=str(gguf_path), + tokenizer="/tmp/tokenizer", + trust_remote_code=True, + revision="gguf-revision", + ) + engine_args.create_model_config() + + assert captured["trust_remote_code"] is False + + +def test_register_keeps_trust_for_explicit_gguf_config_path(tmp_path, monkeypatch): + register() + gguf_path = tmp_path / "model.gguf" + gguf_path.write_bytes(b"GGUF") + config_path = tmp_path / "config-repo" + config_path.mkdir() + captured = {} + + monkeypatch.setattr( + gguf_plugin_module, + "resolve_gguf_config_source", + lambda model, revision=None: "base/repo", + ) + + def fake_model_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr( + arg_utils_module, + "ModelConfig", + fake_model_config, + ) + + engine_args = EngineArgs( + model=str(gguf_path), + tokenizer="/tmp/tokenizer", + trust_remote_code=True, + hf_config_path=str(config_path), + ) + engine_args.create_model_config() + + assert captured["trust_remote_code"] is True + + +def test_register_disables_trust_for_gguf_speculator_config(tmp_path, monkeypatch): + register() + gguf_path = tmp_path / "draft.gguf" + gguf_path.write_bytes(b"GGUF") + captured = {} + + monkeypatch.setattr( + gguf_plugin_module, + "resolve_gguf_config_source", + lambda model, revision=None: "base/repo", + ) + + def fake_model_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr( + arg_utils_module, + "ModelConfig", + fake_model_config, + ) + + engine_args = EngineArgs( + model="verifier/repo", + tokenizer="verifier/repo", + trust_remote_code=True, + speculative_config={"model": str(gguf_path)}, + ) + engine_args.create_model_config() + + assert captured["trust_remote_code"] is False + + +def test_gguf_qkv_shards_are_padded_in_qkv_order(monkeypatch): + register() + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + layer = QKVParallelLinear( + hidden_size=4, + head_size=2, + total_num_heads=2, + total_num_kv_heads=1, + bias=False, + quant_config=OOTGGUFConfig.from_config({}), + disable_tp=True, + ) + + q = torch.full((4, 4), 1, dtype=torch.uint8) + k = torch.full((2, 4), 2, dtype=torch.uint8) + v = torch.full((2, 4), 3, dtype=torch.uint8) + # Load out of canonical order to match GGUF tensor iteration order. + layer.weight_loader_v2(layer.qweight, k, "k") + layer.weight_loader_v2(layer.qweight, q, "q") + layer.weight_loader_v2(layer.qweight, v, "v") + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), "k") + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), "q") + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), "v") + + layer.quant_method.process_weights_after_loading(layer) + + assert layer.qweight.shard_id == ["q", "k", "v"] + assert layer.qweight.shard_offset_map == { + "q": (0, 4, 4), + "k": (4, 6, 4), + "v": (6, 8, 4), + } + assert layer.qweight.numel() == 0 + assert torch.equal(layer.qweight.data_container[0], q) + assert torch.equal(layer.qweight.data_container[1], k) + assert torch.equal(layer.qweight.data_container[2], v) + + +def test_gguf_linear_preserves_cuda_weight_device(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for device placement test") + + register() + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + with torch.device("cuda"): + layer = MergedColumnParallelLinear( + input_size=4, + output_sizes=[4, 4], + bias=False, + quant_config=OOTGGUFConfig.from_config({}), + disable_tp=True, + ) + + layer.weight_loader_v2(layer.qweight, torch.ones((4, 4), dtype=torch.uint8), 0) + layer.weight_loader_v2(layer.qweight, 2 * torch.ones((4, 4), dtype=torch.uint8), 1) + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), 0) + layer.weight_loader_v2(layer.qweight_type, torch.tensor(3, dtype=torch.uint8), 1) + layer.quant_method.process_weights_after_loading(layer) + + assert layer.qweight.device.type == "cuda" + assert [shard.device.type for shard in layer.qweight.data_container] == [ + "cuda", + "cuda", + ] + assert layer.qweight_type.device.type == "cuda" + + +def test_gguf_iq4_xs_batched_linear_uses_mmq_v2(monkeypatch): + import vllm_gguf_plugin.quantization.linear as gguf_linear_module + from vllm_gguf_plugin.quantization.linear import _fused_mul_mat_gguf + + qweight = torch.empty((32, 136), dtype=torch.uint8) + x = torch.empty((17, 256), dtype=torch.bfloat16) + expected = torch.empty((17, 32), dtype=torch.bfloat16) + calls = [] + + def fake_mmq_v2(qweight_arg, x_arg, row_arg): + calls.append((qweight_arg, x_arg, row_arg)) + return expected + + monkeypatch.setattr( + gguf_linear_module.ops, "ggml_mul_mat_a8_iq4_xs_mmq_v2", fake_mmq_v2 + ) + + output = _fused_mul_mat_gguf(x, qweight, WeightType.IQ4_XS) + + assert output is expected + assert calls == [(qweight, x, qweight.shape[0])] + + +def test_gguf_iq4_xs_single_token_linear_keeps_mmvq(monkeypatch): + import vllm_gguf_plugin.quantization.linear as gguf_linear_module + from vllm_gguf_plugin.quantization.linear import _fused_mul_mat_gguf + + qweight = torch.empty((32, 136), dtype=torch.uint8) + x = torch.empty((1, 256), dtype=torch.bfloat16) + expected = torch.empty((1, 32), dtype=torch.bfloat16) + calls = [] + + def fake_mmvq(qweight_arg, x_arg, qweight_type_arg, row_arg): + calls.append((qweight_arg, x_arg, qweight_type_arg, row_arg)) + return expected + + def fail_mmq_v2(*args, **kwargs): + raise AssertionError("IQ4_XS batch-size-1 path must keep MMVQ") + + monkeypatch.setattr(gguf_linear_module.ops, "ggml_mul_mat_vec_a8", fake_mmvq) + monkeypatch.setattr( + gguf_linear_module.ops, "ggml_mul_mat_a8_iq4_xs_mmq_v2", fail_mmq_v2 + ) + + output = _fused_mul_mat_gguf(x, qweight, WeightType.IQ4_XS) + + assert output is expected + assert calls == [(qweight, x, WeightType.IQ4_XS, qweight.shape[0])] + + +def test_gguf_iq4_xs_batched_moe_uses_mmq_v2(monkeypatch): + import vllm.model_executor.layers.fused_moe.fused_moe as fused_moe_module + + import vllm_gguf_plugin.quantization.fused_moe as gguf_moe_module + from vllm_gguf_plugin.quantization.fused_moe import _fused_moe_gguf + + def fake_align(topk_ids, block_size, num_experts): + del num_experts + num_ids = topk_ids.numel() + padded = ((num_ids + block_size - 1) // block_size) * block_size + sorted_token_ids = torch.arange(padded, dtype=torch.int32) + sorted_token_ids[num_ids:] = -1 + expert_ids = torch.zeros(padded // block_size, dtype=torch.int32) + num_tokens_post_padded = torch.tensor([padded], dtype=torch.int32) + return sorted_token_ids, expert_ids, num_tokens_post_padded + + def fake_apply_moe_activation(activation, output, input_): + del activation + output.copy_(input_[..., : output.shape[-1]]) + + calls = [] + + def fake_moe_v2( + x, + weight, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + row, + top_k, + tokens, + ): + calls.append((x.shape, weight.shape, row, top_k, tokens)) + assert sorted_token_ids.dtype == torch.int32 + assert expert_ids.dtype == torch.int32 + assert num_tokens_post_padded.dtype == torch.int32 + return torch.ones((tokens * top_k, row), dtype=x.dtype) + + def fake_moe_sum(input_, output): + output.copy_(input_.sum(dim=1)) + + monkeypatch.setattr(fused_moe_module, "moe_align_block_size", fake_align) + monkeypatch.setattr( + gguf_moe_module, "apply_moe_activation", fake_apply_moe_activation + ) + monkeypatch.setattr(gguf_moe_module.ops, "ggml_moe_a8_iq4_xs_mmq_v2", fake_moe_v2) + monkeypatch.setattr(gguf_moe_module.ops, "moe_sum", fake_moe_sum) + + x = torch.ones((65, 4), dtype=torch.float32) + w1 = torch.empty((4, 8, 4), dtype=torch.uint8) + w2 = torch.empty((4, 4, 4), dtype=torch.uint8) + topk_weights = torch.ones((65, 2), dtype=torch.float32) + topk_ids = torch.zeros((65, 2), dtype=torch.int32) + + output = _fused_moe_gguf( + x, + w1, + w2, + topk_weights, + topk_ids, + WeightType.IQ4_XS, + WeightType.IQ4_XS, + "silu", + ) + + assert output.shape == x.shape + assert calls == [ + (torch.Size([65, 4]), torch.Size([4, 8, 4]), 8, 2, 65), + (torch.Size([130, 4]), torch.Size([4, 4, 4]), 4, 1, 130), + ] + + +def _make_iq4_xs_weight( + num_rows: int, + seed: int, + device: torch.device, +) -> torch.Tensor: + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + weight = torch.randint(0, 256, (num_rows, 136), dtype=torch.uint8, generator=gen) + # block_iq4_xs layout: + # half d, uint16 scales_h, uint8 scales_l[4], uint8 qs[128]. + # Use d=1 and 6-bit sub-block scale 33, which decodes to a small positive + # scale while still exercising the IQ4_XS scale unpacking path. + weight[:, 0:2] = torch.tensor([0x00, 0x3C], dtype=torch.uint8) + weight[:, 2:4] = torch.tensor([0xAA, 0xAA], dtype=torch.uint8) + weight[:, 4:8] = 0x11 + return weight.to(device=device) + + +def test_gguf_iq4_xs_batched_moe_matches_slow_reference_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for GGUF MoE kernel comparison") + if not ( + hasattr(torch.ops, "_C_gguf") + and hasattr(torch.ops._C_gguf, "ggml_moe_a8_iq4_xs_mmq_v2") + and hasattr(torch.ops._C_gguf, "ggml_mul_mat_vec_a8") + ): + pytest.skip("GGUF CUDA extension with IQ4_XS MoE v2 is not available") + + from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, + apply_moe_activation, + ) + + from vllm_gguf_plugin.quantization.fused_moe import _fused_moe_gguf + from vllm_gguf_plugin.quantization.linear import _fused_mul_mat_gguf + + device = torch.device("cuda") + dtype = torch.float16 + num_tokens = 65 + hidden_size = 256 + intermediate_size = 8 + num_experts = 4 + top_k = 2 + torch.manual_seed(0) + + x = torch.randn((num_tokens, hidden_size), dtype=dtype, device=device) * 0.05 + w1 = torch.stack( + [ + _make_iq4_xs_weight(intermediate_size * 2, 100 + expert, device) + for expert in range(num_experts) + ] + ) + w2 = torch.stack( + [ + _make_iq4_xs_weight(hidden_size, 200 + expert, device) + for expert in range(num_experts) + ] + ) + topk_ids = torch.tensor( + [ + [token % num_experts, (token + 1) % num_experts] + for token in range(num_tokens) + ], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor([0.65, 0.35], dtype=dtype, device=device).repeat( + num_tokens, + 1, + ) + + out = _fused_moe_gguf( + x, + w1, + w2, + topk_weights, + topk_ids, + WeightType.IQ4_XS, + WeightType.IQ4_XS, + "silu", + ) + + activation_enum = MoEActivation.from_str("silu") + ref = torch.empty_like(out) + for token_idx in range(num_tokens): + token_out = None + token_x = x[token_idx : token_idx + 1] + for route_idx in range(top_k): + expert_idx = int(topk_ids[token_idx, route_idx].item()) + hidden = _fused_mul_mat_gguf( + token_x, + w1[expert_idx], + WeightType.IQ4_XS, + ) + activated = torch.empty( + (1, intermediate_size), + dtype=hidden.dtype, + device=device, + ) + apply_moe_activation(activation_enum, activated, hidden) + projected = _fused_mul_mat_gguf( + activated, + w2[expert_idx], + WeightType.IQ4_XS, + ).mul(topk_weights[token_idx, route_idx]) + token_out = projected if token_out is None else token_out + projected + ref[token_idx] = token_out + + torch.testing.assert_close(out, ref, atol=5e-1, rtol=5e-2) diff --git a/vllm_gguf_plugin/config_parser.py b/vllm_gguf_plugin/config_parser.py index 9fa67cb7..57378481 100644 --- a/vllm_gguf_plugin/config_parser.py +++ b/vllm_gguf_plugin/config_parser.py @@ -4,17 +4,24 @@ from transformers import PretrainedConfig from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES +from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME +from vllm.logger import init_logger from vllm.transformers_utils.config import HFConfigParser from vllm.transformers_utils.config_parser_base import ConfigParserBase +from vllm.transformers_utils.repo_utils import file_or_path_exists from .gguf_utils import ( check_gguf_file, + get_gguf_file_path_from_hf, is_gguf, is_remote_gguf, maybe_patch_hf_config_from_gguf, + resolve_gguf_config_source, split_remote_gguf, ) +logger = init_logger(__name__) + class GGUFConfigParser(ConfigParserBase): def parse( @@ -26,7 +33,61 @@ def parse( **kwargs, ) -> tuple[dict, PretrainedConfig]: original_model = model - resolved_model = self._resolve_config_source(model) + gguf_path = None + if (gguf_file := kwargs.pop("gguf_file", None)) is not None: + candidate = Path(model) / gguf_file + if check_gguf_file(candidate): + original_model = candidate + gguf_path = candidate + + resolved_model = self._resolve_config_source(model, revision=revision) + + if gguf_path is not None or check_gguf_file(model): + gguf_path = gguf_path or Path(model) + resolved_model = self._resolve_config_source( + gguf_path, + revision=revision, + ) + gguf_repo = gguf_path.parent + if resolved_model != gguf_repo: + logger.warning_once( + "Disabling `trust_remote_code` because GGUF metadata " + "redirected config loading from %s to %s. Pass an " + "explicit `--hf-config-path` to opt in.", + gguf_repo, + resolved_model, + ) + trust_remote_code = False + revision = None + elif not file_or_path_exists( + gguf_repo, + HF_CONFIG_NAME, + revision=revision, + ): + kwargs["gguf_file"] = gguf_path.name + elif is_remote_gguf(model): + repo_id, quant_type = split_remote_gguf(model) + if resolved_model != repo_id: + logger.warning_once( + "Disabling `trust_remote_code` because GGUF metadata " + "redirected config loading from %s to %s. Pass an " + "explicit `--hf-config-path` to opt in.", + repo_id, + resolved_model, + ) + trust_remote_code = False + revision = None + elif not file_or_path_exists( + repo_id, + HF_CONFIG_NAME, + revision=revision, + ): + kwargs["gguf_file"] = get_gguf_file_path_from_hf( + repo_id, + quant_type, + revision=revision, + ) + config_dict, config = HFConfigParser().parse( resolved_model, trust_remote_code=trust_remote_code, @@ -39,6 +100,18 @@ def parse( config_dict["norm_topk_prob"] = True config.update({"norm_topk_prob": True}) + if is_gguf(original_model): + config = maybe_patch_hf_config_from_gguf(str(original_model), config) + + if config.model_type in ("gemma4_assistant", "gemma4_mtp"): + config_dict["architectures"] = ["Gemma4MTPModel"] + config.update({"architectures": ["Gemma4MTPModel"]}) + return config_dict, config + + if config.architectures: + config_dict["architectures"] = list(config.architectures) + return config_dict, config + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: raise RuntimeError(f"Can't get gguf config for {config.model_type}.") @@ -46,16 +119,11 @@ def parse( config_dict["architectures"] = [model_type] config.update({"architectures": [model_type]}) - if is_gguf(original_model): - config = maybe_patch_hf_config_from_gguf(str(original_model), config) - return config_dict, config @staticmethod - def _resolve_config_source(model: str | Path) -> str | Path: - if check_gguf_file(model): - return Path(model).parent - if is_remote_gguf(model): - repo_id, _ = split_remote_gguf(model) - return repo_id - return model + def _resolve_config_source( + model: str | Path, + revision: str | None = None, + ) -> str | Path: + return resolve_gguf_config_source(model, revision=revision) diff --git a/vllm_gguf_plugin/csrc/gguf/gguf_kernel.cu b/vllm_gguf_plugin/csrc/gguf/gguf_kernel.cu index faee237b..6e7807d3 100644 --- a/vllm_gguf_plugin/csrc/gguf/gguf_kernel.cu +++ b/vllm_gguf_plugin/csrc/gguf/gguf_kernel.cu @@ -14,6 +14,7 @@ #include "dequantize.cuh" #include "mmvq.cuh" #include "mmq.cuh" +#include "mmq_v2.cuh" #include "moe.cuh" #include "moe_vec.cuh" @@ -289,6 +290,104 @@ Tensor ggml_mul_mat_a8(Tensor W, // quant weight return Y; } +Tensor ggml_mul_mat_a8_q4_0_mmq_v2(Tensor W, // quant weight + Tensor X, // input + int64_t row) { + const int col = X.sizes()[1]; + const int batch = X.sizes()[0]; + const int padded = (col + 512 - 1) / 512 * 512; + const int batch_padded = (batch + VLLM_GGUF_MMQ_X_Q4_0 - 1) / + VLLM_GGUF_MMQ_X_Q4_0 * VLLM_GGUF_MMQ_X_Q4_0; + constexpr int q8_1_mmq_ints = sizeof(block_q8_1_mmq_v2) / sizeof(int); + const int quant_blocks = padded / (4 * QK8_1); + const int32_t device_idx = X.get_device_index(); + const DeviceGuard device_guard(device_idx); + Tensor Y = torch::stable::new_empty(W, {batch, row}, X.scalar_type()); + Tensor quant_X = torch::stable::new_empty( + W, {quant_blocks * batch_padded * q8_1_mmq_ints}, ScalarType::Int); + cudaStream_t stream = get_current_cuda_stream(device_idx); + + VLLM_DISPATCH_FLOATING_TYPES( + X.scalar_type(), "ggml_mul_mat_a8_q4_0_mmq_v2", [&] { + vllm_gguf_quantize_mmq_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, padded, + batch, batch_padded, stream); + vllm_gguf_mul_mat_q4_0_q8_1_mmq_v2_cuda( + (void*)W.data_ptr(), (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), col, row, batch, batch_padded, row, + stream); + }); + return Y; +} + +Tensor ggml_mul_mat_a8_iq4_xs_mmq_v2(Tensor W, // quant weight + Tensor X, // input + int64_t row) { + const int col = X.sizes()[1]; + const int batch = X.sizes()[0]; + const int padded = (col + 512 - 1) / 512 * 512; + const int batch_padded = (batch + VLLM_GGUF_MMQ_X_IQ4_XS - 1) / + VLLM_GGUF_MMQ_X_IQ4_XS * VLLM_GGUF_MMQ_X_IQ4_XS; + constexpr int q8_1_mmq_ints = sizeof(block_q8_1_mmq_v2) / sizeof(int); + const int quant_blocks = padded / (4 * QK8_1); + const int32_t device_idx = X.get_device_index(); + const DeviceGuard device_guard(device_idx); + Tensor Y = torch::stable::new_empty(W, {batch, row}, X.scalar_type()); + Tensor quant_X = torch::stable::new_empty( + W, {quant_blocks * batch_padded * q8_1_mmq_ints}, ScalarType::Int); + cudaStream_t stream = get_current_cuda_stream(device_idx); + + VLLM_DISPATCH_FLOATING_TYPES( + X.scalar_type(), "ggml_mul_mat_a8_iq4_xs_mmq_v2", [&] { + vllm_gguf_quantize_mmq_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, padded, + batch, batch_padded, stream); + vllm_gguf_mul_mat_iq4_xs_q8_1_mmq_v2_cuda( + (void*)W.data_ptr(), (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), col, row, batch, batch_padded, row, + stream); + }); + return Y; +} + +Tensor ggml_moe_a8_iq4_xs_mmq_v2( + Tensor X, // input + Tensor W, // expert weights + Tensor sorted_token_ids, Tensor expert_ids, Tensor num_tokens_post_padded, + int64_t row, int64_t top_k, int64_t tokens) { + const int col = X.sizes()[1]; + const int batch = X.sizes()[0]; + const int padded = (col + 512 - 1) / 512 * 512; + const int sorted_tokens_padded = sorted_token_ids.sizes()[0]; + constexpr int q8_1_mmq_ints = sizeof(block_q8_1_mmq_v2) / sizeof(int); + const int quant_blocks = padded / (4 * QK8_1); + const int32_t device_idx = X.get_device_index(); + const DeviceGuard device_guard(device_idx); + Tensor Y = torch::stable::new_empty(W, {tokens * top_k, row}, + X.scalar_type()); + Tensor quant_X = torch::stable::new_empty( + W, {quant_blocks * sorted_tokens_padded * q8_1_mmq_ints}, + ScalarType::Int); + cudaStream_t stream = get_current_cuda_stream(device_idx); + + VLLM_DISPATCH_FLOATING_TYPES( + X.scalar_type(), "ggml_moe_a8_iq4_xs_mmq_v2", [&] { + vllm_gguf_quantize_moe_mmq_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), + (int*)sorted_token_ids.data_ptr(), + (int*)num_tokens_post_padded.data_ptr(), col, padded, batch, top_k, + sorted_tokens_padded, stream); + vllm_gguf_moe_iq4_xs_q8_1_mmq_v2_cuda( + (void*)W.data_ptr(), (void*)quant_X.data_ptr(), + (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), + (int*)expert_ids.data_ptr(), + (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, + batch, sorted_tokens_padded, row, top_k, sorted_tokens_padded, + stream); + }); + return Y; +} + Tensor ggml_moe_a8(Tensor X, // input Tensor W, // expert weights Tensor sorted_token_ids, Tensor expert_ids, diff --git a/vllm_gguf_plugin/csrc/gguf/mma_v2.cuh b/vllm_gguf_plugin/csrc/gguf/mma_v2.cuh new file mode 100644 index 00000000..5d9f7819 --- /dev/null +++ b/vllm_gguf_plugin/csrc/gguf/mma_v2.cuh @@ -0,0 +1,130 @@ +#pragma once + +#include + +#include + +#ifndef GGML_CUDA_CC_TURING + #define GGML_CUDA_CC_TURING 750 +#endif + +#ifndef GGML_CUDA_CC_AMPERE + #define GGML_CUDA_CC_AMPERE 800 +#endif + +#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && \ + __CUDA_ARCH__ >= GGML_CUDA_CC_TURING + #define VLLM_GGUF_TURING_MMA_AVAILABLE +#endif + +static __device__ __forceinline__ void vllm_gguf_no_device_code() { +#ifdef __CUDA_ARCH__ + asm("trap;"); +#endif +} + +namespace vllm_gguf_mma { + +enum data_layout { + DATA_LAYOUT_I_MAJOR = 0, +}; + +template +struct tile {}; + +template +struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + static constexpr int ne = I * J / 32; + + T x[ne] = {0}; + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 8) + (threadIdx.x / 4); + } else { + vllm_gguf_no_device_code(); + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 8) { + return ((threadIdx.x % 4) * 2) + (l % 2); + } else { + vllm_gguf_no_device_code(); + return -1; + } + } +}; + +template +static __device__ __forceinline__ void load_generic(tile& t, + const T* __restrict__ xs0, + const int stride) { +#pragma unroll + for (int l = 0; l < t.ne; ++l) { + t.x[l] = xs0[t.get_i(l) * stride + t.get_j(l)]; + } +} + +template +static __device__ __forceinline__ void load_ldmatrix(tile<16, 8, T, dl>& t, + const T* __restrict__ xs0, + const int stride) { +#ifdef VLLM_GGUF_TURING_MMA_AVAILABLE + int* xi = (int*)t.x; + const int* xs = (const int*)xs0 + (threadIdx.x % t.I) * stride + + (threadIdx.x / t.I) * (t.J / 2); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) + : "l"(xs)); +#else + load_generic(t, xs0, stride); +#endif +} + +static __device__ __forceinline__ void mma(tile<16, 8, int>& D, + const tile<16, 8, int>& A, + const tile<8, 8, int>& B) { +#ifdef VLLM_GGUF_TURING_MMA_AVAILABLE + #if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[0]), "r"(A.x[1]), "r"(A.x[2]), "r"(A.x[3]), "r"(B.x[0]), + "r"(B.x[1])); + #else + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 " + "{%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[0]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 " + "{%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[1]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 " + "{%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[2]), "r"(B.x[1])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 " + "{%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[3]), "r"(B.x[1])); + #endif +#else + (void)D; + (void)A; + (void)B; + vllm_gguf_no_device_code(); +#endif +} + +} // namespace vllm_gguf_mma diff --git a/vllm_gguf_plugin/csrc/gguf/mmq_v2.cuh b/vllm_gguf_plugin/csrc/gguf/mmq_v2.cuh new file mode 100644 index 00000000..0ea155ac --- /dev/null +++ b/vllm_gguf_plugin/csrc/gguf/mmq_v2.cuh @@ -0,0 +1,870 @@ +#pragma once + +#include +#include + +#include "mma_v2.cuh" + +#define VLLM_GGUF_MMQ_ITER_K 256 +#define VLLM_GGUF_MMQ_NWARPS 8 +#define VLLM_GGUF_MMQ_X_Q4_0 64 +#define VLLM_GGUF_MMQ_X_IQ4_XS 64 +#define VLLM_GGUF_MOE_MMQ_X_IQ4_XS 8 +#define VLLM_GGUF_MMQ_Y 128 +#define VLLM_GGUF_MMQ_TILE_NE_K 32 +#define VLLM_GGUF_MMQ_TILE_Y_K \ + (VLLM_GGUF_MMQ_TILE_NE_K + VLLM_GGUF_MMQ_TILE_NE_K / QI8_1) +#define VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 \ + (2 * VLLM_GGUF_MMQ_TILE_NE_K + 2 * VLLM_GGUF_MMQ_TILE_NE_K / QI8_0 + 4) +#define VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128 + +// IQ4_XS keeps the same on-disk block_iq4_xs layout as the existing vLLM +// MMVQ/dequant paths, but modern MMQ uses the llama.cpp b4267+ QR/QI tiling. +// Keep these constants local to v2 so the legacy IQ4_XS paths stay unchanged. +#define VLLM_GGUF_MMQ_QR4_XS 2 +#define VLLM_GGUF_MMQ_QI4_XS (QK_K / (4 * VLLM_GGUF_MMQ_QR4_XS)) + +struct block_q8_1_mmq_v2 { + half2 ds4[4]; + int8_t qs[4 * QK8_1]; +}; + +static_assert(sizeof(block_q8_1_mmq_v2) == 4 * sizeof(block_q8_1), + "Unexpected block_q8_1_mmq_v2 size"); + +static constexpr __host__ __device__ int vllm_gguf_mmq_pad(const int x, + const int n) { + return ((x + n - 1) / n) * n; +} + +static constexpr __device__ int vllm_gguf_mmq_get_granularity_device( + const int mmq_x) { +#ifdef VLLM_GGUF_TURING_MMA_AVAILABLE + return mmq_x >= 48 ? 16 : 8; +#else + return 8; +#endif +} + +template +static __global__ void vllm_gguf_quantize_mmq_q8_1( + const scalar_t* __restrict__ x, void* __restrict__ vy, const int kx, + const int kx_padded, const int batch, const int batch_padded) { + const int i0 = (blockDim.x * blockIdx.y + threadIdx.x) * 4; // K offset + if (i0 >= kx_padded) { + return; + } + + const int i1 = blockIdx.x; // token row + + float xi0 = 0.0f; + float xi1 = 0.0f; + float xi2 = 0.0f; + float xi3 = 0.0f; + if (i1 < batch) { + if (i0 + 0 < kx) { + xi0 = static_cast(x[i1 * kx + i0 + 0]); + } + if (i0 + 1 < kx) { + xi1 = static_cast(x[i1 * kx + i0 + 1]); + } + if (i0 + 2 < kx) { + xi2 = static_cast(x[i1 * kx + i0 + 2]); + } + if (i0 + 3 < kx) { + xi3 = static_cast(x[i1 * kx + i0 + 3]); + } + } + + float amax = fabsf(xi0); + amax = fmaxf(amax, fabsf(xi1)); + amax = fmaxf(amax, fabsf(xi2)); + amax = fmaxf(amax, fabsf(xi3)); + float sum = xi0 + xi1 + xi2 + xi3; + +#pragma unroll + for (int offset = 4; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset)); + sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset); + } + + const float d = amax == 0.0f ? 0.0f : amax / 127.0f; + const float d_inv = amax == 0.0f ? 0.0f : 127.0f / amax; + + char4 q; + q.x = static_cast(roundf(xi0 * d_inv)); + q.y = static_cast(roundf(xi1 * d_inv)); + q.z = static_cast(roundf(xi2 * d_inv)); + q.w = static_cast(roundf(xi3 * d_inv)); + + block_q8_1_mmq_v2* y = (block_q8_1_mmq_v2*)vy; + const int ib = (i0 / (4 * QK8_1)) * batch_padded + i1; + const int iqs = i0 % (4 * QK8_1); + + char4* yqs4 = (char4*)y[ib].qs; + yqs4[iqs / 4] = q; + + if (iqs % 32 == 0) { + y[ib].ds4[iqs / 32] = make_half2(__float2half(d), __float2half(sum)); + } +} + +template +static void vllm_gguf_quantize_mmq_q8_1_cuda(const scalar_t* x, void* vy, + const int kx, const int kx_padded, + const int batch, + const int batch_padded, + cudaStream_t stream) { + const int block_num_y = + (kx_padded + 4 * VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / + (4 * VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(batch_padded, block_num_y, 1); + const dim3 block_size(VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + vllm_gguf_quantize_mmq_q8_1<<>>( + x, vy, kx, kx_padded, batch, batch_padded); +} + +template +static __global__ void vllm_gguf_quantize_moe_mmq_q8_1( + const scalar_t* __restrict__ x, void* __restrict__ vy, + const int* __restrict__ sorted_token_ids, + const int* __restrict__ num_tokens_post_padded, const int kx, + const int kx_padded, const int batch, const int top_k, + const int sorted_tokens_padded) { + const int i0 = (blockDim.x * blockIdx.y + threadIdx.x) * 4; // K offset + if (i0 >= kx_padded) { + return; + } + + const int sorted_row = blockIdx.x; + const bool valid_sorted_row = sorted_row < num_tokens_post_padded[0]; + const int token_off = valid_sorted_row ? sorted_token_ids[sorted_row] : -1; + const int token = token_off >= 0 ? token_off / top_k : -1; + + float xi0 = 0.0f; + float xi1 = 0.0f; + float xi2 = 0.0f; + float xi3 = 0.0f; + if (token >= 0 && token < batch) { + if (i0 + 0 < kx) { + xi0 = static_cast(x[token * kx + i0 + 0]); + } + if (i0 + 1 < kx) { + xi1 = static_cast(x[token * kx + i0 + 1]); + } + if (i0 + 2 < kx) { + xi2 = static_cast(x[token * kx + i0 + 2]); + } + if (i0 + 3 < kx) { + xi3 = static_cast(x[token * kx + i0 + 3]); + } + } + + float amax = fabsf(xi0); + amax = fmaxf(amax, fabsf(xi1)); + amax = fmaxf(amax, fabsf(xi2)); + amax = fmaxf(amax, fabsf(xi3)); + float sum = xi0 + xi1 + xi2 + xi3; + +#pragma unroll + for (int offset = 4; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset)); + sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset); + } + + const float d = amax == 0.0f ? 0.0f : amax / 127.0f; + const float d_inv = amax == 0.0f ? 0.0f : 127.0f / amax; + + char4 q; + q.x = static_cast(roundf(xi0 * d_inv)); + q.y = static_cast(roundf(xi1 * d_inv)); + q.z = static_cast(roundf(xi2 * d_inv)); + q.w = static_cast(roundf(xi3 * d_inv)); + + block_q8_1_mmq_v2* y = (block_q8_1_mmq_v2*)vy; + const int ib = (i0 / (4 * QK8_1)) * sorted_tokens_padded + sorted_row; + const int iqs = i0 % (4 * QK8_1); + + char4* yqs4 = (char4*)y[ib].qs; + yqs4[iqs / 4] = q; + + if (iqs % 32 == 0) { + y[ib].ds4[iqs / 32] = make_half2(__float2half(d), __float2half(sum)); + } +} + +template +static void vllm_gguf_quantize_moe_mmq_q8_1_cuda( + const scalar_t* x, void* vy, const int* sorted_token_ids, + const int* num_tokens_post_padded, const int kx, const int kx_padded, + const int batch, const int top_k, const int sorted_tokens_padded, + cudaStream_t stream) { + const int block_num_y = + (kx_padded + 4 * VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / + (4 * VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(sorted_tokens_padded, block_num_y, 1); + const dim3 block_size(VLLM_GGUF_CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + vllm_gguf_quantize_moe_mmq_q8_1<<>>( + x, vy, sorted_token_ids, num_tokens_post_padded, kx, kx_padded, batch, + top_k, sorted_tokens_padded); +} + +template +static __device__ __forceinline__ void vllm_gguf_load_tiles_q4_0_v2( + const char* __restrict__ x, int* __restrict__ x_tile, const int kbx0, + const int i_max, const int stride) { + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int warp_size = WARP_SIZE_GGUF; + + int* x_qs = (int*)x_tile; + float* x_df = (float*)(x_qs + 2 * VLLM_GGUF_MMQ_TILE_NE_K); + + constexpr int threads_per_row = VLLM_GGUF_MMQ_ITER_K / (4 * QR4_0); + constexpr int nrows = warp_size / threads_per_row; + const int txi = + warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_0; + const int kqsx = txi % QI4_0; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows * nwarps) { + int i = + i0 + (nrows == 1 ? threadIdx.y + : threadIdx.y * nrows + threadIdx.x / threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_0* bxi = (const block_q4_0*)x + kbx0 + i * stride + kbx; + const int qs0 = get_int_b2(bxi->qs, kqsx); + + x_qs[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + kbx * (2 * QI4_0) + kqsx + 0] = + __vsubss4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808); + x_qs[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + kbx * (2 * QI4_0) + kqsx + + QI4_0] = __vsubss4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808); + } + + constexpr int blocks_per_tile_x_row = VLLM_GGUF_MMQ_TILE_NE_K / QI4_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = + i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_0* bxi = (const block_q4_0*)x + kbx0 + i * stride + kbxd; + x_df[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = + __half2float(bxi->d); + } +} + +static __device__ __forceinline__ int2 +vllm_gguf_get_int_from_table_16_v2(const int q4, const uint8_t* values) { + int v1; + int v2; + get_int_from_table_16(static_cast(q4), values, v1, v2); + return make_int2(v1, v2); +} + +template +static __device__ __forceinline__ void vllm_gguf_load_tiles_iq4_xs_v2( + const char* __restrict__ x, int* __restrict__ x_tile, const int kbx0, + const int i_max, const int stride) { + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int warp_size = WARP_SIZE_GGUF; + + int* x_qs = (int*)x_tile; + float* x_df = (float*)(x_qs + 2 * VLLM_GGUF_MMQ_TILE_NE_K); + + constexpr int threads_per_row = + VLLM_GGUF_MMQ_ITER_K / (4 * VLLM_GGUF_MMQ_QR4_XS); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + const uint8_t* values = (const uint8_t*)kvalues_iq4nl; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows * nwarps) { + int i = + i0 + (nrows == 1 ? threadIdx.y + : threadIdx.y * nrows + threadIdx.x / threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_iq4_xs* bxi = (const block_iq4_xs*)x + kbx0 + i * stride; + const int aux_q4 = get_int_b4(bxi->qs, kqsx); + const int2 v = vllm_gguf_get_int_from_table_16_v2(aux_q4, values); + const int k0 = 8 * (kqsx / 4) + kqsx % 4; + + x_qs[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x; + x_qs[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + k0 + 4] = v.y; + } + + constexpr int scale_groups = VLLM_GGUF_MMQ_TILE_NE_K / 4; + constexpr int rows_per_warp = warp_size / scale_groups; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / scale_groups; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq4_xs* bxi = (const block_iq4_xs*)x + kbx0 + i * stride; + const int scale_idx = threadIdx.x % scale_groups; + const int ls = + ((bxi->scales_l[scale_idx / 2] >> (4 * (scale_idx % 2))) & 0x0F) | + (((bxi->scales_h >> (2 * scale_idx)) & 0x03) << 4); + + x_df[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + scale_idx] = + __half2float(bxi->d) * (ls - 32); + } +} + +template +static __device__ __forceinline__ void vllm_gguf_vec_dot_q8_0_q8_1_mma_v2( + const int* __restrict__ x, const int* __restrict__ y, + float* __restrict__ sum, const int k00) { + using namespace vllm_gguf_mma; + + typedef tile<16, 8, int> tile_A; + typedef tile<8, 8, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int granularity = vllm_gguf_mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp / tile_C::I; + + y += (threadIdx.y % ntx) * (tile_C::J * VLLM_GGUF_MMQ_TILE_Y_K); + + const int* x_qs = (const int*)x; + const float* x_df = (const float*)x_qs + 2 * VLLM_GGUF_MMQ_TILE_NE_K; + const int* y_qs = (const int*)y + 4; + const half2* y_ds = (const half2*)y; + + tile_A A[ntx][VLLM_GGUF_MMQ_TILE_NE_K / QI8_0]; + float dA[ntx][tile_C::ne / 2][VLLM_GGUF_MMQ_TILE_NE_K / QI8_0]; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < VLLM_GGUF_MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + load_ldmatrix( + A[n][k01 / QI8_0], + x_qs + (i0 + n * tile_A::I) * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + k0, + VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne / 2; ++l) { + const int i = i0 + n * tile_A::I + tile_C::get_i(2 * l); + +#pragma unroll + for (int k01 = 0; k01 < VLLM_GGUF_MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + dA[n][l][k01 / QI8_0] = + x_df[i * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0 + k0 / QI8_0]; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < VLLM_GGUF_MMQ_TILE_NE_K; k01 += QI8_0) { + tile_B B; + float dB[tile_C::ne / 2]; + + load_generic(B, y_qs + j0 * VLLM_GGUF_MMQ_TILE_Y_K + k01, + VLLM_GGUF_MMQ_TILE_Y_K); + +#pragma unroll + for (int l = 0; l < tile_C::ne / 2; ++l) { + const int j = j0 + tile_C::get_j(l); + dB[l] = __low2float(y_ds[j * VLLM_GGUF_MMQ_TILE_Y_K + k01 / QI8_1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n][k01 / QI8_0], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0 / tile_C::J + n) * tile_C::ne + l] += + C.x[l] * dA[n][l / 2][k01 / QI8_0] * dB[l % 2]; + } + } + } + } +} + +template +static __device__ __forceinline__ void vllm_gguf_mmq_write_back_mma_v2( + const float* __restrict__ sum, scalar_t* __restrict__ dst, const int stride, + const int i_max, const int j_max) { + using namespace vllm_gguf_mma; + + typedef tile<16, 8, int> tile_C; + constexpr int granularity = vllm_gguf_mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp / tile_C::I; + + const int i0 = (threadIdx.y / ntx) * (ntx * tile_C::I); + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) { +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); + if (j > j_max) { + continue; + } + + const int i = i0 + n * tile_C::I + tile_C::get_i(l); + if (need_check && i > i_max) { + continue; + } + + dst[j * stride + i] = + static_cast(sum[(j0 / tile_C::J + n) * tile_C::ne + l]); + } + } + } +} + +template +static __device__ __forceinline__ void vllm_gguf_load_moe_q8_1_tile_mma_v2( + const int* __restrict__ y, int* __restrict__ tile_y, const int col_dst_0, + const int block_y, const int ncols_y_padded) { + constexpr int warp_size = WARP_SIZE_GGUF; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int sz = sizeof(block_q8_1_mmq_v2) / sizeof(int); + const int* by0 = y + (block_y * ncols_y_padded + col_dst_0) * sz; + +#pragma unroll + for (int l0 = 0; l0 < mmq_x * VLLM_GGUF_MMQ_TILE_Y_K; + l0 += nwarps * warp_size) { + const int l = l0 + threadIdx.y * warp_size + threadIdx.x; + if (l < mmq_x * VLLM_GGUF_MMQ_TILE_Y_K) { + const int k = l % VLLM_GGUF_MMQ_TILE_Y_K; + tile_y[l] = by0[(l / VLLM_GGUF_MMQ_TILE_Y_K) * sz + k]; + } + } +} + +template +static __device__ __forceinline__ void vllm_gguf_moe_mmq_write_back_mma_v2( + const float* __restrict__ sum, scalar_t* __restrict__ dst, + const int* __restrict__ sorted_token_ids, const int col_dst_0, + const int row_dst_0, const int ncols_dst, const int stride, + const int i_max) { + using namespace vllm_gguf_mma; + + typedef tile<16, 8, int> tile_C; + constexpr int granularity = vllm_gguf_mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp / tile_C::I; + + const int i0 = (threadIdx.y / ntx) * (ntx * tile_C::I); + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) { +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); + if (j >= mmq_x) { + continue; + } + const int col_dst = sorted_token_ids[col_dst_0 + j]; + if (col_dst < 0 || col_dst >= ncols_dst) { + continue; + } + + const int i = i0 + n * tile_C::I + tile_C::get_i(l); + if (need_check && i > i_max) { + continue; + } + + dst[col_dst * stride + row_dst_0 + i] = + static_cast(sum[(j0 / tile_C::J + n) * tile_C::ne + l]); + } + } + } +} + +template +static __device__ __forceinline__ void vllm_gguf_mul_mat_q4_0_process_tile_v2( + const char* __restrict__ x, const int offset_x, const int* __restrict__ y, + scalar_t* __restrict__ dst, const int stride_row_x, const int ncols_y, + const int stride_col_dst, const int tile_x_max_i, const int tile_y_max_j, + const int kb0_start, const int kb0_stop) { + constexpr int warp_size = WARP_SIZE_GGUF; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + constexpr int blocks_per_iter = VLLM_GGUF_MMQ_ITER_K / QK4_0; + constexpr int sz = sizeof(block_q8_1_mmq_v2) / sizeof(int); + constexpr int tile_y_ints = + vllm_gguf_mmq_pad(mmq_x * VLLM_GGUF_MMQ_TILE_Y_K, nwarps * warp_size); + + extern __shared__ int data_mul_mat_q_v2[]; + int* tile_y = data_mul_mat_q_v2; + int* tile_x = tile_y + tile_y_ints; + + float sum[mmq_x * mmq_y / (nwarps * warp_size)] = {0.0f}; + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + vllm_gguf_load_tiles_q4_0_v2(x, tile_x, offset_x + kb0, + tile_x_max_i, stride_row_x); + + { + const int* by0 = y + ncols_y * (kb0 * QK4_0 / (4 * QK8_1)) * sz; +#pragma unroll + for (int l0 = 0; l0 < mmq_x * VLLM_GGUF_MMQ_TILE_Y_K; + l0 += nwarps * warp_size) { + const int l = l0 + threadIdx.y * warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + vllm_gguf_vec_dot_q8_0_q8_1_mma_v2(tile_x, tile_y, sum, 0); + __syncthreads(); + + { + const int* by0 = y + ncols_y * ((kb0 * QK4_0 / (4 * QK8_1)) * sz + sz); +#pragma unroll + for (int l0 = 0; l0 < mmq_x * VLLM_GGUF_MMQ_TILE_Y_K; + l0 += nwarps * warp_size) { + const int l = l0 + threadIdx.y * warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + vllm_gguf_vec_dot_q8_0_q8_1_mma_v2(tile_x, tile_y, sum, + VLLM_GGUF_MMQ_TILE_NE_K); + __syncthreads(); + } + + vllm_gguf_mmq_write_back_mma_v2( + sum, dst, stride_col_dst, tile_x_max_i, tile_y_max_j); +} + +template +static __device__ __forceinline__ void vllm_gguf_mul_mat_iq4_xs_process_tile_v2( + const char* __restrict__ x, const int offset_x, const int* __restrict__ y, + scalar_t* __restrict__ dst, const int stride_row_x, const int ncols_y, + const int stride_col_dst, const int tile_x_max_i, const int tile_y_max_j, + const int kb0_start, const int kb0_stop) { + constexpr int warp_size = WARP_SIZE_GGUF; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + constexpr int blocks_per_iter = VLLM_GGUF_MMQ_ITER_K / QK_K; + constexpr int sz = sizeof(block_q8_1_mmq_v2) / sizeof(int); + constexpr int tile_y_ints = + vllm_gguf_mmq_pad(mmq_x * VLLM_GGUF_MMQ_TILE_Y_K, nwarps * warp_size); + + extern __shared__ int data_mul_mat_q_v2[]; + int* tile_y = data_mul_mat_q_v2; + int* tile_x = tile_y + tile_y_ints; + + float sum[mmq_x * mmq_y / (nwarps * warp_size)] = {0.0f}; + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + vllm_gguf_load_tiles_iq4_xs_v2( + x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + + { + const int* by0 = y + ncols_y * (kb0 * QK_K / (4 * QK8_1)) * sz; +#pragma unroll + for (int l0 = 0; l0 < mmq_x * VLLM_GGUF_MMQ_TILE_Y_K; + l0 += nwarps * warp_size) { + const int l = l0 + threadIdx.y * warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + vllm_gguf_vec_dot_q8_0_q8_1_mma_v2(tile_x, tile_y, sum, 0); + __syncthreads(); + + { + const int* by0 = y + ncols_y * ((kb0 * QK_K / (4 * QK8_1)) * sz + sz); +#pragma unroll + for (int l0 = 0; l0 < mmq_x * VLLM_GGUF_MMQ_TILE_Y_K; + l0 += nwarps * warp_size) { + const int l = l0 + threadIdx.y * warp_size + threadIdx.x; + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + vllm_gguf_vec_dot_q8_0_q8_1_mma_v2(tile_x, tile_y, sum, + VLLM_GGUF_MMQ_TILE_NE_K); + __syncthreads(); + } + + vllm_gguf_mmq_write_back_mma_v2( + sum, dst, stride_col_dst, tile_x_max_i, tile_y_max_j); +} + +template +static __device__ __forceinline__ void vllm_gguf_moe_iq4_xs_process_tile_v2( + const char* __restrict__ x, const int offset_x, const int* __restrict__ y, + scalar_t* __restrict__ dst, const int* __restrict__ sorted_token_ids, + const int row_dst_0, const int col_dst_0, const int stride_row_x, + const int ncols_y, const int ncols_y_padded, const int top_k, + const int ncols_dst, const int stride_col_dst, const int tile_x_max_i, + const int kb0_start, const int kb0_stop) { + constexpr int warp_size = WARP_SIZE_GGUF; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + constexpr int blocks_per_iter = VLLM_GGUF_MMQ_ITER_K / QK_K; + constexpr int tile_y_ints = + vllm_gguf_mmq_pad(mmq_x * VLLM_GGUF_MMQ_TILE_Y_K, nwarps * warp_size); + + extern __shared__ int data_mul_mat_q_v2[]; + int* tile_y = data_mul_mat_q_v2; + int* tile_x = tile_y + tile_y_ints; + + float sum[mmq_x * mmq_y / (nwarps * warp_size)] = {0.0f}; + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + vllm_gguf_load_tiles_iq4_xs_v2( + x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + + vllm_gguf_load_moe_q8_1_tile_mma_v2( + y, tile_y, col_dst_0, kb0 * QK_K / (4 * QK8_1), ncols_y_padded); + + __syncthreads(); + vllm_gguf_vec_dot_q8_0_q8_1_mma_v2(tile_x, tile_y, sum, 0); + __syncthreads(); + + vllm_gguf_load_moe_q8_1_tile_mma_v2( + y, tile_y, col_dst_0, kb0 * QK_K / (4 * QK8_1) + 1, ncols_y_padded); + + __syncthreads(); + vllm_gguf_vec_dot_q8_0_q8_1_mma_v2(tile_x, tile_y, sum, + VLLM_GGUF_MMQ_TILE_NE_K); + __syncthreads(); + } + + vllm_gguf_moe_mmq_write_back_mma_v2( + sum, dst, sorted_token_ids, col_dst_0, row_dst_0, ncols_dst, + stride_col_dst, tile_x_max_i); +} + +template +static __global__ __launch_bounds__( + WARP_SIZE_GGUF* VLLM_GGUF_MMQ_NWARPS, + 1) void vllm_gguf_mul_mat_q4_0_v2(const void* __restrict__ vx, + const int* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, const int nrows_x, + const int ncols_y, + const int ncols_y_padded, + const int nrows_dst) { + constexpr int mmq_x = VLLM_GGUF_MMQ_X_Q4_0; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + + const int it = blockIdx.x; + const int jt = blockIdx.y; + const int row_dst_0 = it * mmq_y; + const int col_dst_0 = jt * mmq_x; + const int blocks_per_row_x = ncols_x / QK4_0; + const int tile_x_max_i = nrows_x - row_dst_0 - 1; + const int tile_y_max_j = ncols_y - col_dst_0 - 1; + const int offset_x = row_dst_0 * blocks_per_row_x; + constexpr int sz = sizeof(block_q8_1_mmq_v2) / sizeof(int); + const int offset_y = col_dst_0 * sz; + + vllm_gguf_mul_mat_q4_0_process_tile_v2( + (const char*)vx, offset_x, vy + offset_y, + dst + col_dst_0 * nrows_dst + row_dst_0, blocks_per_row_x, ncols_y_padded, + nrows_dst, tile_x_max_i, tile_y_max_j, 0, blocks_per_row_x); +} + +template +static __global__ __launch_bounds__( + WARP_SIZE_GGUF* VLLM_GGUF_MMQ_NWARPS, + 1) void vllm_gguf_mul_mat_iq4_xs_v2(const void* __restrict__ vx, + const int* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, const int nrows_x, + const int ncols_y, + const int ncols_y_padded, + const int nrows_dst) { + constexpr int mmq_x = VLLM_GGUF_MMQ_X_IQ4_XS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + + const int it = blockIdx.x; + const int jt = blockIdx.y; + const int row_dst_0 = it * mmq_y; + const int col_dst_0 = jt * mmq_x; + const int blocks_per_row_x = ncols_x / QK_K; + const int tile_x_max_i = nrows_x - row_dst_0 - 1; + const int tile_y_max_j = ncols_y - col_dst_0 - 1; + const int offset_x = row_dst_0 * blocks_per_row_x; + constexpr int sz = sizeof(block_q8_1_mmq_v2) / sizeof(int); + const int offset_y = col_dst_0 * sz; + + vllm_gguf_mul_mat_iq4_xs_process_tile_v2( + (const char*)vx, offset_x, vy + offset_y, + dst + col_dst_0 * nrows_dst + row_dst_0, blocks_per_row_x, ncols_y_padded, + nrows_dst, tile_x_max_i, tile_y_max_j, 0, blocks_per_row_x); +} + +template +static __global__ __launch_bounds__( + WARP_SIZE_GGUF* VLLM_GGUF_MMQ_NWARPS, + 1) void vllm_gguf_moe_iq4_xs_v2(const void* __restrict__ vx, + const int* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* __restrict__ sorted_token_ids, + const int* __restrict__ expert_ids, + const int* __restrict__ num_tokens_post_padded, + const int exp_stride, const int ncols_x, + const int nrows_x, const int ncols_y, + const int ncols_y_padded, + const int nrows_dst, const int top_k) { + constexpr int mmq_x = VLLM_GGUF_MOE_MMQ_X_IQ4_XS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + + const int row_dst_0 = blockIdx.x * mmq_y; + const int col_dst_0 = blockIdx.y * mmq_x; + + const int exp_idx = expert_ids[blockIdx.y]; + if (exp_idx > 255 || exp_idx < 0) { + return; + } + if (col_dst_0 >= num_tokens_post_padded[0]) { + return; + } + + const char* x = (const char*)vx + exp_idx * exp_stride; + const int blocks_per_row_x = ncols_x / QK_K; + const int tile_x_max_i = nrows_x - row_dst_0 - 1; + const int offset_x = row_dst_0 * blocks_per_row_x; + const int ncols_dst = ncols_y * top_k; + + vllm_gguf_moe_iq4_xs_process_tile_v2( + x, offset_x, vy, dst, sorted_token_ids, row_dst_0, col_dst_0, + blocks_per_row_x, ncols_y, ncols_y_padded, top_k, ncols_dst, nrows_dst, + tile_x_max_i, 0, blocks_per_row_x); +} + +template +static void vllm_gguf_mul_mat_q4_0_q8_1_mmq_v2_cuda( + const void* vx, const void* vy, scalar_t* dst, const int ncols_x, + const int nrows_x, const int ncols_y, const int ncols_y_padded, + const int nrows_dst, cudaStream_t stream) { + constexpr int mmq_x = VLLM_GGUF_MMQ_X_Q4_0; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int tile_y_ints = vllm_gguf_mmq_pad(mmq_x * VLLM_GGUF_MMQ_TILE_Y_K, + nwarps * WARP_SIZE_GGUF); + constexpr int shared_mem = + (tile_y_ints + mmq_y * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0) * sizeof(int); + + const dim3 block_nums((nrows_x + mmq_y - 1) / mmq_y, + (ncols_y + mmq_x - 1) / mmq_x, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + vllm_gguf_mul_mat_q4_0_v2 + <<>>( + vx, (const int*)vy, dst, ncols_x, nrows_x, ncols_y, ncols_y_padded, + nrows_dst); + } else { + constexpr bool need_check = true; + vllm_gguf_mul_mat_q4_0_v2 + <<>>( + vx, (const int*)vy, dst, ncols_x, nrows_x, ncols_y, ncols_y_padded, + nrows_dst); + } +} + +template +static void vllm_gguf_moe_iq4_xs_q8_1_mmq_v2_cuda( + const void* vx, const void* vy, scalar_t* dst, const int* sorted_token_ids, + const int* expert_ids, const int* num_tokens_post_padded, + const int exp_stride, const int ncols_x, const int nrows_x, + const int ncols_y, const int ncols_y_padded, const int nrows_dst, + const int top_k, const int tokens_post_padded, cudaStream_t stream) { + constexpr int mmq_x = VLLM_GGUF_MOE_MMQ_X_IQ4_XS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int tile_y_ints = vllm_gguf_mmq_pad(mmq_x * VLLM_GGUF_MMQ_TILE_Y_K, + nwarps * WARP_SIZE_GGUF); + constexpr int shared_mem = + (tile_y_ints + mmq_y * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0) * sizeof(int); + + const dim3 block_nums((nrows_x + mmq_y - 1) / mmq_y, + tokens_post_padded / mmq_x, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + vllm_gguf_moe_iq4_xs_v2 + <<>>( + vx, (const int*)vy, dst, sorted_token_ids, expert_ids, + num_tokens_post_padded, exp_stride, ncols_x, nrows_x, ncols_y, + ncols_y_padded, nrows_dst, top_k); + } else { + constexpr bool need_check = true; + vllm_gguf_moe_iq4_xs_v2 + <<>>( + vx, (const int*)vy, dst, sorted_token_ids, expert_ids, + num_tokens_post_padded, exp_stride, ncols_x, nrows_x, ncols_y, + ncols_y_padded, nrows_dst, top_k); + } +} + +template +static void vllm_gguf_mul_mat_iq4_xs_q8_1_mmq_v2_cuda( + const void* vx, const void* vy, scalar_t* dst, const int ncols_x, + const int nrows_x, const int ncols_y, const int ncols_y_padded, + const int nrows_dst, cudaStream_t stream) { + constexpr int mmq_x = VLLM_GGUF_MMQ_X_IQ4_XS; + constexpr int mmq_y = VLLM_GGUF_MMQ_Y; + constexpr int nwarps = VLLM_GGUF_MMQ_NWARPS; + constexpr int tile_y_ints = vllm_gguf_mmq_pad(mmq_x * VLLM_GGUF_MMQ_TILE_Y_K, + nwarps * WARP_SIZE_GGUF); + constexpr int shared_mem = + (tile_y_ints + mmq_y * VLLM_GGUF_MMQ_MMA_TILE_X_K_Q8_0) * sizeof(int); + + const dim3 block_nums((nrows_x + mmq_y - 1) / mmq_y, + (ncols_y + mmq_x - 1) / mmq_x, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + vllm_gguf_mul_mat_iq4_xs_v2 + <<>>( + vx, (const int*)vy, dst, ncols_x, nrows_x, ncols_y, ncols_y_padded, + nrows_dst); + } else { + constexpr bool need_check = true; + vllm_gguf_mul_mat_iq4_xs_v2 + <<>>( + vx, (const int*)vy, dst, ncols_x, nrows_x, ncols_y, ncols_y_padded, + nrows_dst); + } +} diff --git a/vllm_gguf_plugin/csrc/torch_bindings.cpp b/vllm_gguf_plugin/csrc/torch_bindings.cpp index ebc4ee68..8493932e 100644 --- a/vllm_gguf_plugin/csrc/torch_bindings.cpp +++ b/vllm_gguf_plugin/csrc/torch_bindings.cpp @@ -12,11 +12,17 @@ Tensor ggml_dequantize(Tensor W, int64_t type, int64_t m, int64_t n, std::optional dtype); Tensor ggml_mul_mat_vec_a8(Tensor W, Tensor X, int64_t type, int64_t row); Tensor ggml_mul_mat_a8(Tensor W, Tensor X, int64_t type, int64_t row); +Tensor ggml_mul_mat_a8_q4_0_mmq_v2(Tensor W, Tensor X, int64_t row); +Tensor ggml_mul_mat_a8_iq4_xs_mmq_v2(Tensor W, Tensor X, int64_t row); Tensor ggml_moe_a8(Tensor X, Tensor W, Tensor sorted_token_ids, Tensor expert_ids, Tensor num_tokens_post_padded, int64_t type, int64_t row, int64_t top_k, int64_t tokens); Tensor ggml_moe_a8_vec(Tensor X, Tensor W, Tensor topk_ids, int64_t top_k, int64_t type, int64_t row, int64_t tokens); +Tensor ggml_moe_a8_iq4_xs_mmq_v2(Tensor X, Tensor W, Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, int64_t row, + int64_t top_k, int64_t tokens); int64_t ggml_moe_get_block_size(int64_t type); STABLE_TORCH_LIBRARY(_C_gguf, ops) { @@ -28,6 +34,12 @@ STABLE_TORCH_LIBRARY(_C_gguf, ops) { "-> Tensor"); ops.def( "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor"); + ops.def( + "ggml_mul_mat_a8_q4_0_mmq_v2(Tensor W, Tensor X, SymInt row) " + "-> Tensor"); + ops.def( + "ggml_mul_mat_a8_iq4_xs_mmq_v2(Tensor W, Tensor X, SymInt row) " + "-> Tensor"); ops.def( "ggml_moe_a8(Tensor X, Tensor W, " "Tensor sorted_token_ids, Tensor expert_ids, Tensor " @@ -37,6 +49,11 @@ STABLE_TORCH_LIBRARY(_C_gguf, ops) { "ggml_moe_a8_vec(Tensor X, Tensor W, " "Tensor topk_ids, int top_k, " "int type, SymInt row, SymInt tokens) -> Tensor"); + ops.def( + "ggml_moe_a8_iq4_xs_mmq_v2(Tensor X, Tensor W, " + "Tensor sorted_token_ids, Tensor expert_ids, Tensor " + "num_tokens_post_padded, " + "SymInt row, SymInt top_k, SymInt tokens) -> Tensor"); ops.def("ggml_moe_get_block_size(int type) -> int"); } @@ -44,8 +61,14 @@ STABLE_TORCH_LIBRARY_IMPL(_C_gguf, CUDA, ops) { ops.impl("ggml_dequantize", TORCH_BOX(&ggml_dequantize)); ops.impl("ggml_mul_mat_vec_a8", TORCH_BOX(&ggml_mul_mat_vec_a8)); ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8)); + ops.impl("ggml_mul_mat_a8_q4_0_mmq_v2", + TORCH_BOX(&ggml_mul_mat_a8_q4_0_mmq_v2)); + ops.impl("ggml_mul_mat_a8_iq4_xs_mmq_v2", + TORCH_BOX(&ggml_mul_mat_a8_iq4_xs_mmq_v2)); ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8)); ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec)); + ops.impl("ggml_moe_a8_iq4_xs_mmq_v2", + TORCH_BOX(&ggml_moe_a8_iq4_xs_mmq_v2)); } STABLE_TORCH_LIBRARY_IMPL(_C_gguf, CompositeExplicitAutograd, ops) { diff --git a/vllm_gguf_plugin/gguf_tokenizer_builder.py b/vllm_gguf_plugin/gguf_tokenizer_builder.py new file mode 100644 index 00000000..cf4e722f --- /dev/null +++ b/vllm_gguf_plugin/gguf_tokenizer_builder.py @@ -0,0 +1,729 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Build a HF-compatible tokenizer directory from GGUF embedded metadata.""" + +import hashlib +import json +import os +from contextlib import suppress +from os import PathLike +from pathlib import Path +from shutil import copyfile +from typing import Any + +import gguf +from transformers import PreTrainedTokenizerFast +from transformers.integrations.ggml import ( + GGUF_TOKENIZER_MAPPING, + convert_gguf_tokenizer, +) +from vllm.logger import init_logger + +from .gguf_utils import ( + _gguf_reader_value, + _gguf_scalar_value, + check_gguf_file, + detect_gguf_multimodal, +) + +logger = init_logger(__name__) + +_TOKENIZER_CACHE_ENV = "VLLM_GGUF_TOKENIZER_CACHE" +_DEFAULT_TOKENIZER_CACHE = "~/.cache/vllm-gguf-plugin/tokenizers" +_PROCESSOR_SIDECAR_FILES = ( + "processor_config.json", + "preprocessor_config.json", + "video_preprocessor_config.json", + "image_processor_config.json", +) +_PROCESSOR_DEFAULT_SOURCE = ( + "non-GGUF processor defaults mirrored from Unsloth Qwen3.6/Gemma4 " + "GGUF sidecars observed on 2026-06-11; GGUF metadata values take " + "precedence when present" +) + +_TOKENIZER_ARCH_ALIASES = { + "qwen35": "qwen3", + "qwen3_5": "qwen3", + "qwen35moe": "qwen3_moe", + "qwen3_5_moe": "qwen3_moe", + "gemma4": "gemma3_text", + "gemma4-assistant": "gemma3_text", + "gemma4_assistant": "gemma3_text", +} + + +def _decode_value(value: Any) -> Any: + value = _gguf_scalar_value(value) + if isinstance(value, bytes): + with suppress(UnicodeDecodeError): + return value.decode("utf-8") + return value + + +def _decode_sequence(value: Any) -> Any: + if value is None or isinstance(value, (str, bytes)): + return _decode_value(value) + with suppress(AttributeError): + value = value.tolist() + if isinstance(value, tuple): + value = list(value) + if isinstance(value, list): + return [_decode_value(item) for item in value] + return _decode_value(value) + + +def _gguf_architecture(reader: gguf.GGUFReader) -> str | None: + value = _decode_value(_gguf_reader_value(reader, "general.architecture")) + return value if isinstance(value, str) else None + + +def _tokenizer_cache_root() -> Path: + return Path( + os.environ.get(_TOKENIZER_CACHE_ENV, _DEFAULT_TOKENIZER_CACHE) + ).expanduser() + + +def _cache_key(model_path: Path) -> str | None: + try: + stat = model_path.stat() + except OSError as e: + logger.debug("Failed to stat GGUF tokenizer source %s: %s", model_path, e) + return None + raw_key = f"{model_path.resolve()}:{stat.st_size}:{stat.st_mtime_ns}" + return hashlib.sha256(raw_key.encode("utf-8")).hexdigest()[:24] + + +def _extract_tokenizer_dict(reader: gguf.GGUFReader) -> dict[str, Any]: + tokenizer_dict: dict[str, Any] = {} + field_mapping = GGUF_TOKENIZER_MAPPING["tokenizer"] + for gguf_suffix, hf_name in field_mapping.items(): + value = _gguf_reader_value(reader, f"tokenizer.{gguf_suffix}") + if value is not None: + tokenizer_dict[hf_name] = _decode_sequence(value) + token_type_value = _gguf_reader_value(reader, "tokenizer.ggml.token_type") + if token_type_value is not None: + tokenizer_dict["token_type"] = _decode_sequence(token_type_value) + return tokenizer_dict + + +def _extract_tokenizer_config(reader: gguf.GGUFReader) -> dict[str, Any]: + tokenizer_config: dict[str, Any] = {} + field_mapping = GGUF_TOKENIZER_MAPPING["tokenizer_config"] + for gguf_suffix, hf_name in field_mapping.items(): + value = _gguf_reader_value(reader, f"tokenizer.{gguf_suffix}") + if value is not None: + tokenizer_config[hf_name] = _decode_sequence(value) + return tokenizer_config + + +def _token_by_id(tokens: list[Any], token_id: Any) -> str | None: + if token_id is None: + return None + if isinstance(token_id, (list, tuple)): + if not token_id: + return None + token_id = token_id[0] + with suppress(TypeError, ValueError, IndexError): + token = tokens[int(token_id)] + if isinstance(token, str): + return token + return None + + +def _special_token_kwargs(tokenizer_dict: dict[str, Any]) -> dict[str, str]: + tokens = tokenizer_dict.get("tokens") + if not isinstance(tokens, list): + return {} + special_ids = { + "bos_token": tokenizer_dict.get("bos_token_id"), + "eos_token": tokenizer_dict.get("eos_token_id"), + "pad_token": tokenizer_dict.get("pad_token_id"), + "unk_token": tokenizer_dict.get("unk_token_id"), + } + return { + key: token + for key, token_id in special_ids.items() + if (token := _token_by_id(tokens, token_id)) is not None + } + + +def _local_config_path(model_path: Path) -> Path | None: + """Return the nearest local HF config next to a GGUF file, if present.""" + for candidate in (model_path.parent, model_path.parent.parent): + config_path = candidate / "config.json" + if config_path.is_file(): + return config_path + return None + + +def _read_special_token_ids_from_config(config: dict[str, Any]) -> dict[str, Any]: + text_config = config.get("text_config") + if not isinstance(text_config, dict): + text_config = {} + + special_ids: dict[str, Any] = {} + for token_name, config_key in { + "bos_token": "bos_token_id", + "eos_token": "eos_token_id", + "pad_token": "pad_token_id", + "unk_token": "unk_token_id", + }.items(): + token_id = config.get(config_key) + if token_id is None: + token_id = text_config.get(config_key) + if token_id is not None: + special_ids[token_name] = token_id + return special_ids + + +def _local_config_special_token_kwargs( + model_path: Path, + tokenizer_dict: dict[str, Any], +) -> dict[str, str]: + tokens = tokenizer_dict.get("tokens") + if not isinstance(tokens, list): + return {} + + config_path = _local_config_path(model_path) + if config_path is None: + return {} + + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except Exception as e: + logger.debug("Failed to read local GGUF config %s: %s", config_path, e) + return {} + + return { + token_name: token + for token_name, token_id in _read_special_token_ids_from_config( + config, + ).items() + if (token := _token_by_id(tokens, token_id)) is not None + } + + +_GEMMA4_MODEL_SPECIFIC_TOKENS = { + "audio_token": "<|audio|>", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eot_token": "", + "escape_token": '<|"|>', + "etc_token": "", + "etd_token": "", + "etr_token": "", + "image_token": "<|image|>", + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>", + "video_token": "<|video|>", +} + +_QWEN_MM_SPECIAL_TOKENS = ( + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>", +) + + +def _tokens_present( + tokenizer_dict: dict[str, Any], + candidates: tuple[str, ...], +) -> list[str]: + tokens = tokenizer_dict.get("tokens") + if not isinstance(tokens, list): + return [] + token_set = {token for token in tokens if isinstance(token, str)} + return [token for token in candidates if token in token_set] + + +def _gemma4_model_specific_special_tokens( + tokenizer_dict: dict[str, Any], +) -> dict[str, str]: + return { + name: token + for name, token in _GEMMA4_MODEL_SPECIFIC_TOKENS.items() + if token in _tokens_present(tokenizer_dict, (token,)) + } + + +def _append_additional_special_tokens( + tokenizer_config: dict[str, Any], + tokens: list[str], +) -> bool: + if not tokens: + return False + + existing = tokenizer_config.get("additional_special_tokens") + if existing is None: + additional_tokens: list[Any] = [] + elif isinstance(existing, list): + additional_tokens = list(existing) + else: + additional_tokens = [existing] + + existing_tokens: set[str] = set() + for item in additional_tokens: + if isinstance(item, str): + existing_tokens.add(item) + elif isinstance(item, dict) and isinstance(item.get("content"), str): + existing_tokens.add(item["content"]) + + changed = False + for token in tokens: + if token in existing_tokens: + continue + additional_tokens.append(token) + existing_tokens.add(token) + changed = True + + if changed: + tokenizer_config["additional_special_tokens"] = additional_tokens + return changed + + +def _remove_additional_special_tokens( + tokenizer_config: dict[str, Any], + tokens: set[str], +) -> bool: + if not tokens: + return False + + existing = tokenizer_config.get("additional_special_tokens") + if existing is None: + return False + additional_tokens = list(existing) if isinstance(existing, list) else [existing] + + filtered_tokens: list[Any] = [] + changed = False + for item in additional_tokens: + token = None + if isinstance(item, str): + token = item + elif isinstance(item, dict) and isinstance(item.get("content"), str): + token = item["content"] + if token in tokens: + changed = True + continue + filtered_tokens.append(item) + + if changed: + tokenizer_config["additional_special_tokens"] = filtered_tokens + return changed + + +_GGUF_SPECIAL_TOKEN_TYPES = { + int(gguf.TokenType.CONTROL), + int(gguf.TokenType.USER_DEFINED), +} + + +def _gguf_special_control_tokens( + tokenizer_dict: dict[str, Any], + special_token_kwargs: dict[str, str] | None = None, +) -> list[str]: + """Restore GGUF control/user-defined tokens as HF special tokens.""" + tokens = tokenizer_dict.get("tokens") + token_types = tokenizer_dict.get("token_type") + if not isinstance(tokens, list) or not isinstance(token_types, list): + return [] + + named_special_tokens = set(_special_token_kwargs(tokenizer_dict).values()) + if special_token_kwargs: + named_special_tokens.update(special_token_kwargs.values()) + special_tokens: list[str] = [] + for token, token_type in zip(tokens, token_types, strict=False): + if not isinstance(token, str) or token in named_special_tokens: + continue + with suppress(TypeError, ValueError): + if int(token_type) in _GGUF_SPECIAL_TOKEN_TYPES: + special_tokens.append(token) + return special_tokens + + +def _patch_tokenizer_config_from_gguf( + cache_dir: Path, + architecture: str, + tokenizer_dict: dict[str, Any], + model_path: Path, +) -> None: + tokenizer_config_path = cache_dir / "tokenizer_config.json" + if not tokenizer_config_path.is_file(): + return + + try: + tokenizer_config = json.loads(tokenizer_config_path.read_text(encoding="utf-8")) + except Exception as e: + logger.debug("Failed to read tokenizer config %s: %s", tokenizer_config_path, e) + return + + changed = False + special_token_kwargs = _local_config_special_token_kwargs( + model_path, + tokenizer_dict, + ) + for token_name, token in special_token_kwargs.items(): + if tokenizer_config.get(token_name) != token: + tokenizer_config[token_name] = token + changed = True + changed |= _remove_additional_special_tokens( + tokenizer_config, + set(special_token_kwargs.values()), + ) + + changed |= _append_additional_special_tokens( + tokenizer_config, + _gguf_special_control_tokens(tokenizer_dict, special_token_kwargs), + ) + + if architecture in {"gemma4", "gemma4-assistant", "gemma4_assistant"}: + model_specific_tokens = _gemma4_model_specific_special_tokens(tokenizer_dict) + if model_specific_tokens: + tokenizer_config["processor_class"] = "Gemma4Processor" + tokenizer_config["model_specific_special_tokens"] = model_specific_tokens + for name, token in model_specific_tokens.items(): + tokenizer_config.setdefault(name, token) + tokenizer_config["extra_special_tokens"] = { + **( + tokenizer_config["extra_special_tokens"] + if isinstance(tokenizer_config.get("extra_special_tokens"), dict) + else {} + ), + **model_specific_tokens, + } + changed = True + + if architecture in {"qwen35", "qwen3_5", "qwen35moe", "qwen3_5_moe"}: + mm_tokens = _tokens_present(tokenizer_dict, _QWEN_MM_SPECIAL_TOKENS) + changed |= _append_additional_special_tokens(tokenizer_config, mm_tokens) + if "<|image_pad|>" in mm_tokens: + tokenizer_config.setdefault("image_token", "<|image_pad|>") + changed = True + if "<|video_pad|>" in mm_tokens: + tokenizer_config.setdefault("video_token", "<|video_pad|>") + changed = True + + if changed: + tokenizer_config_path.write_text( + json.dumps(tokenizer_config, indent=2) + "\n", + encoding="utf-8", + ) + + +def _read_int(reader: gguf.GGUFReader, key: str) -> int | None: + value = _gguf_scalar_value(_gguf_reader_value(reader, key)) + if value is None: + return None + with suppress(TypeError, ValueError): + return int(value) + return None + + +def _copy_local_processor_sidecars(model_path: Path, cache_dir: Path) -> None: + """Copy local sidecar files next to the GGUF, without network fallback.""" + for filename in _PROCESSOR_SIDECAR_FILES: + source = model_path.parent / filename + target = cache_dir / filename + if target.is_file() or not source.is_file(): + continue + try: + cache_dir.mkdir(parents=True, exist_ok=True) + copyfile(source, target) + except Exception as e: + logger.debug("Failed to copy local GGUF sidecar %s: %s", source, e) + + +def _write_json_if_missing( + cache_dir: Path, + filename: str, + data: dict[str, Any], +) -> None: + target = cache_dir / filename + if target.is_file(): + return + cache_dir.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8") + + +def _read_mmproj_reader(model_path: Path) -> gguf.GGUFReader | None: + mmproj_path = detect_gguf_multimodal(str(model_path)) + if mmproj_path is None: + return None + try: + return gguf.GGUFReader(str(mmproj_path)) + except Exception as e: + logger.debug("Failed to read GGUF mmproj sidecar %s: %s", mmproj_path, e) + return None + + +def _qwen_image_processor_config(reader: gguf.GGUFReader) -> dict[str, Any]: + # See _PROCESSOR_DEFAULT_SOURCE for the provenance of non-GGUF defaults. + return { + "do_convert_rgb": True, + "do_normalize": True, + "do_rescale": True, + "do_resize": True, + "image_mean": [0.5, 0.5, 0.5], + "image_processor_type": "Qwen2VLImageProcessor", + "image_std": [0.5, 0.5, 0.5], + "merge_size": _read_int(reader, "clip.vision.spatial_merge_size") or 2, + "patch_size": _read_int(reader, "clip.vision.patch_size") or 16, + "resample": 3, + "rescale_factor": 1.0 / 255.0, + "size": { + "longest_edge": 16777216, + "shortest_edge": 65536, + }, + "temporal_patch_size": ( + _read_int(reader, "clip.vision.temporal_patch_size") or 2 + ), + } + + +def _qwen_video_processor_config(reader: gguf.GGUFReader) -> dict[str, Any]: + # See _PROCESSOR_DEFAULT_SOURCE for the provenance of non-GGUF defaults. + config = _qwen_image_processor_config(reader) + config.update( + { + "do_sample_frames": True, + "fps": 2, + "max_frames": 768, + "min_frames": 4, + "return_metadata": False, + "size": { + "longest_edge": 25165824, + "shortest_edge": 4096, + }, + "video_processor_type": "Qwen3VLVideoProcessor", + } + ) + config.pop("image_processor_type", None) + return config + + +def _gemma4_image_processor_config(reader: gguf.GGUFReader) -> dict[str, Any]: + # See _PROCESSOR_DEFAULT_SOURCE for the provenance of non-GGUF defaults. + patch_size = _read_int(reader, "clip.vision.patch_size") or 16 + image_seq_length = _read_int(reader, "clip.vision.default_output_length") or 280 + return { + "do_convert_rgb": True, + "do_normalize": False, + "do_rescale": True, + "do_resize": True, + "image_mean": [0.0, 0.0, 0.0], + "image_processor_type": "Gemma4ImageProcessor", + "image_seq_length": image_seq_length, + "image_std": [1.0, 1.0, 1.0], + "max_soft_tokens": image_seq_length, + "patch_size": patch_size, + "pooling_kernel_size": 3, + "resample": 3, + "rescale_factor": 1.0 / 255.0, + } + + +def _gemma4_video_processor_config(reader: gguf.GGUFReader) -> dict[str, Any]: + # See _PROCESSOR_DEFAULT_SOURCE for the provenance of non-GGUF defaults. + patch_size = _read_int(reader, "clip.vision.patch_size") or 16 + return { + "do_convert_rgb": True, + "do_normalize": True, + "do_rescale": True, + "do_resize": True, + "do_sample_frames": True, + "image_mean": [0.0, 0.0, 0.0], + "image_std": [1.0, 1.0, 1.0], + "max_soft_tokens": 70, + "num_frames": 32, + "patch_size": patch_size, + "pooling_kernel_size": 3, + "resample": 3, + "rescale_factor": 1.0 / 255.0, + "return_metadata": False, + "video_processor_type": "Gemma4VideoProcessor", + } + + +def _gemma4_feature_extractor_config() -> dict[str, Any]: + # See _PROCESSOR_DEFAULT_SOURCE for the provenance of non-GGUF defaults. + return { + "dither": 0.0, + "feature_extractor_type": "Gemma4AudioFeatureExtractor", + "feature_size": 128, + "fft_length": 512, + "fft_overdrive": False, + "frame_length": 320, + "hop_length": 160, + "input_scale_factor": 1.0, + "max_frequency": 8000.0, + "mel_floor": 0.001, + "min_frequency": 0.0, + "padding_side": "right", + "padding_value": 0.0, + "per_bin_mean": None, + "per_bin_stddev": None, + "preemphasis": 0.0, + "preemphasis_htk_flavor": True, + "return_attention_mask": True, + "sampling_rate": 16000, + } + + +def _processor_sidecars_from_metadata( + architecture: str, + mmproj_reader: gguf.GGUFReader, +) -> dict[str, dict[str, Any]]: + if architecture in {"qwen35", "qwen3_5", "qwen35moe", "qwen3_5_moe"}: + image_config = _qwen_image_processor_config(mmproj_reader) + video_config = _qwen_video_processor_config(mmproj_reader) + return { + "processor_config.json": { + "image_processor": image_config, + "processor_class": "Qwen3VLProcessor", + "video_processor": video_config, + }, + "preprocessor_config.json": image_config, + "video_preprocessor_config.json": video_config, + } + + if architecture == "gemma4": + image_config = _gemma4_image_processor_config(mmproj_reader) + video_config = _gemma4_video_processor_config(mmproj_reader) + return { + "processor_config.json": { + "audio_ms_per_token": 40, + "audio_seq_length": 750, + "feature_extractor": _gemma4_feature_extractor_config(), + "image_processor": image_config, + "image_seq_length": image_config["image_seq_length"], + "processor_class": "Gemma4Processor", + "video_processor": video_config, + }, + "preprocessor_config.json": image_config, + "video_preprocessor_config.json": video_config, + } + + return {} + + +def _materialize_processor_sidecars( + model_path: Path, + cache_dir: Path, + architecture: str | None = None, +) -> None: + """Materialize processor sidecars without implicit HF Hub access.""" + _copy_local_processor_sidecars(model_path, cache_dir) + mmproj_reader = _read_mmproj_reader(model_path) + if mmproj_reader is None: + return + + if architecture is None: + try: + architecture = _gguf_architecture(gguf.GGUFReader(str(model_path))) + except Exception as e: + logger.debug("Failed to read GGUF architecture from %s: %s", model_path, e) + return + if architecture is None: + return + + for filename, data in _processor_sidecars_from_metadata( + architecture, + mmproj_reader, + ).items(): + _write_json_if_missing(cache_dir, filename, data) + + +def _patch_cached_tokenizer_from_gguf( + model_path: Path, + cache_dir: Path, +) -> str | None: + try: + reader = gguf.GGUFReader(str(model_path)) + architecture = _gguf_architecture(reader) + if architecture is None: + return None + tokenizer_dict = _extract_tokenizer_dict(reader) + except Exception as e: + logger.debug( + "Failed to read cached GGUF tokenizer metadata %s: %s", model_path, e + ) + return None + + _patch_tokenizer_config_from_gguf( + cache_dir, + architecture, + tokenizer_dict, + model_path, + ) + _materialize_processor_sidecars(model_path, cache_dir, architecture) + return architecture + + +def build_tokenizer_from_gguf(model: str | PathLike) -> str | None: + """Materialize a tokenizer directory from GGUF embedded metadata. + + Returns the cache directory path on success. Returns ``None`` if the GGUF + tokenizer is unsupported or incomplete, so callers can use existing + tokenizer fallback behavior. + """ + model_path = Path(model) + if not check_gguf_file(model_path): + return None + + cache_key = _cache_key(model_path) + if cache_key is None: + return None + cache_dir = _tokenizer_cache_root() / cache_key + if (cache_dir / "tokenizer.json").is_file(): + _patch_cached_tokenizer_from_gguf(model_path, cache_dir) + return str(cache_dir) + + try: + reader = gguf.GGUFReader(str(model_path)) + architecture = _gguf_architecture(reader) + if architecture is None: + return None + tokenizer_architecture = _TOKENIZER_ARCH_ALIASES.get( + architecture, + architecture, + ) + tokenizer_dict = _extract_tokenizer_dict(reader) + tokenizer_config = _extract_tokenizer_config(reader) + backend_tokenizer, additional_kwargs = convert_gguf_tokenizer( + tokenizer_architecture, + tokenizer_dict, + ) + special_token_kwargs = _special_token_kwargs(tokenizer_dict) + special_token_kwargs.update( + _local_config_special_token_kwargs(model_path, tokenizer_dict) + ) + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend_tokenizer, + **additional_kwargs, + **special_token_kwargs, + ) + if chat_template := tokenizer_config.get("chat_template"): + tokenizer.chat_template = chat_template + except Exception as e: + logger.debug("Failed to build tokenizer from GGUF %s: %s", model_path, e) + return None + + cache_dir.mkdir(parents=True, exist_ok=True) + tokenizer.save_pretrained(cache_dir) + _patch_tokenizer_config_from_gguf( + cache_dir, + architecture, + tokenizer_dict, + model_path, + ) + _materialize_processor_sidecars(model_path, cache_dir, architecture) + return str(cache_dir) diff --git a/vllm_gguf_plugin/gguf_utils.py b/vllm_gguf_plugin/gguf_utils.py index f4156316..b84ec210 100644 --- a/vllm_gguf_plugin/gguf_utils.py +++ b/vllm_gguf_plugin/gguf_utils.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GGUF utility functions.""" +from contextlib import suppress from functools import cache from os import PathLike from pathlib import Path +from typing import Any import gguf import regex as re @@ -12,7 +14,11 @@ from gguf.quants import GGMLQuantizationType from transformers import Gemma3Config, PretrainedConfig, SiglipVisionConfig from vllm.logger import init_logger -from vllm.transformers_utils.repo_utils import list_filtered_repo_files +from vllm.transformers_utils.repo_utils import ( + file_or_path_exists, + hf_api, + list_filtered_repo_files, +) logger = init_logger(__name__) @@ -83,6 +89,14 @@ def is_nonstandard_gguf_quant_type(quant_type: str) -> bool: # Common suffixes used in GGUF file naming conventions # e.g., Q4_K_M, Q3_K_S, Q5_K_L, Q2_K_XL _GGUF_QUANT_SUFFIXES = ("_M", "_S", "_L", "_XL", "_XS", "_XXS") +_HF_CONFIG_FILES = ("config.json",) +_HF_REPO_ID_PATTERN = re.compile( + r"^[a-zA-Z0-9][a-zA-Z0-9._-]*/[a-zA-Z0-9][a-zA-Z0-9._-]*$" +) +_HF_REPO_URL_PATTERN = re.compile( + r"^https?://huggingface\.co/" + r"([a-zA-Z0-9][a-zA-Z0-9._-]*/[a-zA-Z0-9][a-zA-Z0-9._-]*)" +) def is_valid_gguf_quant_type(gguf_quant_type: str) -> bool: @@ -129,6 +143,308 @@ def split_remote_gguf(model: str | Path) -> tuple[str, str]: ) +def _normalize_base_model_ids(base_model: Any) -> list[str]: + if base_model is None: + return [] + if isinstance(base_model, str): + return [base_model] if base_model else [] + if isinstance(base_model, (list, tuple, set)): + return [model_id for model_id in base_model if isinstance(model_id, str)] + return [] + + +def _normalize_hf_repo_id(value: Any) -> str | None: + if not isinstance(value, str): + return None + + value = value.strip() + if value.endswith(".git"): + value = value[:-4] + + if _HF_REPO_ID_PATTERN.fullmatch(value): + return value + + match = _HF_REPO_URL_PATTERN.match(value) + if match: + repo_id = match.group(1) + if repo_id.endswith(".git"): + repo_id = repo_id[:-4] + return repo_id + + return None + + +@cache +def _get_remote_gguf_base_model_ids( + repo_id: str, + revision: str | None = None, +) -> tuple[str, ...]: + try: + info = hf_api().model_info(repo_id, revision=revision) + except Exception as e: + logger.debug("Failed to inspect GGUF model card for %s: %s", repo_id, e) + return () + + card_data = getattr(info, "card_data", None) + base_model = getattr(card_data, "base_model", None) + if base_model is None and isinstance(card_data, dict): + base_model = card_data.get("base_model") + + base_model_ids: list[str] = [] + for value in _normalize_base_model_ids(base_model): + if normalized_repo_id := _normalize_hf_repo_id(value): + base_model_ids.append(normalized_repo_id) + + return tuple(dict.fromkeys(base_model_ids)) + + +def _gguf_field_value(field: Any) -> Any: + try: + return field.contents() + except Exception as e: + logger.debug("Failed to read GGUF metadata field: %s", e) + return None + + +def _gguf_reader_value(reader: gguf.GGUFReader, key: str) -> Any: + field = reader.get_field(key) + if field is None: + return None + return _gguf_field_value(field) + + +def _gguf_scalar_value(value: Any) -> Any: + if value is None: + return None + if isinstance(value, (str, bytes)): + return value + try: + return value.item() + except (AttributeError, ValueError, TypeError): + pass + with suppress(AttributeError): + value = value.tolist() + if isinstance(value, (list, tuple)): + if len(value) != 1: + return None + return _gguf_scalar_value(value[0]) + return value + + +def _update_config(config: PretrainedConfig, values: dict[str, Any]) -> None: + values = {key: value for key, value in values.items() if value is not None} + if values: + config.update(values) + + +def _gguf_sequence_edge(value: Any, *, first: bool) -> Any: + if value is None: + return None + if isinstance(value, (str, bytes)): + return value + try: + return value[0] if first else value[-1] + except (TypeError, KeyError, IndexError): + return value + + +def _gguf_int_value(reader: gguf.GGUFReader, key: str) -> int | None: + value = _gguf_scalar_value(_gguf_reader_value(reader, key)) + if value is None: + return None + with suppress(TypeError, ValueError): + return int(value) + return None + + +def _gguf_float_value(reader: gguf.GGUFReader, key: str) -> float | None: + value = _gguf_scalar_value(_gguf_reader_value(reader, key)) + if value is None: + return None + with suppress(TypeError, ValueError): + return float(value) + return None + + +def _gguf_string_value(reader: gguf.GGUFReader, key: str) -> str | None: + value = _gguf_scalar_value(_gguf_reader_value(reader, key)) + if isinstance(value, bytes): + with suppress(UnicodeDecodeError): + return value.decode("utf-8") + return None + if isinstance(value, str): + return value + if value is None: + return None + return str(value) + + +def _gguf_int_list_value(reader: gguf.GGUFReader, key: str) -> list[int] | None: + value = _gguf_reader_value(reader, key) + if value is None or isinstance(value, (str, bytes)): + return None + with suppress(AttributeError): + value = value.tolist() + if isinstance(value, tuple): + value = list(value) + if not isinstance(value, list): + return None + + values: list[int] = [] + for item in value: + scalar = _gguf_scalar_value(item) + with suppress(TypeError, ValueError): + values.append(int(scalar)) + continue + return None + return values + + +def _qwen35_rope_parameters_from_gguf( + reader: gguf.GGUFReader, + prefix: str, + partial_rotary_factor: float | None, +) -> dict[str, Any]: + rope_parameters: dict[str, Any] = { + "rope_type": "default", + "rope_theta": _gguf_float_value(reader, f"{prefix}.rope.freq_base"), + "partial_rotary_factor": partial_rotary_factor, + } + + mrope_section = _gguf_int_list_value( + reader, + f"{prefix}.rope.dimension_sections", + ) + if mrope_section: + mrope_section = list(mrope_section) + while mrope_section and mrope_section[-1] == 0: + mrope_section.pop() + if mrope_section: + rope_parameters["mrope_section"] = mrope_section + rope_parameters["mrope_interleaved"] = True + + return {key: value for key, value in rope_parameters.items() if value is not None} + + +def _qwen35_text_config_updates_from_gguf( + reader: gguf.GGUFReader, + prefix: str, +) -> dict[str, Any]: + head_dim = _gguf_int_value(reader, f"{prefix}.attention.key_length") + rope_dim = _gguf_int_value(reader, f"{prefix}.rope.dimension_count") + partial_rotary_factor = None + if head_dim and rope_dim: + partial_rotary_factor = rope_dim / head_dim + + rope_parameters = _qwen35_rope_parameters_from_gguf( + reader, + prefix, + partial_rotary_factor, + ) + updates: dict[str, Any] = { + # These are Qwen3.5 HF/vLLM config defaults. Current GGUF metadata does + # not carry them, while llama.cpp handles the same behavior in model + # code. Keep the defaults explicit so the GGUF-derived config preserves + # the base model contract. + "attn_output_gate": True, + "mamba_ssm_dtype": "float32", + "mtp_use_dedicated_embeddings": False, + "partial_rotary_factor": partial_rotary_factor, + "rope_theta": _gguf_float_value(reader, f"{prefix}.rope.freq_base"), + "full_attention_interval": _gguf_int_value( + reader, + f"{prefix}.full_attention_interval", + ), + "output_gate_type": _gguf_string_value(reader, f"{prefix}.output_gate_type"), + } + if "mrope_section" in rope_parameters: + updates["rope_parameters"] = rope_parameters + return updates + + +@cache +def _get_local_gguf_base_model_ids(model: str | Path) -> tuple[str, ...]: + try: + reader = gguf.GGUFReader(str(model)) + except Exception as e: + logger.debug("Failed to inspect GGUF metadata for %s: %s", model, e) + return () + + base_model_ids: list[str] = [] + for key, field in reader.fields.items(): + if not (key.startswith("general.base_model.") and key.endswith(".repo_url")): + continue + if repo_id := _normalize_hf_repo_id(_gguf_field_value(field)): + base_model_ids.append(repo_id) + + return tuple(dict.fromkeys(base_model_ids)) + + +def _source_has_any_file( + model: str | Path, + filenames: tuple[str, ...], + revision: str | None = None, +) -> bool: + return any(file_or_path_exists(model, filename, revision) for filename in filenames) + + +def _local_gguf_source_candidates(source: Path) -> tuple[Path, ...]: + parent = source.parent + if parent == source: + return (source,) + return (source, parent) + + +def _resolve_gguf_hf_source( + model: str | Path, + filenames: tuple[str, ...], + revision: str | None = None, +) -> str | Path: + local_sources: tuple[Path, ...] = () + if is_remote_gguf(model): + source: str | Path + source, _ = split_remote_gguf(model) + base_model_ids = list( + _get_remote_gguf_base_model_ids(source, revision=revision) + ) + elif check_gguf_file(model): + source = Path(model).parent + local_sources = _local_gguf_source_candidates(source) + base_model_ids = list(_get_local_gguf_base_model_ids(model)) + else: + return model + + for local_source in local_sources: + if _source_has_any_file(local_source, filenames, revision=revision): + return local_source + if not local_sources and _source_has_any_file(source, filenames, revision=revision): + return source + + for base_model in base_model_ids: + if _source_has_any_file(base_model, filenames, revision=None): + logger.warning_once( + "GGUF metadata redirects HF config loading from %s to base " + "model '%s'. `trust_remote_code` is not inherited across " + "implicit GGUF base-model redirects; pass an explicit " + "`--hf-config-path` to opt in for that repository.", + model, + base_model, + ) + return base_model + + return source + + +def resolve_gguf_config_source( + model: str | Path, + revision: str | None = None, +) -> str | Path: + """Resolve where a GGUF model should load its HF config from.""" + if is_gguf(model): + return _resolve_gguf_hf_source(model, _HF_CONFIG_FILES, revision=revision) + return model + + def is_local_gguf_quant(model: str | Path) -> bool: """Check if the model is a local path with GGUF quant type. @@ -220,11 +536,16 @@ def extract_vision_config_from_gguf(mmproj_path: str) -> "SiglipVisionConfig | N # Detect projector type to apply model-specific parameters projector_type = None - projector_type_field = reader.get_field(Keys.Clip.PROJECTOR_TYPE) - if projector_type_field: + projector_type_value = _gguf_scalar_value( + _gguf_reader_value(reader, Keys.Clip.PROJECTOR_TYPE) + ) + if projector_type_value is not None: try: - projector_type = bytes(projector_type_field.parts[-1]).decode("utf-8") - except (AttributeError, UnicodeDecodeError) as e: + if isinstance(projector_type_value, bytes): + projector_type = projector_type_value.decode("utf-8") + else: + projector_type = str(projector_type_value) + except UnicodeDecodeError as e: logger.warning("Failed to decode projector type from GGUF: %s", e) # Map GGUF field constants to SiglipVisionConfig parameters. @@ -243,15 +564,23 @@ def extract_vision_config_from_gguf(mmproj_path: str) -> "SiglipVisionConfig | N # Extract and validate all required fields config_params = {} for gguf_key, (param_name, dtype) in VISION_CONFIG_FIELDS.items(): - field = reader.get_field(gguf_key) - if field is None: + value = _gguf_scalar_value(_gguf_reader_value(reader, gguf_key)) + if value is None: logger.warning( "Missing required vision config field '%s' in mmproj.gguf", gguf_key, ) return None # Extract scalar value from GGUF field and convert to target type - config_params[param_name] = dtype(field.parts[-1]) + try: + config_params[param_name] = dtype(value) + except (TypeError, ValueError) as e: + logger.warning( + "Invalid vision config field '%s' in mmproj.gguf: %s", + gguf_key, + e, + ) + return None # Apply model-specific parameters based on projector type if projector_type == VisionProjectorType.GEMMA3: @@ -335,6 +664,188 @@ def maybe_patch_hf_config_from_gguf( text_config = hf_config.get_text_config() text_config.update({"tie_word_embeddings": not has_lm_head}) + if check_gguf_file(model): + try: + reader = gguf.GGUFReader(str(model)) + except Exception as e: + logger.debug("Failed to inspect GGUF metadata for %s: %s", model, e) + else: + architecture = _gguf_reader_value(reader, "general.architecture") + if architecture == "qwen35": + text_config = hf_config.get_text_config() + is_multimodal = ( + detect_gguf_multimodal(model) is not None + or getattr(hf_config, "vision_config", None) is not None + ) + _update_config( + hf_config, + { + "model_type": "qwen3_5", + "architectures": [ + "Qwen3_5ForConditionalGeneration" + if is_multimodal + else "Qwen3_5ForCausalLM" + ], + }, + ) + if text_config is not hf_config: + _update_config( + text_config, + { + "model_type": "qwen3_5_text", + }, + ) + _update_config( + text_config, + _qwen35_text_config_updates_from_gguf(reader, "qwen35"), + ) + nextn_layers = _gguf_reader_value( + reader, "qwen35.nextn_predict_layers" + ) + block_count = _gguf_reader_value(reader, "qwen35.block_count") + if nextn_layers: + updates: dict[str, Any] = { + "mtp_num_hidden_layers": int(nextn_layers), + "num_nextn_predict_layers": int(nextn_layers), + } + if block_count is not None: + updates["num_hidden_layers"] = max( + int(block_count) - int(nextn_layers), + 0, + ) + _update_config(text_config, updates) + elif architecture == "qwen35moe": + text_config = hf_config.get_text_config() + is_multimodal = ( + detect_gguf_multimodal(model) is not None + or getattr(hf_config, "vision_config", None) is not None + ) + _update_config( + hf_config, + { + "model_type": "qwen3_5_moe", + "architectures": [ + "Qwen3_5MoeForConditionalGeneration" + if is_multimodal + else "Qwen3_5MoeForCausalLM" + ], + }, + ) + if text_config is not hf_config: + _update_config( + text_config, + { + "model_type": "qwen3_5_moe_text", + }, + ) + _update_config( + text_config, + _qwen35_text_config_updates_from_gguf(reader, "qwen35moe"), + ) + nextn_layers = _gguf_reader_value( + reader, "qwen35moe.nextn_predict_layers" + ) + block_count = _gguf_reader_value(reader, "qwen35moe.block_count") + if nextn_layers: + updates: dict[str, Any] = { + "mtp_num_hidden_layers": int(nextn_layers), + "num_nextn_predict_layers": int(nextn_layers), + } + if block_count is not None: + updates["num_hidden_layers"] = max( + int(block_count) - int(nextn_layers), + 0, + ) + _update_config(text_config, updates) + elif architecture == "gemma4-assistant": + prefix = "gemma4-assistant" + text_config = hf_config.get_text_config() + nextn_layers = _gguf_reader_value( + reader, f"{prefix}.nextn_predict_layers" + ) + block_count = _gguf_reader_value(reader, f"{prefix}.block_count") + head_count_kv = _gguf_reader_value( + reader, f"{prefix}.attention.head_count_kv" + ) + layer_types = None + sliding_pattern = _gguf_reader_value( + reader, f"{prefix}.attention.sliding_window_pattern" + ) + if sliding_pattern: + layer_types = [ + "sliding_attention" if is_sliding else "full_attention" + for is_sliding in sliding_pattern + ] + rope_theta = _gguf_reader_value(reader, f"{prefix}.rope.freq_base") + rope_theta_swa = _gguf_reader_value( + reader, f"{prefix}.rope.freq_base_swa" + ) + + _update_config( + hf_config, + { + "model_type": "gemma4_assistant", + "architectures": ["Gemma4MTPModel"], + "backbone_hidden_size": _gguf_reader_value( + reader, f"{prefix}.embedding_length_out" + ), + "n_predict": 1, + }, + ) + _update_config( + text_config, + { + "model_type": "gemma4_assistant", + "hidden_size": _gguf_reader_value( + reader, f"{prefix}.embedding_length" + ), + "backbone_hidden_size": _gguf_reader_value( + reader, f"{prefix}.embedding_length_out" + ), + "num_hidden_layers": block_count, + "num_nextn_predict_layers": nextn_layers, + "mtp_num_hidden_layers": nextn_layers, + "intermediate_size": _gguf_reader_value( + reader, f"{prefix}.feed_forward_length" + ), + "num_attention_heads": _gguf_reader_value( + reader, f"{prefix}.attention.head_count" + ), + "num_key_value_heads": _gguf_sequence_edge( + head_count_kv, + first=True, + ), + "num_global_key_value_heads": _gguf_sequence_edge( + head_count_kv, + first=False, + ), + "head_dim": _gguf_reader_value( + reader, f"{prefix}.attention.key_length_swa" + ), + "global_head_dim": _gguf_reader_value( + reader, f"{prefix}.attention.key_length" + ), + "sliding_window": _gguf_reader_value( + reader, f"{prefix}.attention.sliding_window" + ), + "layer_types": layer_types, + "attention_k_eq_v": True, + "attention_bias": False, + "num_kv_shared_layers": 0, + "rope_local_base_freq": rope_theta_swa, + "rope_parameters": { + "sliding_attention": { + "rope_type": "default", + "rope_theta": rope_theta_swa, + }, + "full_attention": { + "rope_type": "default", + "rope_theta": rope_theta, + }, + }, + }, + ) + # Patch multimodal config if mmproj.gguf exists mmproj_path = detect_gguf_multimodal(model) if mmproj_path is not None: diff --git a/vllm_gguf_plugin/loader.py b/vllm_gguf_plugin/loader.py index 64dc9f38..65a08edd 100644 --- a/vllm_gguf_plugin/loader.py +++ b/vllm_gguf_plugin/loader.py @@ -8,6 +8,7 @@ from huggingface_hub import hf_hub_download from vllm.config import ModelConfig, VllmConfig from vllm.config.load import LoadConfig +from vllm.config.utils import replace from vllm.logger import init_logger from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.utils import ( @@ -84,7 +85,7 @@ def load_model( ) -> nn.Module: device_config = vllm_config.device_config adapter = self._prepare_adapter(model_config) - vllm_config.model_config.hf_config = model_config.hf_config + vllm_config = replace(vllm_config, model_config=model_config) logger.debug( "GGUF unquantized modules: %s", adapter.load_spec.unquantized_modules ) @@ -96,7 +97,10 @@ def load_model( target_device = torch.device(device_config.device) with set_default_torch_dtype(model_config.dtype): with target_device: - model = initialize_model(vllm_config=vllm_config, prefix=prefix) + model = initialize_model( + vllm_config=vllm_config, + prefix=prefix, + ) model.load_weights( adapter.prepare_weights(model_config), ) diff --git a/vllm_gguf_plugin/ops.py b/vllm_gguf_plugin/ops.py index e90a8297..970b73a1 100644 --- a/vllm_gguf_plugin/ops.py +++ b/vllm_gguf_plugin/ops.py @@ -43,6 +43,26 @@ def _ggml_mul_mat_a8_fake( ) -> torch.Tensor: return torch.empty((X.size(0), row), dtype=X.dtype, device=W.device) + if hasattr(torch.ops._C_gguf, "ggml_mul_mat_a8_q4_0_mmq_v2"): + + @register_fake("_C_gguf::ggml_mul_mat_a8_q4_0_mmq_v2") + def _ggml_mul_mat_a8_q4_0_mmq_v2_fake( + W: torch.Tensor, + X: torch.Tensor, + row: torch.SymInt, + ) -> torch.Tensor: + return torch.empty((X.size(0), row), dtype=X.dtype, device=W.device) + + if hasattr(torch.ops._C_gguf, "ggml_mul_mat_a8_iq4_xs_mmq_v2"): + + @register_fake("_C_gguf::ggml_mul_mat_a8_iq4_xs_mmq_v2") + def _ggml_mul_mat_a8_iq4_xs_mmq_v2_fake( + W: torch.Tensor, + X: torch.Tensor, + row: torch.SymInt, + ) -> torch.Tensor: + return torch.empty((X.size(0), row), dtype=X.dtype, device=W.device) + @register_fake("_C_gguf::ggml_moe_a8") def _ggml_moe_a8_fake( X: torch.Tensor, @@ -55,9 +75,22 @@ def _ggml_moe_a8_fake( top_k: torch.SymInt, tokens: torch.SymInt, ) -> torch.Tensor: - return torch.empty( - (X.size(0) * top_k, row), dtype=torch.float16, device=W.device - ) + return torch.empty((X.size(0) * top_k, row), dtype=X.dtype, device=W.device) + + if hasattr(torch.ops._C_gguf, "ggml_moe_a8_iq4_xs_mmq_v2"): + + @register_fake("_C_gguf::ggml_moe_a8_iq4_xs_mmq_v2") + def _ggml_moe_a8_iq4_xs_mmq_v2_fake( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + row: torch.SymInt, + top_k: torch.SymInt, + tokens: torch.SymInt, + ) -> torch.Tensor: + return torch.empty((tokens * top_k, row), dtype=X.dtype, device=W.device) if hasattr(torch.ops, "_C_gguf") and hasattr(torch.ops._C_gguf, "ggml_moe_a8_vec"): @@ -99,6 +132,22 @@ def ggml_mul_mat_a8( return torch.ops._C_gguf.ggml_mul_mat_a8(W, X, quant_type, row) +def ggml_mul_mat_a8_q4_0_mmq_v2( + W: torch.Tensor, + X: torch.Tensor, + row: int, +) -> torch.Tensor: + return torch.ops._C_gguf.ggml_mul_mat_a8_q4_0_mmq_v2(W, X, row) + + +def ggml_mul_mat_a8_iq4_xs_mmq_v2( + W: torch.Tensor, + X: torch.Tensor, + row: int, +) -> torch.Tensor: + return torch.ops._C_gguf.ggml_mul_mat_a8_iq4_xs_mmq_v2(W, X, row) + + def ggml_moe_a8( X: torch.Tensor, W: torch.Tensor, @@ -141,5 +190,27 @@ def ggml_moe_get_block_size(quant_type: int) -> int: return torch.ops._C_gguf.ggml_moe_get_block_size(quant_type) +def ggml_moe_a8_iq4_xs_mmq_v2( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + row: int, + top_k: int, + tokens: int, +) -> torch.Tensor: + return torch.ops._C_gguf.ggml_moe_a8_iq4_xs_mmq_v2( + X, + W, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + row, + top_k, + tokens, + ) + + def moe_sum(input: torch.Tensor, output: torch.Tensor) -> None: torch.ops._moe_C.moe_sum(input, output) diff --git a/vllm_gguf_plugin/plugin.py b/vllm_gguf_plugin/plugin.py index 005b1403..dc74e4a9 100644 --- a/vllm_gguf_plugin/plugin.py +++ b/vllm_gguf_plugin/plugin.py @@ -1,15 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 +import sys from functools import wraps from pathlib import Path +import huggingface_hub +import vllm.config.model as model_config_module import vllm.engine.arg_utils as arg_utils_module +import vllm.model_executor.layers.quantization as quantization_module import vllm.transformers_utils.config as config_module +from transformers import PretrainedConfig +from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME from vllm.config.load import LoadConfig from vllm.engine.arg_utils import EngineArgs from vllm.model_executor.layers.quantization import ( - QUANTIZATION_METHODS, - get_quantization_config, register_quantization_config, ) from vllm.model_executor.model_loader import ( @@ -18,9 +22,20 @@ register_model_loader, ) from vllm.transformers_utils.config import get_config_parser, register_config_parser +from vllm.transformers_utils.repo_utils import file_or_path_exists +from vllm.transformers_utils.utils import without_trust_remote_code from .config_parser import GGUFConfigParser -from .gguf_utils import check_gguf_file, is_gguf, is_remote_gguf, split_remote_gguf +from .gguf_tokenizer_builder import build_tokenizer_from_gguf +from .gguf_utils import ( + check_gguf_file, + get_gguf_file_path_from_hf, + is_gguf, + is_remote_gguf, + maybe_patch_hf_config_from_gguf, + resolve_gguf_config_source, + split_remote_gguf, +) from .loader import GGUFModelLoader from .quantization import GGUFConfig @@ -34,21 +49,41 @@ def _is_gguf_reference(model: str | None) -> bool: return model.endswith(".gguf") or is_remote_gguf(model) or is_gguf(model) -def _get_gguf_config_source( - model: str, - tokenizer: str | None, - hf_config_path: str | None, -) -> str: - if hf_config_path is not None: - return hf_config_path - if tokenizer is not None and not _is_gguf_reference(tokenizer): - return tokenizer +def _uses_gguf_derived_config_source( + model: str | None, + revision: str | None = None, +) -> bool: + if not _is_gguf_reference(model): + return False + + if check_gguf_file(model): + gguf_repo = Path(model).parent + resolved_source = resolve_gguf_config_source(model, revision=revision) + if resolved_source != gguf_repo: + return True + return not file_or_path_exists(gguf_repo, HF_CONFIG_NAME, revision=revision) + if is_remote_gguf(model): repo_id, _ = split_remote_gguf(model) - return repo_id - if check_gguf_file(model): - return str(Path(model).parent) - return model + resolved_source = resolve_gguf_config_source(model, revision=revision) + if resolved_source != repo_id: + return True + return not file_or_path_exists(repo_id, HF_CONFIG_NAME, revision=revision) + + return False + + +def _get_gguf_config_probe_model(engine_args: EngineArgs) -> str | None: + if _is_gguf_reference(engine_args.model): + return engine_args.model + + speculative_config = getattr(engine_args, "speculative_config", None) + if isinstance(speculative_config, dict): + speculative_model = speculative_config.get("model") + if isinstance(speculative_model, str) and _is_gguf_reference(speculative_model): + return speculative_model + + return None def _patch_engine_args() -> None: @@ -59,8 +94,30 @@ def _patch_engine_args() -> None: @wraps(original_create_model_config) def create_model_config(self, *args, **kwargs): + if ( + self.trust_remote_code + and self.hf_config_path is None + and (gguf_config_model := _get_gguf_config_probe_model(self)) is not None + and _uses_gguf_derived_config_source( + gguf_config_model, + revision=self.revision, + ) + ): + config_module.logger.warning_once( + "Disabling `trust_remote_code` because model config was " + "selected from a GGUF-derived config source. Pass an " + "explicit `--hf-config-path` to opt in for that repository.", + ) + self.trust_remote_code = False + if _is_gguf_reference(self.model): gguf_model = self.model + if ( + self.tokenizer is None + and check_gguf_file(gguf_model) + and (tokenizer_path := build_tokenizer_from_gguf(gguf_model)) + ): + self.tokenizer = tokenizer_path if self.quantization is None: self.quantization = "gguf" if self.load_format == "auto": @@ -71,17 +128,23 @@ def create_model_config(self, *args, **kwargs): self.model_weights = gguf_model if self.served_model_name is None: self.served_model_name = [gguf_model] - self.model = _get_gguf_config_source( - gguf_model, - self.tokenizer if isinstance(self.tokenizer, str) else None, - self.hf_config_path, - ) return original_create_model_config(self, *args, **kwargs) EngineArgs.create_model_config = create_model_config EngineArgs._gguf_create_model_config_patched = True +def _patch_gguf_config_helpers() -> None: + if getattr(model_config_module, "_gguf_config_helpers_patched", False): + return + + # ModelConfig imports this helper by value, so update that binding too. + model_config_module.maybe_patch_hf_config_from_gguf = ( + maybe_patch_hf_config_from_gguf + ) + model_config_module._gguf_config_helpers_patched = True + + def _patch_speculator_probe() -> None: if getattr(arg_utils_module, "_gguf_speculator_probe_patched", False): return @@ -89,10 +152,82 @@ def _patch_speculator_probe() -> None: original_maybe_override = arg_utils_module.maybe_override_with_speculators @wraps(original_maybe_override) - def maybe_override_with_speculators(model, tokenizer, *args, **kwargs): - if _is_gguf_reference(model): - return model, tokenizer, kwargs.get("vllm_speculative_config") - return original_maybe_override(model, tokenizer, *args, **kwargs) + def maybe_override_with_speculators( + model, + tokenizer, + trust_remote_code, + revision=None, + hf_config_path=None, + vllm_speculative_config=None, + hf_token=None, + **kwargs, + ): + if not _is_gguf_reference(model): + return original_maybe_override( + model=model, + tokenizer=tokenizer, + trust_remote_code=trust_remote_code, + revision=revision, + hf_config_path=hf_config_path, + vllm_speculative_config=vllm_speculative_config, + hf_token=hf_token, + **kwargs, + ) + + if check_gguf_file(model): + if hf_config_path is None: + gguf_repo = Path(model).parent + gguf_model_repo = resolve_gguf_config_source( + model, + revision=revision, + ) + if gguf_model_repo != gguf_repo: + revision = None + elif not file_or_path_exists( + gguf_repo, + HF_CONFIG_NAME, + revision=revision, + ): + kwargs["gguf_file"] = Path(model).name + else: + gguf_model_repo = Path(model).parent + elif is_remote_gguf(model): + repo_id, quant_type = split_remote_gguf(model) + gguf_model_repo = resolve_gguf_config_source(model, revision=revision) + if gguf_model_repo != repo_id: + revision = None + elif not file_or_path_exists(repo_id, HF_CONFIG_NAME, revision=revision): + kwargs["gguf_file"] = get_gguf_file_path_from_hf( + repo_id, + quant_type, + revision=revision, + ) + else: + return model, tokenizer, vllm_speculative_config + + kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE + config_source = hf_config_path or gguf_model_repo + config_dict, _ = PretrainedConfig.get_config_dict( + config_source, + revision=revision, + token=hf_token, + **without_trust_remote_code(kwargs), + ) + speculators_config = config_dict.get("speculators_config") + if speculators_config is None: + return model, tokenizer, vllm_speculative_config + + from vllm.transformers_utils.configs.speculators.base import ( + SpeculatorsConfig, + ) + + speculative_config = SpeculatorsConfig.extract_vllm_speculative_config( + config_dict=config_dict + ) + speculative_config["model"] = model + + verifier_model = speculators_config["verifier"]["name_or_path"] + return verifier_model, verifier_model, speculative_config arg_utils_module.maybe_override_with_speculators = maybe_override_with_speculators config_module.maybe_override_with_speculators = maybe_override_with_speculators @@ -100,13 +235,36 @@ def maybe_override_with_speculators(model, tokenizer, *args, **kwargs): config_module._gguf_speculator_probe_patched = True +def _patch_quantization_config_lookup() -> None: + if getattr(quantization_module, "_gguf_config_lookup_patched", False): + return + + original_get_quantization_config = quantization_module.get_quantization_config + + @wraps(original_get_quantization_config) + def get_quantization_config(quantization: str): + if quantization == "gguf": + return GGUFConfig + return original_get_quantization_config(quantization) + + quantization_module.get_quantization_config = get_quantization_config + + for module_name in ("vllm.model_executor.model_loader.weight_utils",): + module = sys.modules.get(module_name) + if ( + module is not None + and getattr(module, "get_quantization_config", None) + is original_get_quantization_config + ): + module.get_quantization_config = get_quantization_config + + quantization_module._gguf_config_lookup_patched = True + + def register() -> None: """Register the out-of-tree GGUF integration.""" - if ( - "gguf" not in QUANTIZATION_METHODS - or get_quantization_config("gguf") is not GGUFConfig - ): - register_quantization_config("gguf")(GGUFConfig) + register_quantization_config("gguf")(GGUFConfig) + _patch_quantization_config_lookup() if "gguf" not in _LOAD_FORMAT_TO_MODEL_LOADER or not isinstance( get_model_loader(LoadConfig(load_format="gguf")), GGUFModelLoader @@ -119,5 +277,6 @@ def register() -> None: parser = None if not isinstance(parser, GGUFConfigParser): register_config_parser("gguf")(GGUFConfigParser) + _patch_gguf_config_helpers() _patch_engine_args() _patch_speculator_probe() diff --git a/vllm_gguf_plugin/quantization/config.py b/vllm_gguf_plugin/quantization/config.py index 16812032..db87646e 100644 --- a/vllm_gguf_plugin/quantization/config.py +++ b/vllm_gguf_plugin/quantization/config.py @@ -63,9 +63,10 @@ def from_config(cls, config: dict[str, Any]) -> "GGUFConfig": @classmethod def override_quantization_method( - cls, hf_quant_cfg: dict[str, Any], user_quant: str | None + cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config: Any = None ) -> "QuantizationMethods | None": del hf_quant_cfg + del hf_config if user_quant == "gguf": return "gguf" return None diff --git a/vllm_gguf_plugin/quantization/fused_moe.py b/vllm_gguf_plugin/quantization/fused_moe.py index fb717c92..8fe7486c 100644 --- a/vllm_gguf_plugin/quantization/fused_moe.py +++ b/vllm_gguf_plugin/quantization/fused_moe.py @@ -4,6 +4,7 @@ from functools import partial import torch +from gguf import GGMLQuantizationType as WeightType from vllm.model_executor.layers.fused_moe import ( RoutedExperts, ) @@ -28,7 +29,9 @@ _gguf_moe_weight_loader, _gguf_moe_weight_type_loader, ) -from .utils import MMQ_QUANT_TYPES, MMVQ_QUANT_TYPES, logger +from .utils import GGUF_PLUGIN_TORCH_LIB, MMQ_QUANT_TYPES, MMVQ_QUANT_TYPES, logger + +IQ4_XS_MOE_MMQ_V2_BLOCK_SIZE = 8 def _fused_moe_gguf( @@ -53,7 +56,108 @@ def act(inp: torch.Tensor): from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size out_hidden_states = torch.empty_like(x) + + def _slow_moe_fallback() -> None: + from . import fused_mul_mat_gguf as fused_mul_mat_gguf_op + + logger.warning_once( + "There is no support for fast MoE kernel " + "for current quantization method. " + "Falling back to slow implementation. " + ) + for tok, (w, idx) in enumerate(zip(topk_weights, topk_ids)): + inp = x[tok].reshape((1,) + x.shape[1:]) + current_hidden_state = None + for ww, ii in zip(w, idx): + out = fused_mul_mat_gguf_op(inp, w1[ii], qweight_type) + out = act(out) + current_state = fused_mul_mat_gguf_op(out, w2[ii], qweight_type2).mul_( + ww + ) + if current_hidden_state is None: + current_hidden_state = current_state + else: + current_hidden_state.add_(current_state) + out_hidden_states[tok] = current_hidden_state + + def _fused_moe_gguf_batched_leg( + leg_x: torch.Tensor, + leg_weight: torch.Tensor, + leg_type: int, + leg_row: int, + leg_top_k: int, + leg_tokens: int, + num_experts: int, + ) -> torch.Tensor | None: + if leg_type == WeightType.IQ4_XS: + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, IQ4_XS_MOE_MMQ_V2_BLOCK_SIZE, num_experts + ) + return ops.ggml_moe_a8_iq4_xs_mmq_v2( + leg_x, + leg_weight, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + leg_row, + leg_top_k, + leg_tokens, + ) + if leg_type in MMQ_QUANT_TYPES: + block_size = ops.ggml_moe_get_block_size(leg_type) + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, block_size, num_experts + ) + return ops.ggml_moe_a8( + leg_x, + leg_weight, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + leg_type, + leg_row, + leg_top_k, + leg_tokens, + ) + if leg_type in MMVQ_QUANT_TYPES: + return ops.ggml_moe_a8_vec( + leg_x, leg_weight, topk_ids, leg_top_k, leg_type, leg_row, leg_tokens + ) + return None + + has_iq4_xs_leg = ( + qweight_type == WeightType.IQ4_XS or qweight_type2 == WeightType.IQ4_XS + ) if ( + has_iq4_xs_leg + and qweight_type in MMVQ_QUANT_TYPES + and qweight_type2 in MMVQ_QUANT_TYPES + and x.shape[0] > 64 + ): + num_tokens, _ = x.shape + E, N, _ = w1.shape + top_k = topk_ids.shape[1] + + out = _fused_moe_gguf_batched_leg(x, w1, qweight_type, N, top_k, num_tokens, E) + if out is not None: + out = act(out) + out = _fused_moe_gguf_batched_leg( + out, + w2, + qweight_type2, + w2.shape[1], + 1, + num_tokens * top_k, + E, + ) + if out is not None: + out = out.reshape(num_tokens, top_k, w2.shape[1]).mul_( + topk_weights.view(num_tokens, top_k, 1) + ) + ops.moe_sum(out, out_hidden_states) + else: + _slow_moe_fallback() + elif ( qweight_type2 in MMQ_QUANT_TYPES and qweight_type in MMQ_QUANT_TYPES and x.shape[0] > 64 @@ -109,27 +213,7 @@ def act(inp: torch.Tensor): ) ops.moe_sum(out, out_hidden_states) else: - from . import fused_mul_mat_gguf as fused_mul_mat_gguf_op - - logger.warning_once( - "There is no support for fast MoE kernel " - "for current quantization method. " - "Falling back to slow implementation. " - ) - for tok, (w, idx) in enumerate(zip(topk_weights, topk_ids)): - inp = x[tok].reshape((1,) + x.shape[1:]) - current_hidden_state = None - for ww, ii in zip(w, idx): - out = fused_mul_mat_gguf_op(inp, w1[ii], qweight_type) - out = act(out) - current_state = fused_mul_mat_gguf_op(out, w2[ii], qweight_type2).mul_( - ww - ) - if current_hidden_state is None: - current_hidden_state = current_state - else: - current_hidden_state.add_(current_state) - out_hidden_states[tok] = current_hidden_state + _slow_moe_fallback() return out_hidden_states @@ -152,8 +236,9 @@ def _fused_moe_gguf_fake( op_name="_fused_moe_gguf", op_func=_fused_moe_gguf, fake_impl=_fused_moe_gguf_fake, + target_lib=GGUF_PLUGIN_TORCH_LIB, ) - fused_moe_gguf = torch.ops.vllm._fused_moe_gguf + fused_moe_gguf = torch.ops.vllm_gguf_plugin._fused_moe_gguf except AttributeError as error: raise error diff --git a/vllm_gguf_plugin/quantization/linear.py b/vllm_gguf_plugin/quantization/linear.py index cedbab44..d8a3c6b0 100644 --- a/vllm_gguf_plugin/quantization/linear.py +++ b/vllm_gguf_plugin/quantization/linear.py @@ -24,6 +24,7 @@ ) from .utils import ( DEQUANT_TYPES, + GGUF_PLUGIN_TORCH_LIB, IMATRIX_QUANT_TYPES, MMQ_QUANT_TYPES, MMVQ_QUANT_TYPES, @@ -44,6 +45,8 @@ def _fused_mul_mat_gguf( return x @ qweight.T if x.shape[0] <= mmvq_safe and qweight_type in MMVQ_QUANT_TYPES: y = ops.ggml_mul_mat_vec_a8(qweight, x, qweight_type, qweight.shape[0]) + elif qweight_type == WeightType.IQ4_XS: + y = ops.ggml_mul_mat_a8_iq4_xs_mmq_v2(qweight, x, qweight.shape[0]) elif qweight_type in MMQ_QUANT_TYPES: y = ops.ggml_mul_mat_a8(qweight, x, qweight_type, qweight.shape[0]) elif qweight_type in DEQUANT_TYPES: @@ -70,8 +73,9 @@ def _fused_mul_mat_gguf_fake( op_name="_fused_mul_mat_gguf", op_func=_fused_mul_mat_gguf, fake_impl=_fused_mul_mat_gguf_fake, + target_lib=GGUF_PLUGIN_TORCH_LIB, ) - fused_mul_mat_gguf = torch.ops.vllm._fused_mul_mat_gguf + fused_mul_mat_gguf = torch.ops.vllm_gguf_plugin._fused_mul_mat_gguf except AttributeError as error: raise error @@ -143,7 +147,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module): raise ValueError( f"Unsupported GGUF quantization type {qweight_type} in layer {layer}." ) - self._create_padded_weight_param(layer) + self._create_multi_shard_weight_param(layer) def _materialize_gguf_parameters(self, layer: torch.nn.Module) -> None: self._materialize_qweight(layer) @@ -155,8 +159,8 @@ def _materialize_qweight(self, layer: torch.nn.Module) -> None: def _materialize_qweight_type(self, layer: torch.nn.Module) -> None: _materialize_gguf_weight_type_parameter(layer, "qweight_type") - def _create_padded_weight_param(self, layer: torch.nn.Module): - """Create padded weight parameter for GGUF MergedLinear layer.""" + def _create_multi_shard_weight_param(self, layer: torch.nn.Module): + """Keep GGUF MergedLinear shards separate to avoid GPU peak memory spikes.""" qweight = layer.qweight shard_id_map = qweight.shard_id_map shard_id = qweight.shard_id @@ -166,43 +170,38 @@ def _create_padded_weight_param(self, layer: torch.nn.Module): f"Data container has mixed dtypes: {dtype}" ) dtype = next(iter(dtype)) - padded_side = max(x.size(1) for x in data_container) - concat_side = sum(x.size(0) for x in data_container) - padded_data = torch.zeros( - (concat_side, padded_side), dtype=dtype, device=qweight.device - ) + target_device = qweight.device shard_offset_map = dict[str, tuple[int, int, int]]() ordered_shard_ids = _gguf_ordered_shard_ids(shard_id) + ordered_data_container = [] + ordered_shard_id_map = dict[int | str, int]() current_offset = 0 for idx in ordered_shard_ids: id_in_container = shard_id_map[idx] + ordered_data_container.append(data_container[id_in_container]) + ordered_shard_id_map[idx] = len(ordered_data_container) - 1 start = current_offset end = start + data_container[id_in_container].size(0) size = data_container[id_in_container].size(1) - padded_data[start:end, :size] = data_container[id_in_container] shard_offset_map[idx] = (start, end, size) current_offset = end - padded_param = GGUFWeightParameter( - data=padded_data, + qweight.data_container.clear() + qweight.shard_id.clear() + qweight.shard_id_map.clear() + sharded_param = GGUFWeightParameter( + data=torch.empty(0, dtype=dtype, device=target_device), weight_loader=qweight.weight_loader, input_dim=qweight.input_dim, output_dim=qweight.output_dim, tensor_shape=qweight.tensor_shape, ) - padded_param.data_container = [] - padded_param.shard_id = ordered_shard_ids - padded_param.shard_id_map = dict(qweight.shard_id_map) + sharded_param.data_container = ordered_data_container + sharded_param.shard_id = ordered_shard_ids + sharded_param.shard_id_map = ordered_shard_id_map if hasattr(qweight, "ignore_warning"): - padded_param.ignore_warning = qweight.ignore_warning - set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map}) - qweight.data_container.clear() - qweight.shard_id.clear() - qweight.shard_id_map.clear() - if qweight.data.numel() > 0: - qweight.data = torch.empty( - 0, dtype=qweight.dtype, device=qweight.device - ) - layer.register_parameter("qweight", padded_param) + sharded_param.ignore_warning = qweight.ignore_warning + set_weight_attrs(sharded_param, {"shard_offset_map": shard_offset_map}) + layer.register_parameter("qweight", sharded_param) def apply( self, @@ -217,6 +216,19 @@ def apply( shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id qweight = layer.qweight fallback_wtype = layer.qweight_type.weight_type + if qweight.data_container: + result = [] + for idx in shard_id: + qweight_type = layer.qweight_type.shard_weight_type.get( + idx, fallback_wtype + ) + shard = qweight.data_container[qweight.shard_id_map[idx]] + result.append(fused_mul_mat_gguf_op(x, shard, qweight_type)) + out = torch.cat(result, axis=1) + if bias is not None: + out.add_(bias) + return out + shard_weight_types = [ layer.qweight_type.shard_weight_type.get(idx, fallback_wtype) for idx in shard_id diff --git a/vllm_gguf_plugin/quantization/params.py b/vllm_gguf_plugin/quantization/params.py index 1ea8b97a..12e77b93 100644 --- a/vllm_gguf_plugin/quantization/params.py +++ b/vllm_gguf_plugin/quantization/params.py @@ -18,15 +18,85 @@ def _clone_loaded_weight(loaded_weight: torch.Tensor) -> torch.Tensor: return loaded_weight.detach().clone() +def _split_gguf_tuple_shard_weight( + layer: torch.nn.Module, + param, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...], +) -> tuple[torch.Tensor, ...] | None: + output_dim = getattr(param, "output_dim", None) + output_sizes = getattr(layer, "output_sizes", None) + if output_dim is None or output_sizes is None: + return None + + shard_sizes = [output_sizes[idx] for idx in loaded_shard_id] + packed_dim = getattr(param, "packed_dim", None) + packed_factor = getattr(param, "packed_factor", 1) + if packed_dim == output_dim: + if any(size % packed_factor != 0 for size in shard_sizes): + return None + shard_sizes = [size // packed_factor for size in shard_sizes] + + if loaded_weight.size(output_dim) != sum(shard_sizes): + return None + return torch.split(loaded_weight, shard_sizes, dim=output_dim) + + +def _call_gguf_base_loader( + base_loader, + param, + loaded_weight: torch.Tensor, + loaded_shard_id=None, +): + if loaded_shard_id is None: + base_loader(param, loaded_weight) + else: + base_loader(param, loaded_weight, loaded_shard_id) + + def _resolve_gguf_weight_loader( layer: torch.nn.Module, fallback_weight_loader=None, ): - return ( + base_loader = ( layer.weight_loader_v2 if hasattr(layer, "weight_loader_v2") else fallback_weight_loader ) + if base_loader is None: + return fallback_weight_loader + + def _gguf_weight_loader_v2(param, loaded_weight, loaded_shard_id=None): + if not isinstance(loaded_shard_id, tuple): + _call_gguf_base_loader( + base_loader, param, loaded_weight, loaded_shard_id + ) + return + + if hasattr(layer, "validate_shard_id"): + layer.validate_shard_id(loaded_shard_id) + + if hasattr(param, "shard_weight_type"): + for shard_id in loaded_shard_id: + _call_gguf_base_loader(base_loader, param, loaded_weight, shard_id) + return + + if hasattr(param, "data_container"): + weight_shards = _split_gguf_tuple_shard_weight( + layer, param, loaded_weight, loaded_shard_id + ) + if weight_shards is not None: + for shard_id, weight_shard in zip( + loaded_shard_id, weight_shards, strict=True + ): + _call_gguf_base_loader( + base_loader, param, weight_shard, shard_id + ) + return + + _call_gguf_base_loader(base_loader, param, loaded_weight, loaded_shard_id) + + return _gguf_weight_loader_v2 def _resolve_gguf_weight_type_loader( @@ -178,11 +248,15 @@ def _gguf_moe_weight_loader( base_weight_loader, param: Parameter | UninitializedParameter, loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, + weight_name: str | None = None, + shard_id: str | None = None, + expert_id: int | None = None, return_success: bool = False, ) -> bool | None: + if shard_id is None: + _store_gguf_loaded_weight(param, loaded_weight) + return True if return_success else None + _materialize_gguf_moe_param(layer, param, loaded_weight, shard_id) return base_weight_loader( param, @@ -197,9 +271,9 @@ def _gguf_moe_weight_loader( def _gguf_moe_weight_type_loader( param: Parameter | UninitializedParameter, loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, + weight_name: str | None = None, + shard_id: str | None = None, + expert_id: int | None = None, return_success: bool = False, ) -> bool | None: del weight_name, expert_id diff --git a/vllm_gguf_plugin/quantization/utils.py b/vllm_gguf_plugin/quantization/utils.py index d6dfa6c1..db24dc55 100644 --- a/vllm_gguf_plugin/quantization/utils.py +++ b/vllm_gguf_plugin/quantization/utils.py @@ -5,10 +5,13 @@ from types import MappingProxyType from gguf import GGMLQuantizationType as WeightType +from torch.library import Library from vllm.logger import init_logger logger = init_logger(__name__) +GGUF_PLUGIN_TORCH_LIB = Library("vllm_gguf_plugin", "FRAGMENT") + def is_layer_skipped_gguf( prefix: str, diff --git a/vllm_gguf_plugin/quantization/vocal_embeds.py b/vllm_gguf_plugin/quantization/vocal_embeds.py index 0bdcdedd..7c6ee6da 100644 --- a/vllm_gguf_plugin/quantization/vocal_embeds.py +++ b/vllm_gguf_plugin/quantization/vocal_embeds.py @@ -20,7 +20,7 @@ _materialize_gguf_weight_parameter, _materialize_gguf_weight_type_parameter, ) -from .utils import DEQUANT_TYPES, UNQUANTIZED_TYPES +from .utils import DEQUANT_TYPES, GGUF_PLUGIN_TORCH_LIB, UNQUANTIZED_TYPES def _apply_gguf_embedding( @@ -61,8 +61,9 @@ def _apply_gguf_embedding_fake( op_name="_apply_gguf_embedding", op_func=_apply_gguf_embedding, fake_impl=_apply_gguf_embedding_fake, + target_lib=GGUF_PLUGIN_TORCH_LIB, ) - apply_gguf_embedding = torch.ops.vllm._apply_gguf_embedding + apply_gguf_embedding = torch.ops.vllm_gguf_plugin._apply_gguf_embedding except AttributeError as error: raise error diff --git a/vllm_gguf_plugin/weights_adapter/__init__.py b/vllm_gguf_plugin/weights_adapter/__init__.py index 2207bc4a..9a4f4c0c 100644 --- a/vllm_gguf_plugin/weights_adapter/__init__.py +++ b/vllm_gguf_plugin/weights_adapter/__init__.py @@ -4,9 +4,13 @@ from .base import BaseGGUFWeightsAdapter from .default import GGUFWeightsAdapter from .gemma3 import Gemma3GGUFAdapter +from .gemma4 import Gemma4GGUFAdapter +from .qwen3_5 import Qwen3_5GGUFAdapter _ADAPTER_REGISTRY: list[type[GGUFWeightsAdapter]] = [ Gemma3GGUFAdapter, + Gemma4GGUFAdapter, + Qwen3_5GGUFAdapter, ] @@ -22,5 +26,7 @@ def get_weights_adapter(config) -> GGUFWeightsAdapter: "BaseGGUFWeightsAdapter", "GGUFWeightsAdapter", "Gemma3GGUFAdapter", + "Gemma4GGUFAdapter", + "Qwen3_5GGUFAdapter", "get_weights_adapter", ] diff --git a/vllm_gguf_plugin/weights_adapter/default.py b/vllm_gguf_plugin/weights_adapter/default.py index 9ab049ba..284fc36c 100644 --- a/vllm_gguf_plugin/weights_adapter/default.py +++ b/vllm_gguf_plugin/weights_adapter/default.py @@ -11,10 +11,10 @@ import gguf import regex import torch -from transformers import AutoModelForCausalLM +from transformers import AutoModelForCausalLM, AutoModelForImageTextToText from vllm.logger import init_logger -from ..gguf_utils import maybe_patch_hf_config_from_gguf +from ..gguf_utils import detect_gguf_multimodal, maybe_patch_hf_config_from_gguf from ..weight_utils import ( get_gguf_extra_tensor_names, get_gguf_weight_type_map, @@ -28,6 +28,323 @@ logger = init_logger(__name__) +_GGUF_MODEL_TYPE_ALIASES = { + "gemma4": "gemma3", +} + + +def _gguf_arch_model_type(model_type: str) -> str: + return _GGUF_MODEL_TYPE_ALIASES.get(model_type, model_type) + + +def _gguf_name_with_suffix(base_name: str, suffix: str) -> str: + return f"{base_name}.{suffix}" if suffix else base_name + + +def _get_vision_num_layers(config: PretrainedConfig) -> int: + vision_num_layers = getattr(config.vision_config, "num_hidden_layers", None) + if vision_num_layers is None: + vision_num_layers = config.vision_config.depth + return vision_num_layers + + +def _is_multimodal_config(config: PretrainedConfig) -> bool: + return hasattr(config, "vision_config") and config.vision_config is not None + + +def _uses_multimodal_weight_layout( + config: PretrainedConfig, + architectures: list[str] | None = None, +) -> bool: + if not _is_multimodal_config(config): + return False + + if architectures is None: + architectures = getattr(config, "architectures", None) + if not architectures: + return True + + return any("ConditionalGeneration" in arch for arch in architectures) + + +def _is_gemma4_mtp_config(config: PretrainedConfig) -> bool: + return config.model_type in ("gemma4_assistant", "gemma4_mtp") + + +def _get_mtp_num_layers(text_config: PretrainedConfig) -> int: + return int( + getattr( + text_config, + "mtp_num_hidden_layers", + getattr(text_config, "num_nextn_predict_layers", 0), + ) + or 0 + ) + + +def _add_gemma4_mtp_gguf_mappings( + config: PretrainedConfig, + gguf_to_hf_name_map: dict[str, str], +) -> None: + text_config = config.get_text_config() + num_layers = int(getattr(text_config, "num_hidden_layers", 0) or 0) + + gguf_to_hf_name_map.update( + { + "token_embd.weight": "model.embed_tokens.weight", + "output_norm.weight": "model.norm.weight", + "nextn.pre_projection.weight": "model.pre_projection.weight", + "nextn.post_projection.weight": "model.post_projection.weight", + } + ) + + # Gemma4 MTP attention is Q-only and reuses the target model KV cache, + # so GGUF assistant blocks intentionally do not map attn_k/attn_v. + for idx in range(num_layers): + layer_prefix = f"model.layers.{idx}" + gguf_to_hf_name_map.update( + { + f"blk.{idx}.attn_norm.weight": ( + f"{layer_prefix}.input_layernorm.weight" + ), + f"blk.{idx}.attn_q.weight": (f"{layer_prefix}.self_attn.q_proj.weight"), + f"blk.{idx}.attn_output.weight": ( + f"{layer_prefix}.self_attn.o_proj.weight" + ), + f"blk.{idx}.attn_q_norm.weight": ( + f"{layer_prefix}.self_attn.q_norm.weight" + ), + f"blk.{idx}.post_attention_norm.weight": ( + f"{layer_prefix}.post_attention_layernorm.weight" + ), + f"blk.{idx}.ffn_norm.weight": ( + f"{layer_prefix}.pre_feedforward_layernorm.weight" + ), + f"blk.{idx}.post_ffw_norm.weight": ( + f"{layer_prefix}.post_feedforward_layernorm.weight" + ), + f"blk.{idx}.ffn_gate.weight": (f"{layer_prefix}.mlp.gate_proj.weight"), + f"blk.{idx}.ffn_up.weight": f"{layer_prefix}.mlp.up_proj.weight", + f"blk.{idx}.ffn_down.weight": (f"{layer_prefix}.mlp.down_proj.weight"), + f"blk.{idx}.layer_output_scale.weight": ( + f"{layer_prefix}.layer_scalar" + ), + } + ) + + +def _add_qwen3_5_mtp_gguf_mappings( + config: PretrainedConfig, + gguf_to_hf_name_map: dict[str, str], + sideload_params: list[re.Pattern], +) -> None: + text_config = config.get_text_config() + num_mtp_layers = _get_mtp_num_layers(text_config) + if num_mtp_layers <= 0: + return + + base_layer = int(getattr(text_config, "num_hidden_layers", 0) or 0) + shared_nextn_gguf_idx = base_layer + gguf_to_hf_name_map.update( + { + f"blk.{shared_nextn_gguf_idx}.nextn.eh_proj.weight": "mtp.fc.weight", + f"blk.{shared_nextn_gguf_idx}.nextn.enorm.weight": ( + "mtp.pre_fc_norm_embedding.weight" + ), + f"blk.{shared_nextn_gguf_idx}.nextn.hnorm.weight": ( + "mtp.pre_fc_norm_hidden.weight" + ), + f"blk.{shared_nextn_gguf_idx}.nextn.shared_head_norm.weight": ( + "mtp.norm.weight" + ), + f"blk.{shared_nextn_gguf_idx}.nextn.embed_tokens.weight": ( + "mtp.embed_tokens.weight" + ), + } + ) + for mtp_idx in range(num_mtp_layers): + gguf_idx = base_layer + mtp_idx + layer_prefix = f"mtp.layers.{mtp_idx}" + gguf_to_hf_name_map.update( + { + f"blk.{gguf_idx}.attn_norm.weight": ( + f"{layer_prefix}.input_layernorm.weight" + ), + f"blk.{gguf_idx}.post_attention_norm.weight": ( + f"{layer_prefix}.post_attention_layernorm.weight" + ), + f"blk.{gguf_idx}.attn_q.weight": ( + f"{layer_prefix}.self_attn.q_proj.weight" + ), + f"blk.{gguf_idx}.attn_k.weight": ( + f"{layer_prefix}.self_attn.k_proj.weight" + ), + f"blk.{gguf_idx}.attn_v.weight": ( + f"{layer_prefix}.self_attn.v_proj.weight" + ), + f"blk.{gguf_idx}.attn_output.weight": ( + f"{layer_prefix}.self_attn.o_proj.weight" + ), + f"blk.{gguf_idx}.attn_q_norm.weight": ( + f"{layer_prefix}.self_attn.q_norm.weight" + ), + f"blk.{gguf_idx}.attn_k_norm.weight": ( + f"{layer_prefix}.self_attn.k_norm.weight" + ), + f"blk.{gguf_idx}.ffn_gate.weight": ( + f"{layer_prefix}.mlp.gate_proj.weight" + ), + f"blk.{gguf_idx}.ffn_up.weight": ( + f"{layer_prefix}.mlp.up_proj.weight" + ), + f"blk.{gguf_idx}.ffn_down.weight": ( + f"{layer_prefix}.mlp.down_proj.weight" + ), + f"blk.{gguf_idx}.ffn_gate_inp.weight": ( + f"{layer_prefix}.mlp.gate.weight" + ), + f"blk.{gguf_idx}.ffn_gate_inp_shexp.weight": ( + f"{layer_prefix}.mlp.shared_expert_gate.weight" + ), + f"blk.{gguf_idx}.ffn_gate_shexp.weight": ( + f"{layer_prefix}.mlp.shared_expert.gate_proj.weight" + ), + f"blk.{gguf_idx}.ffn_up_shexp.weight": ( + f"{layer_prefix}.mlp.shared_expert.up_proj.weight" + ), + f"blk.{gguf_idx}.ffn_down_shexp.weight": ( + f"{layer_prefix}.mlp.shared_expert.down_proj.weight" + ), + f"blk.{gguf_idx}.ffn_gate_exps.weight": ( + f"{layer_prefix}.mlp.experts.0.gate_proj.weight" + ), + f"blk.{gguf_idx}.ffn_up_exps.weight": ( + f"{layer_prefix}.mlp.experts.0.up_proj.weight" + ), + f"blk.{gguf_idx}.ffn_down_exps.weight": ( + f"{layer_prefix}.mlp.experts.0.down_proj.weight" + ), + } + ) + sideload_params.append( + regex.compile( + f"mtp\\.layers\\.{mtp_idx}" + r"\.mlp\.experts\.[0-9]+\.(gate|up|down)_proj\.weight" + ) + ) + + +def _add_gemma4_gguf_mappings( + config: PretrainedConfig, + gguf_to_hf_name_map: dict[str, str], + sideload_params: list[re.Pattern], +) -> None: + text_config = config.get_text_config() + uses_multimodal_layout = _uses_multimodal_weight_layout(config) + layer_prefix_base = ( + "model.language_model.layers" if uses_multimodal_layout else "model.layers" + ) + for idx in range(text_config.num_hidden_layers): + layer_prefix = f"{layer_prefix_base}.{idx}" + gguf_to_hf_name_map[f"blk.{idx}.layer_output_scale.weight"] = ( + f"{layer_prefix}.layer_scalar" + ) + gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_inp.scale"] = ( + f"{layer_prefix}.router.scale" + ) + gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_inp.weight"] = ( + f"{layer_prefix}.router.proj.weight" + ) + gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.scale"] = ( + f"{layer_prefix}.router.per_expert_scale" + ) + gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_up_exps.weight"] = ( + f"{layer_prefix}.experts.gate_up_proj.weight" + ) + gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.weight"] = ( + f"{layer_prefix}.experts.down_proj.weight" + ) + gguf_to_hf_name_map[f"blk.{idx}.post_ffw_norm_1.weight"] = ( + f"{layer_prefix}.post_feedforward_layernorm_1.weight" + ) + gguf_to_hf_name_map[f"blk.{idx}.post_ffw_norm_2.weight"] = ( + f"{layer_prefix}.post_feedforward_layernorm_2.weight" + ) + gguf_to_hf_name_map[f"blk.{idx}.pre_ffw_norm_2.weight"] = ( + f"{layer_prefix}.pre_feedforward_layernorm_2.weight" + ) + sideload_params.extend( + [ + regex.compile( + f"{re.escape(layer_prefix_base)}\\.{idx}" + r"\.experts\.(gate_up_proj|down_proj)(\.weight)?" + ), + ] + ) + + if uses_multimodal_layout: + gguf_to_hf_name_map.update( + { + "v.std_bias": "model.vision_tower.std_bias", + "v.std_scale": "model.vision_tower.std_scale", + "v.patch_embd.weight": ( + "model.vision_tower.patch_embedder.input_proj.weight" + ), + "v.position_embd.weight": ( + "model.vision_tower.patch_embedder.position_embedding_table" + ), + "mm.input_projection.weight": ( + "model.embed_vision.embedding_projection.weight" + ), + } + ) + + for idx in range(_get_vision_num_layers(config)): + vision_prefix = f"model.vision_tower.encoder.layers.{idx}" + gguf_to_hf_name_map.update( + { + f"v.blk.{idx}.attn_q.weight": ( + f"{vision_prefix}.self_attn.q_proj.linear.weight" + ), + f"v.blk.{idx}.attn_k.weight": ( + f"{vision_prefix}.self_attn.k_proj.linear.weight" + ), + f"v.blk.{idx}.attn_v.weight": ( + f"{vision_prefix}.self_attn.v_proj.linear.weight" + ), + f"v.blk.{idx}.attn_out.weight": ( + f"{vision_prefix}.self_attn.o_proj.linear.weight" + ), + f"v.blk.{idx}.attn_q_norm.weight": ( + f"{vision_prefix}.self_attn.q_norm.weight" + ), + f"v.blk.{idx}.attn_k_norm.weight": ( + f"{vision_prefix}.self_attn.k_norm.weight" + ), + f"v.blk.{idx}.ffn_gate.weight": ( + f"{vision_prefix}.mlp.gate_proj.linear.weight" + ), + f"v.blk.{idx}.ffn_up.weight": ( + f"{vision_prefix}.mlp.up_proj.linear.weight" + ), + f"v.blk.{idx}.ffn_down.weight": ( + f"{vision_prefix}.mlp.down_proj.linear.weight" + ), + f"v.blk.{idx}.ln1.weight": ( + f"{vision_prefix}.input_layernorm.weight" + ), + f"v.blk.{idx}.attn_post_norm.weight": ( + f"{vision_prefix}.post_attention_layernorm.weight" + ), + f"v.blk.{idx}.ln2.weight": ( + f"{vision_prefix}.pre_feedforward_layernorm.weight" + ), + f"v.blk.{idx}.ffn_post_norm.weight": ( + f"{vision_prefix}.post_feedforward_layernorm.weight" + ), + } + ) + class GGUFWeightsAdapter(BaseGGUFWeightsAdapter): """Default adapter for GGUF models.""" @@ -46,17 +363,22 @@ def build_name_map(self, model_config: ModelConfig) -> dict[str, str]: config = model_config.hf_config text_config = config.get_text_config() model_type = config.model_type - is_multimodal = ( - hasattr(config, "vision_config") and config.vision_config is not None - ) + is_multimodal = _uses_multimodal_weight_layout(config) + orig_model_type = model_type gguf_to_hf_name_map: dict[str, str] = {} sideload_params: list[re.Pattern] = [] + if _is_gemma4_mtp_config(config): + _add_gemma4_mtp_gguf_mappings(config, gguf_to_hf_name_map) + return gguf_to_hf_name_map + if model_type == "cohere": model_type = "command-r" if model_type == "gemma3_text": model_type = "gemma3" + if model_type == "gemma4": + _add_gemma4_gguf_mappings(config, gguf_to_hf_name_map, sideload_params) if model_type in ("deepseek_v3", "deepseek_v2"): model_type = "deepseek2" for idx in range(config.num_hidden_layers): @@ -78,24 +400,62 @@ def build_name_map(self, model_config: ModelConfig) -> dict[str, str]: r"\.mlp\.experts\.[0-9]+\.(gate|up|down)_proj\.weight" ) ) - if model_type in ("qwen2_moe", "qwen3_moe"): + if model_type == "qwen3_5": + model_type = "qwen35" + if model_type in ("qwen2_moe", "qwen3_moe", "qwen3_5_moe"): model_type = model_type.replace("_", "") - for idx in range(config.num_hidden_layers): + if is_multimodal and model_type == "qwen35moe": + layer_prefix = "model.language_model.layers" + else: + layer_prefix = "model.layers" + for idx in range(text_config.num_hidden_layers): gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.weight"] = ( - f"model.layers.{idx}.mlp.experts.0.down_proj.weight" + f"{layer_prefix}.{idx}.mlp.experts.0.down_proj.weight" ) gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_exps.weight"] = ( - f"model.layers.{idx}.mlp.experts.0.gate_proj.weight" + f"{layer_prefix}.{idx}.mlp.experts.0.gate_proj.weight" ) gguf_to_hf_name_map[f"blk.{idx}.ffn_up_exps.weight"] = ( - f"model.layers.{idx}.mlp.experts.0.up_proj.weight" + f"{layer_prefix}.{idx}.mlp.experts.0.up_proj.weight" ) sideload_params.append( regex.compile( - f"model\\.layers\\.{idx}" + f"{re.escape(layer_prefix)}\\.{idx}" r"\.mlp\.experts\.[0-9]+\.(gate|up|down)_proj\.weight" ) ) + if orig_model_type in ("qwen3_5", "qwen3_5_moe"): + layer_prefix = ( + "model.language_model.layers" if is_multimodal else "model.layers" + ) + layer_types = getattr(text_config, "layer_types", []) + for idx, layer_type in enumerate(layer_types): + if layer_type == "linear_attention": + gguf_to_hf_name_map[f"blk.{idx}.ssm_dt.bias"] = ( + f"{layer_prefix}.{idx}.linear_attn.dt_bias" + ) + _add_qwen3_5_mtp_gguf_mappings(config, gguf_to_hf_name_map, sideload_params) + if orig_model_type in ("qwen3_5", "qwen3_5_moe") and is_multimodal: + gguf_to_hf_name_map.update( + { + "token_embd.weight": ( + "model.language_model.embed_tokens.weight" + ), + "v.patch_embd.weight.1": ("model.visual.patch_embed.proj.weight.1"), + "v.post_ln.weight": "model.visual.merger.norm.weight", + "v.post_ln.bias": "model.visual.merger.norm.bias", + "mm.0.weight": "model.visual.merger.linear_fc1.weight", + "mm.0.bias": "model.visual.merger.linear_fc1.bias", + "mm.2.weight": "model.visual.merger.linear_fc2.weight", + "mm.2.bias": "model.visual.merger.linear_fc2.bias", + } + ) + sideload_params.extend( + [ + regex.compile(r"model\.visual\.merger\.norm\.weight"), + regex.compile(r"model\.visual\.merger\.norm\.bias"), + ] + ) if model_type == "olmoe": for idx in range(config.num_hidden_layers): gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.weight"] = ( @@ -143,7 +503,7 @@ def build_name_map(self, model_config: ModelConfig) -> dict[str, str]: arch = None for key, value in gguf.MODEL_ARCH_NAMES.items(): - if value == model_type: + if value == _gguf_arch_model_type(model_type): arch = key break if arch is None: @@ -154,14 +514,18 @@ def build_name_map(self, model_config: ModelConfig) -> dict[str, str]: if is_multimodal: mm_proj_arch = gguf.MODEL_ARCH.MMPROJ vision_name_map = gguf.get_tensor_name_map( - mm_proj_arch, config.vision_config.num_hidden_layers + mm_proj_arch, _get_vision_num_layers(config) ) else: vision_name_map = None with torch.device("meta"): - dummy_model = AutoModelForCausalLM.from_config( - config, trust_remote_code=model_config.trust_remote_code + auto_cls = ( + AutoModelForImageTextToText if is_multimodal else AutoModelForCausalLM + ) + auto_config = config if is_multimodal else text_config + dummy_model = auto_cls.from_config( + auto_config, trust_remote_code=model_config.trust_remote_code ) state_dict = dummy_model.state_dict() @@ -206,7 +570,7 @@ def find_hf_name_in_tensor_map(hf_name: str) -> str | None: gguf_name = text_name_map.get_name(base_name) if gguf_name is None: return None - return gguf_name + "." + suffix + return _gguf_name_with_suffix(gguf_name, suffix) unmapped_params = [] for hf_name in state_dict: @@ -255,17 +619,34 @@ def _get_all_gguf_files(model_path: str) -> list[str]: logger.info("Discovered %d GGUF shard files", len(files)) return files if files else [model_path] - def update_tie_word_embeddings( + def _get_weight_sources( self, model_path: str, hf_config: PretrainedConfig, + use_multimodal_weight_layout: bool | None = None, + ) -> list[str]: + gguf_files = self._get_all_gguf_files(model_path) + if use_multimodal_weight_layout is None: + use_multimodal_weight_layout = _uses_multimodal_weight_layout(hf_config) + if use_multimodal_weight_layout: + mm_proj_path = detect_gguf_multimodal(model_path) + if mm_proj_path is not None: + mm_proj_file = os.fspath(mm_proj_path) + if mm_proj_file not in gguf_files: + gguf_files.append(mm_proj_file) + return gguf_files + + def update_tie_word_embeddings( + self, + gguf_files: list[str], + hf_config: PretrainedConfig, gguf_to_hf_name_map: dict[str, str], ) -> None: if "lm_head.weight" not in gguf_to_hf_name_map.values(): return all_extra_names = [] - for gguf_file in self._get_all_gguf_files(model_path): + for gguf_file in gguf_files: all_extra_names.extend( get_gguf_extra_tensor_names(gguf_file, gguf_to_hf_name_map) ) @@ -273,11 +654,11 @@ def update_tie_word_embeddings( def get_weight_type_map( self, - model_path: str, + gguf_files: list[str], gguf_to_hf_name_map: dict[str, str], ) -> dict[str, str]: weight_type_map = {} - for gguf_file in self._get_all_gguf_files(model_path): + for gguf_file in gguf_files: weight_type_map.update( get_gguf_weight_type_map(gguf_file, gguf_to_hf_name_map) ) @@ -300,12 +681,20 @@ def prepare_loading( model_path, model_config.hf_config ) gguf_to_hf_name_map = self.build_name_map(model_config) + use_multimodal_weight_layout = _uses_multimodal_weight_layout( + model_config.hf_config + ) + gguf_files = self._get_weight_sources( + model_path, + model_config.hf_config, + use_multimodal_weight_layout, + ) self.update_tie_word_embeddings( - model_path, model_config.hf_config, gguf_to_hf_name_map + gguf_files, model_config.hf_config, gguf_to_hf_name_map ) - weight_type_map = self.get_weight_type_map(model_path, gguf_to_hf_name_map) + weight_type_map = self.get_weight_type_map(gguf_files, gguf_to_hf_name_map) self.load_spec = GGUFLoadSpec( - weights_source=self._get_all_gguf_files(model_path), + weights_source=gguf_files, gguf_to_hf_name_map=gguf_to_hf_name_map, unquantized_modules=self.get_unquantized_modules(weight_type_map), ) diff --git a/vllm_gguf_plugin/weights_adapter/gemma4.py b/vllm_gguf_plugin/weights_adapter/gemma4.py new file mode 100644 index 00000000..42c8992c --- /dev/null +++ b/vllm_gguf_plugin/weights_adapter/gemma4.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Iterable + +import torch + +from .default import GGUFWeightsAdapter + + +def _flatten_gemma4_patch_embed_weight(weight: torch.Tensor) -> torch.Tensor: + return weight.flatten(1) if weight.dim() == 4 else weight + + +def _transform_gemma4_weight_name(name: str) -> str: + replacements = { + ".experts.gate_up_proj.qweight_type": ( + ".moe.experts.routed_experts.w13_qweight_type" + ), + ".experts.gate_up_proj.qweight": ( + ".moe.experts.routed_experts.w13_qweight" + ), + ".experts.down_proj.qweight_type": ( + ".moe.experts.routed_experts.w2_qweight_type" + ), + ".experts.down_proj.qweight": ( + ".moe.experts.routed_experts.w2_qweight" + ), + } + for old, new in replacements.items(): + if name.endswith(old): + return name.removesuffix(old) + new + return name + + +class Gemma4GGUFAdapter(GGUFWeightsAdapter): + """Adapter for Gemma4 GGUF models.""" + + @classmethod + def matches(cls, config) -> bool: + return config.model_type in ("gemma4", "gemma4_assistant", "gemma4_mtp") + + def map_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in weights: + name = _transform_gemma4_weight_name(name) + yield name, self.transform_weight(name, weight) + + def transform_weight( + self, + hf_name: str, + weight: torch.Tensor, + ) -> torch.Tensor: + if hf_name == "model.vision_tower.patch_embedder.input_proj.weight": + return _flatten_gemma4_patch_embed_weight(weight) + return weight diff --git a/vllm_gguf_plugin/weights_adapter/qwen3_5.py b/vllm_gguf_plugin/weights_adapter/qwen3_5.py new file mode 100644 index 00000000..65db608c --- /dev/null +++ b/vllm_gguf_plugin/weights_adapter/qwen3_5.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import gguf +import torch + +from .default import GGUFWeightsAdapter + +if TYPE_CHECKING: + from vllm.config import ModelConfig + +_QWEN3_5_PATCH_EMBED_WEIGHT = "model.visual.patch_embed.proj.weight" +_QWEN3_5_PATCH_EMBED_WEIGHT_1 = f"{_QWEN3_5_PATCH_EMBED_WEIGHT}.1" +_QWEIGHT_SUFFIX = ".qweight" +_QWEIGHT_TYPE_SUFFIX = ".qweight_type" + + +def _get_text_config(config): + if hasattr(config, "get_text_config"): + return config.get_text_config() + text_config = getattr(config, "text_config", None) + return text_config if text_config is not None else config + + +def _qwen3_5_linear_attention_dims(config) -> tuple[int, int, int, int] | None: + text_config = _get_text_config(config) + num_k_heads = getattr(text_config, "linear_num_key_heads", None) + num_v_heads = getattr(text_config, "linear_num_value_heads", None) + head_k_dim = getattr(text_config, "linear_key_head_dim", None) + head_v_dim = getattr(text_config, "linear_value_head_dim", None) + if None in (num_k_heads, num_v_heads, head_k_dim, head_v_dim): + return None + + num_k_heads = int(num_k_heads) + num_v_heads = int(num_v_heads) + head_k_dim = int(head_k_dim) + head_v_dim = int(head_v_dim) + if num_k_heads <= 0 or num_v_heads <= 0: + return None + if num_k_heads == num_v_heads: + return None + if num_v_heads % num_k_heads != 0: + return None + return num_k_heads, num_v_heads, head_k_dim, head_v_dim + + +def _grouped_to_tiled_v_heads( + tensor: torch.Tensor, + dim: int, + num_k_heads: int, + num_v_per_k: int, + head_dim: int, +) -> torch.Tensor: + shape = list(tensor.shape) + if dim < 0: + dim += len(shape) + if shape[dim] != num_k_heads * num_v_per_k * head_dim: + return tensor + + 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 _tiled_to_grouped_v_heads( + tensor: torch.Tensor, + dim: int, + num_k_heads: int, + num_v_per_k: int, + head_dim: int, +) -> torch.Tensor: + shape = list(tensor.shape) + if dim < 0: + dim += len(shape) + if shape[dim] != num_k_heads * num_v_per_k * head_dim: + return tensor + + new_shape = shape[:dim] + [num_v_per_k, num_k_heads, 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 _maybe_restore_qwen3_5_gdn_layout( + name: str, + weight: torch.Tensor, + config, +) -> torch.Tensor: + if ".linear_attn." not in name: + return weight + + dims = _qwen3_5_linear_attention_dims(config) + if dims is None: + return weight + num_k_heads, num_v_heads, head_k_dim, head_v_dim = dims + num_v_per_k = num_v_heads // num_k_heads + key_dim = num_k_heads * head_k_dim + value_dim = num_v_heads * head_v_dim + + if ".linear_attn.in_proj_qkv" in name: + if weight.shape[0] != 2 * key_dim + value_dim: + return weight + q = weight[:key_dim] + k = weight[key_dim : 2 * key_dim] + v = weight[2 * key_dim :] + v = _tiled_to_grouped_v_heads(v, 0, num_k_heads, num_v_per_k, head_v_dim) + return torch.cat((q, k, v), dim=0) + + if ".linear_attn.in_proj_z" in name: + return _tiled_to_grouped_v_heads( + weight, 0, num_k_heads, num_v_per_k, head_v_dim + ) + + if ".linear_attn.in_proj_a" in name or ".linear_attn.in_proj_b" in name: + return _tiled_to_grouped_v_heads(weight, 0, num_k_heads, num_v_per_k, 1) + + if name.endswith(".linear_attn.dt_bias") or ".linear_attn.dt_proj" in name: + if weight.dim() == 1: + return _tiled_to_grouped_v_heads( + weight.unsqueeze(-1), 0, num_k_heads, num_v_per_k, 1 + ).squeeze(-1) + return _tiled_to_grouped_v_heads(weight, -1, num_k_heads, num_v_per_k, 1) + + if name.endswith(".linear_attn.A_log"): + if torch.any(weight >= 0): + raise ValueError( + "Qwen3.5 GGUF A_log tensor is expected to store negative " + "-exp(A_log) values" + ) + restored = torch.log(-weight.to(torch.float32)) + if restored.dim() == 1: + return _tiled_to_grouped_v_heads( + restored.unsqueeze(-1), 0, num_k_heads, num_v_per_k, 1 + ).squeeze(-1) + return _tiled_to_grouped_v_heads(restored, -1, num_k_heads, num_v_per_k, 1) + + if ".linear_attn.conv1d" in name: + conv_weight = weight.squeeze(1) if weight.dim() == 3 else weight + if conv_weight.shape[0] != 2 * key_dim + value_dim: + return conv_weight[:, None, :] if conv_weight.dim() == 2 else conv_weight + qk_part = conv_weight[: 2 * key_dim] + v_part = conv_weight[2 * key_dim :] + v_part = _tiled_to_grouped_v_heads( + v_part, 0, num_k_heads, num_v_per_k, head_v_dim + ) + return torch.cat((qk_part, v_part), dim=0)[:, None, :] + + if ".linear_attn.out_proj" in name: + return _tiled_to_grouped_v_heads( + weight, 1, num_k_heads, num_v_per_k, head_v_dim + ) + + return weight + + +def _maybe_reshape_qwen3_5_gguf_weight( + name: str, + weight: torch.Tensor, +) -> torch.Tensor: + if "mlp.shared_expert_gate" in name and weight.dim() == 1: + return weight[None, :] + if "linear_attn.conv1d.weight" in name and weight.dim() == 2: + return weight[:, None, :] + return weight + + +def _maybe_restore_qwen3_5_norm_weight( + name: str, + weight: torch.Tensor, +) -> torch.Tensor: + if not name.endswith("norm.weight") or name.endswith("linear_attn.norm.weight"): + return weight + text_norm_prefixes = ( + "model.layers.", + "model.language_model.layers.", + ) + text_norm_names = ( + "model.norm.weight", + "model.language_model.norm.weight", + ) + if name in text_norm_names or name.startswith(text_norm_prefixes): + # llama.cpp stores Qwen3.5 text RMSNorm weights as weight + 1. + return weight - 1 + return weight + + +def _dequantize_gguf_weight( + weight: torch.Tensor, + qweight_type: gguf.GGMLQuantizationType, +) -> torch.Tensor: + dense = gguf.quants.dequantize(weight.detach().cpu().numpy(), qweight_type) + return torch.from_numpy(dense.copy()) + + +class Qwen3_5GGUFAdapter(GGUFWeightsAdapter): + """Adapter for Qwen3.5 GGUF models.""" + + def __init__(self, config) -> None: + super().__init__(config) + self._forced_dequantized_modules: set[str] = set() + self._qweight_types: dict[str, gguf.GGMLQuantizationType] = {} + + @classmethod + def matches(cls, config) -> bool: + return config.model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_mtp") + + def prepare_loading( + self, + model_path: str, + model_config: ModelConfig, + ): + load_spec = super().prepare_loading(model_path, model_config) + self._forced_dequantized_modules.clear() + if _qwen3_5_linear_attention_dims(self.config) is None: + return load_spec + if load_spec.gguf_to_hf_name_map is None: + return load_spec + + for hf_name in load_spec.gguf_to_hf_name_map.values(): + if hf_name.endswith(".linear_attn.out_proj.weight"): + module_name = hf_name.removesuffix(".weight") + self._forced_dequantized_modules.add(module_name) + if module_name not in load_spec.unquantized_modules: + load_spec.unquantized_modules.append(module_name) + return load_spec + + def map_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + self._qweight_types.clear() + patch_weight: torch.Tensor | None = None + patch_weight_1: torch.Tensor | None = None + + for hf_name, weight in weights: + if hf_name.endswith(_QWEIGHT_TYPE_SUFFIX): + module_name = hf_name.removesuffix(_QWEIGHT_TYPE_SUFFIX) + qweight_type = gguf.GGMLQuantizationType(int(weight.item())) + self._qweight_types[module_name] = qweight_type + if module_name in self._forced_dequantized_modules: + continue + yield hf_name, weight + continue + + if hf_name.endswith(_QWEIGHT_SUFFIX): + module_name = hf_name.removesuffix(_QWEIGHT_SUFFIX) + if module_name in self._forced_dequantized_modules: + qweight_type = self._qweight_types.get(module_name) + if qweight_type is None: + raise ValueError( + "Missing GGUF qweight_type for forced dense tensor " + f"{hf_name}" + ) + hf_name = f"{module_name}.weight" + weight = _dequantize_gguf_weight(weight, qweight_type) + + if hf_name == _QWEN3_5_PATCH_EMBED_WEIGHT: + patch_weight = weight + continue + if hf_name == _QWEN3_5_PATCH_EMBED_WEIGHT_1: + patch_weight_1 = weight + continue + yield hf_name, self.transform_weight(hf_name, weight) + + if patch_weight is None: + if patch_weight_1 is not None: + yield _QWEN3_5_PATCH_EMBED_WEIGHT_1, patch_weight_1 + return + + if patch_weight_1 is not None: + patch_weight = torch.stack((patch_weight, patch_weight_1), dim=2) + yield _QWEN3_5_PATCH_EMBED_WEIGHT, self.transform_weight( + _QWEN3_5_PATCH_EMBED_WEIGHT, patch_weight + ) + + def transform_weight( + self, + hf_name: str, + weight: torch.Tensor, + ) -> torch.Tensor: + weight = _maybe_restore_qwen3_5_gdn_layout(hf_name, weight, self.config) + weight = _maybe_restore_qwen3_5_norm_weight(hf_name, weight) + return _maybe_reshape_qwen3_5_gguf_weight(hf_name, weight)