From f37d00e856266da3dff40b5adc76fad0571b76df Mon Sep 17 00:00:00 2001
From: littledgg <61149469+littledgg@users.noreply.github.com>
Date: Thu, 24 Jul 2025 20:15:00 +0800
Subject: [PATCH] [Model] Provide clearer error for missing KV cache
 quantization scales (#3007)

---
 fastdeploy/model_executor/models/deepseek_v3.py | 1 +
 fastdeploy/model_executor/models/qwen2.py       | 1 +
 fastdeploy/model_executor/models/qwen3.py       | 1 +
 3 files changed, 3 insertions(+)

diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py
index 2c2b0efe1..4d75b03b9 100644
--- a/fastdeploy/model_executor/models/deepseek_v3.py
+++ b/fastdeploy/model_executor/models/deepseek_v3.py
@@ -420,6 +420,7 @@ class DeepseekV3MLAAttention(nn.Layer):
         # NOTE(Ryan):Make sure kv_b_proj_bmm loaded before kv_b_proj,
         # The same weight key will be poped after kv_b_proj.
         self.o_proj.load_state_dict(state_dict)
+        self.mla_attn.load_state_dict(state_dict)
 
 
 class DeepSeekV3DecoderLayer(nn.Layer):
diff --git a/fastdeploy/model_executor/models/qwen2.py b/fastdeploy/model_executor/models/qwen2.py
index 418f3cdfb..af2af00b1 100644
--- a/fastdeploy/model_executor/models/qwen2.py
+++ b/fastdeploy/model_executor/models/qwen2.py
@@ -113,6 +113,7 @@ class Qwen2Attention(nn.Layer):
         """ """
         self.qkv_proj.load_state_dict(state_dict)
         self.o_proj.load_state_dict(state_dict)
+        self.attn.load_state_dict(state_dict)
 
     def forward(
         self,
diff --git a/fastdeploy/model_executor/models/qwen3.py b/fastdeploy/model_executor/models/qwen3.py
index 9dee0f8e4..4b106aea2 100644
--- a/fastdeploy/model_executor/models/qwen3.py
+++ b/fastdeploy/model_executor/models/qwen3.py
@@ -95,6 +95,7 @@ class Qwen3Attention(nn.Layer):
         self.o_proj.load_state_dict(state_dict)
         self.q_norm.load_state_dict(state_dict)
         self.k_norm.load_state_dict(state_dict)
+        self.attn.load_state_dict(state_dict)
 
     def forward(
         self,