From 284b6a5d576ee6c4a5d921d5b94722db6bd9b964 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 06:40:58 +0000 Subject: [PATCH] Fix attention patch: use original vllm imports, only patch forward method Previous version copied the entire file from our local vllm clone which had imports (breakable_cudagraph) missing from the Docker image's vllm. Now we start from the Docker image's original file and only patch the DeepseekV4MultiHeadLatentAttentionWrapper.forward method. --- vllm/patches/deepseek_v4_attention.py | 37 +++++++++------------------ 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/vllm/patches/deepseek_v4_attention.py b/vllm/patches/deepseek_v4_attention.py index e5337ebb..e4846da9 100644 --- a/vllm/patches/deepseek_v4_attention.py +++ b/vllm/patches/deepseek_v4_attention.py @@ -17,7 +17,6 @@ import torch.nn.functional as F from transformers import DeepseekV2Config, DeepseekV3Config import vllm.envs as envs -from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.model_executor.layers.linear import ( ReplicatedLinear, ) @@ -50,7 +49,7 @@ from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.deepseek_compressor import DeepseekCompressor -from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.input_quant_fp8 import ( QuantFP8, @@ -89,8 +88,9 @@ logger = init_logger(__name__) PREFILL_CHUNK_SIZE = 4 +@dataclass # --------------------------------------------------------------------------- -# BF16 inverse RoPE (replaces fused_inv_rope_fp8_quant + FP8 einsum) +# BF16 inverse RoPE (replaces fused_inv_rope_fp8_quant for the O projection) # --------------------------------------------------------------------------- def _apply_inv_rope_bf16( @@ -100,35 +100,23 @@ def _apply_inv_rope_bf16( nope_dim: int = 448, rope_dim: int = 64, ) -> torch.Tensor: - """Apply inverse RoPE to attention output in BF16. - - Pure-PyTorch replacement for fused_inv_rope_fp8_quant. - Only does inverse RoPE (no FP8 quant) since we use NVFP4 for wo_b. - """ + """Apply inverse RoPE to attention output in BF16.""" if rope_dim == 0 or o.numel() == 0: return o half_rope = rope_dim // 2 - cos_all = cos_sin_cache[positions, :half_rope].unsqueeze(1).to(o.dtype) sin_all = cos_sin_cache[positions, half_rope:].unsqueeze(1).to(o.dtype) - o_rope = o[:, :, nope_dim:] o_even = o_rope[:, :, 0::2] o_odd = o_rope[:, :, 1::2] - - # Inverse rotation (conjugate): - # inv[2i] = x[2i] * cos + x[2i+1] * sin - # inv[2i+1] = -x[2i] * sin + x[2i+1] * cos inv_even = o_even * cos_all + o_odd * sin_all inv_odd = -o_even * sin_all + o_odd * cos_all - result = o.clone() result[:, :, nope_dim:][:, :, 0::2] = inv_even result[:, :, nope_dim:][:, :, 1::2] = inv_odd return result -@dataclass class DeepseekV4MLAModules: """Modules used in DeepseekV4 MLA.""" @@ -347,11 +335,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): ) o = o_padded[:, : self.n_local_heads, :] - # === O Projection (patched for NVFP4 + BF16 wo_a) === - # The original path uses fused_inv_rope_fp8_quant + FP8 einsum for wo_a, - # which requires wo_a to have FP8 weights and weight_scale_inv. - # Our checkpoint has wo_a in BF16 and wo_b in NVFP4. - # We replace with: inverse RoPE (BF16) + BMM wo_a + NVFP4 wo_b. + # === O Projection (patched for BF16 wo_a + NVFP4 wo_b) === + # Original path uses fused_inv_rope_fp8_quant + FP8 einsum, which + # requires wo_a.weight_scale_inv (doesn't exist for BF16 wo_a). # Step 1: Inverse RoPE (BF16, pure PyTorch) o_inv = _apply_inv_rope_bf16( @@ -365,13 +351,11 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): # Step 2: wo_a grouped linear (BF16 BMM) # o_inv: (T, n_local_heads, head_dim) # wo_a.weight: (n_local_groups * o_lora_rank, heads_per_group * head_dim) BF16 - heads_per_group = self.n_local_heads // self.n_local_groups - hidden_dim = heads_per_group * self.head_dim + hidden_dim = self.wo_a.weight.shape[1] o_grouped = o_inv.view(num_tokens, self.n_local_groups, hidden_dim) wo_a_w = self.wo_a.weight.view( self.n_local_groups, self.o_lora_rank, hidden_dim ) - # BMM: (G, T, D) @ (G, D, R) → (G, T, R) → (T, G, R) z = torch.bmm( o_grouped.permute(1, 0, 2), wo_a_w.transpose(1, 2), @@ -379,6 +363,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): # Step 3: wo_b (NVFP4 via CuTeDSL) return self.wo_b(z.flatten(1)) + ) + + return self.wo_b(z.flatten(1)) def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: aux_streams = self.aux_stream_list @@ -584,7 +571,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): ) -@eager_break_during_capture def deepseek_v4_attention( hidden_states: torch.Tensor, positions: torch.Tensor, @@ -1141,6 +1127,7 @@ class DeepseekV4Indexer(nn.Module): quant_config=None, prefix=f"{prefix}.weights_proj", ) + self.k_norm = LayerNorm(self.head_dim, eps=1e-6) self.softmax_scale = self.head_dim**-0.5 self.scale_fmt = "ue8m0"