biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix attention patch: use original vllm imports, only patch forward method

Browse Source 
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.
This commit is contained in:
biondizzle
2026-05-19 06:40:58 +00:00
parent 199efe0871
commit 284b6a5d57
1 changed files with 12 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
vllm/patches/deepseek_v4_attention.py
View File

@@ -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"