From 6cf86cd95fe8d1d003f578f98545edb7c563497b Mon Sep 17 00:00:00 2001 From: MOVZX Date: Sun, 18 Oct 2026 17:18:47 +0000 Subject: [PATCH] [Bugfix] Load checkpoint KV-cache scales in the Qwen3.5 family Checkpoints that ship calibrated k_scale/v_scale tensors (e.g. unsloth/Qwen3.8-27B-NVFP4, which carries one pair per full-attention layer) name them ...self_attn.{k,v}_scale, but the runtime parameters live on the inner attention layer (...self_attn.attn.{k,v}_scale). Qwen3_5Model.load_weights hands the iterable straight to AutoWeightsLoader, which performs no scale-name remapping, so the scales are silently dropped and --kv-cache-dtype fp8 runs at scale 1.0 with only a warning_once that fires even on success (the MTP draft layer legitimately has no scales). Route scale-suffixed names through maybe_remap_kv_scale_name first, mirroring gemma4/gpt_oss. No behavior change for checkpoints without scales. Measured on RTX 5090 / Qwen3.8-27B NVFP4 / fp8 KV / MTP k=3 (GSM8K greedy, n=300): 92.3% at scale 1.0 vs 93.7% with scales loaded (bf16-KV control 94.5%). Fixes #54623 Co-Authored-By: Hinata Signed-off-by: MOVZX --- vllm/model_executor/models/qwen3_5.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 67c6fb8821..f1691ba0a5 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -61,6 +61,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.weight_utils import ( + maybe_remap_kv_scale_name, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -296,8 +299,28 @@ class Qwen3_5Model(Qwen3NextModel): n_shared_experts=1, ckpt_prefix="mlp.shared_expert", ) + # Checkpoints that ship calibrated KV-cache scales name them + # `...self_attn.{k,v}_scale`, but the runtime parameters live on the + # inner attention layer (`...self_attn.attn.{k,v}_scale`), and + # AutoWeightsLoader performs no remapping — without this the scales + # are silently dropped and fp8 KV runs at scale 1.0 (#54623). + params_dict = dict(self.named_parameters()) + + def _remap_kv_scale_names( + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in weights: + if name.endswith((".k_scale", ".v_scale", ".q_scale")): + remapped = maybe_remap_kv_scale_name(name, params_dict) + if remapped is None: + continue + name = remapped + yield name, weight + loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loader.load_weights( + _remap_kv_scale_names(weights), mapper=self.hf_to_vllm_mapper + ) class Qwen3_5ForCausalLMBase(