From 6283021142bbf5ee324395dad5e80b8661400329 Mon Sep 17 00:00:00 2001 From: Pavani Majety Date: Thu, 26 Feb 2026 15:38:19 -0800 Subject: [PATCH] [Bugfix] Fix KV Scale loading for MLA Models (#35430) Signed-off-by: Pavani Majety --- vllm/model_executor/layers/quantization/modelopt.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index c0cc35b28..999bb6325 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -12,7 +12,7 @@ from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_fp8_linear_kernel, ) -from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention import Attention, MLAAttention from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -183,7 +183,7 @@ class ModelOptQuantConfigBase(QuantizationConfig): self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": # handle kv-cache first so we can focus only on weight quantization thereafter - if isinstance(layer, Attention): + if isinstance(layer, (Attention, MLAAttention)): return self.KVCacheMethodCls(self) # handle exclusion