[Bugfix][Qwen4Exp] Derive the QSA head_dim from the unsharded head count (#59945)

Signed-off-by: lesj0610 <lesj0610@godoiksan.org>
Co-authored-by: lesj0610 <lesj0610@godoiksan.org>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>

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

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