docker/kyle-ct-quant.patch

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,