File size: 4,614 Bytes
3a03e99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | diff --git a/vllm/model_executor/layers/deepseek_compressor.py b/vllm/model_executor/layers/deepseek_compressor.py
index 48628fec4..5f57d737c 100644
--- a/vllm/model_executor/layers/deepseek_compressor.py
+++ b/vllm/model_executor/layers/deepseek_compressor.py
@@ -221,7 +221,7 @@ class DeepseekCompressor(nn.Module):
[self.coff * self.head_dim, self.coff * self.head_dim],
bias=False,
return_bias=False,
- quant_config=None,
+ quant_config=vllm_config.quant_config,
disable_tp=True,
prefix=f"{prefix}.fused_wkv_wgate",
)
diff --git a/vllm/model_executor/layers/deepseek_v4_attention.py b/vllm/model_executor/layers/deepseek_v4_attention.py
index 93595c373..53ff43af3 100644
--- a/vllm/model_executor/layers/deepseek_v4_attention.py
+++ b/vllm/model_executor/layers/deepseek_v4_attention.py
@@ -405,6 +405,14 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
k_cache_prefix=self.mla_attn.prefix,
)
+ # TODO: For now, model requires fp8 quantization for attention
+ if hasattr(self.wo_a, "weight_scale_inv"):
+ self.wo_scale_name = "weight_scale_inv"
+ elif hasattr(self.wo_a, "weight_scale"):
+ self.wo_scale_name = "weight_scale"
+ else:
+ raise NotImplementedError("DeepSeekV4 requires FP8 attention quantization")
+
def forward(
self,
positions: torch.Tensor,
@@ -455,7 +463,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
)
wo_a_fp8 = self.wo_a.weight
- wo_a_scale = self.wo_a.weight_scale_inv
+ wo_a_scale = getattr(self.wo_a, self.wo_scale_name)
z = _allocate_deepseek_v4_wo_a_output(
num_tokens,
@@ -493,11 +501,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
compressor = self.compressor
def compressor_kv_score() -> torch.Tensor:
- return torch.mm(
- hidden_states,
- compressor.fused_wkv_wgate.weight.T,
- out_dtype=torch.float32,
- )
+ return compressor.fused_wkv_wgate(hidden_states)
aux_fns[0] = compressor_kv_score
@@ -510,11 +514,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
return weights
def indexer_compressor_kv_score() -> torch.Tensor:
- return torch.mm(
- hidden_states,
- indexer.compressor.fused_wkv_wgate.weight.T,
- out_dtype=torch.float32,
- )
+ return indexer.compressor.fused_wkv_wgate(hidden_states)
aux_fns[1] = indexer_weights_proj
aux_fns[2] = indexer_compressor_kv_score
@@ -1850,7 +1850,7 @@ class DeepseekV4Indexer(nn.Module):
hidden_size,
self.n_head,
bias=False,
- quant_config=None,
+ quant_config=quant_config,
prefix=f"{prefix}.weights_proj",
)
self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
diff --git a/vllm/model_executor/models/deepseek_v4.py b/vllm/model_executor/models/deepseek_v4.py
index 1c3adb3ac..17ba3776e 100644
--- a/vllm/model_executor/models/deepseek_v4.py
+++ b/vllm/model_executor/models/deepseek_v4.py
@@ -1021,7 +1021,14 @@ class DeepseekV4Attention(nn.Module):
prefix=f"{prefix}.wo_b",
)
self.softmax_scale = self.head_dim**-0.5
- self.scale_fmt = config.quantization_config["scale_fmt"]
+ # scale_fmt is only used in the indexer (for C4A layers), not in
+ # the main attention. Default to "ue8m0" for compatibility.
+ self.scale_fmt = (
+ config.quantization_config.get("scale_fmt", "ue8m0")
+ if hasattr(config, "quantization_config")
+ and isinstance(config.quantization_config, dict)
+ else "ue8m0"
+ )
self.rope_parameters = config.rope_scaling
@@ -1487,7 +1494,7 @@ class DeepseekV4Model(nn.Module):
first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
if first_layer.ffn.use_mega_moe:
return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
- # Params for weights, fp8 weight scales, fp8 activation scales
+ # Params for unfused moe weights
# (param_name, weight_name, expert_id, shard_id)
return FusedMoE.make_expert_params_mapping(
self,
|