From d9503dc503e75193d2128908b59c0d141d4b5fb1 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Wed, 7 Oct 2026 21:24:40 +0900 Subject: [PATCH] [Bugfix][Qwen4Exp] Derive the QSA head_dim from the unsharded head count (#59945) Signed-off-by: lesj0610 Co-authored-by: lesj0610 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/amd/qsa.py | 2 +- vllm/models/qwen4_exp/nvidia/qsa.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 325e0e3d91..3f03b293b4 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -214,7 +214,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must be divisible by replicated QSA KV heads") self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index a20fcdeec9..0246db8094 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -393,7 +393,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must be divisible by replicated QSA KV heads") self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5