Fix torch.compile crash: move Blackwell path inside custom op boundary
The previous approach called _forward_blackwell() BEFORE the torch.ops.vllm.deepseek_v4_attention custom op, which broke torch.compile (dynamo can't trace the Python functions). Fix: instead of modifying forward(), modify attention_impl() which runs INSIDE the custom op boundary. Detect SM100+ and dispatch to _attention_impl_blackwell() which uses: - fused_qnorm_rope_kv_insert_py() instead of C++ kernel - full_sdpa_attention() instead of FlashMLA Removed dead _forward_blackwell method from forward().
This commit is contained in:
@@ -273,20 +273,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
|
||||
hidden_states: torch.Tensor,
|
||||
llama_4_scaling: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_tokens = hidden_states.shape[0]
|
||||
|
||||
# ── Blackwell (SM100+) path: pure PyTorch, no FlashMLA ────────
|
||||
# FlashMLA and fused CUDA kernels don't work on SM100.
|
||||
# Use our CSA/HCA attention with PyTorch SDPA instead.
|
||||
cap = current_platform.get_device_capability()
|
||||
if cap is not None and cap.major >= 10:
|
||||
return self._forward_blackwell(
|
||||
positions, hidden_states, num_tokens,
|
||||
)
|
||||
|
||||
# ── Original path (SM90 and below) ────────────────────────────
|
||||
# Pre-allocate attention output with FlashMLA-padded head count.
|
||||
# The op writes into `o_padded`; we slice to n_local_heads after.
|
||||
num_tokens = hidden_states.shape[0]
|
||||
o_padded = torch.empty(
|
||||
(num_tokens, self.padded_heads, self.head_dim),
|
||||
dtype=hidden_states.dtype,
|
||||
@@ -389,102 +378,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
|
||||
|
||||
return self.wo_b(z.flatten(1))
|
||||
|
||||
def _forward_blackwell(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
num_tokens: int,
|
||||
) -> torch.Tensor:
|
||||
"""Blackwell (SM100+) attention path using CSA/SDPA.
|
||||
|
||||
Replaces:
|
||||
- torch.ops.vllm.deepseek_v4_attention → pure PyTorch
|
||||
- fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert → pure PyTorch
|
||||
- FlashMLA sparse attention → PyTorch SDPA
|
||||
- FP8 einsum for wo_a → BF16 BMM
|
||||
"""
|
||||
from vllm.model_executor.layers.csa_attention import (
|
||||
fused_qnorm_rope_kv_insert_py,
|
||||
full_sdpa_attention,
|
||||
)
|
||||
|
||||
# 1. Run the GEMM projections (same as non-Blackwell path)
|
||||
qr_kv, kv_score, indexer_kv_score, indexer_weights = (
|
||||
self.attn_gemm_parallel_execute(hidden_states)
|
||||
)
|
||||
qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
|
||||
|
||||
# 2. Fused q/kv RMS norm (same as upstream)
|
||||
qr, kv = fused_q_kv_rmsnorm(
|
||||
qr, kv,
|
||||
self.q_norm.weight.data,
|
||||
self.kv_norm.weight.data,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
# 3. wq_b → full Q
|
||||
q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
|
||||
|
||||
# 4. RoPE on Q + KV cache insert (pure PyTorch)
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
|
||||
if isinstance(attn_metadata, dict):
|
||||
from vllm.v1.attention.backends.mla.sparse_swa import (
|
||||
DeepseekSparseSWAMetadata,
|
||||
)
|
||||
swa_metadata = cast(
|
||||
"DeepseekSparseSWAMetadata | None",
|
||||
attn_metadata.get(self.swa_cache_layer.prefix),
|
||||
)
|
||||
if swa_metadata is not None:
|
||||
swa_kv_cache = self.swa_cache_layer.kv_cache
|
||||
swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
|
||||
fused_qnorm_rope_kv_insert_py(
|
||||
q, kv, swa_kv_cache_2d,
|
||||
swa_metadata.slot_mapping,
|
||||
positions.to(torch.int64),
|
||||
self.rotary_emb.cos_sin_cache,
|
||||
self.eps,
|
||||
swa_metadata.block_size,
|
||||
nope_dim=self.nope_head_dim,
|
||||
rope_dim=self.rope_head_dim,
|
||||
)
|
||||
else:
|
||||
# Dummy run — just apply RoPE to Q
|
||||
half = self.rope_head_dim // 2
|
||||
cos_q = self.rotary_emb.cos_sin_cache[positions, :half].unsqueeze(1).to(q.dtype)
|
||||
sin_q = self.rotary_emb.cos_sin_cache[positions, half:].unsqueeze(1).to(q.dtype)
|
||||
q_rope = q[:, :, self.nope_head_dim:].clone()
|
||||
q[:, :, self.nope_head_dim:][:, :, 0::2] = q_rope[:, :, 0::2] * cos_q - q_rope[:, :, 1::2] * sin_q
|
||||
q[:, :, self.nope_head_dim:][:, :, 1::2] = q_rope[:, :, 0::2] * sin_q + q_rope[:, :, 1::2] * cos_q
|
||||
|
||||
# 5. Attention: use PyTorch SDPA (works on Blackwell)
|
||||
o = full_sdpa_attention(q, kv, self.softmax_scale)
|
||||
|
||||
# 6. wo_a: BF16 inverse RoPE + BMM (same as existing BF16 path)
|
||||
from vllm.model_executor.layers.csa_attention import apply_inv_gptj_rope
|
||||
cos_f32 = self.rotary_emb.cos_sin_cache.to(torch.float32)
|
||||
half = self.rope_head_dim // 2
|
||||
cos_o = cos_f32[positions, :half].unsqueeze(1).to(o.dtype)
|
||||
sin_o = cos_f32[positions, half:].unsqueeze(1).to(o.dtype)
|
||||
o_inv = apply_inv_gptj_rope(o, cos_o, sin_o, self.nope_head_dim)
|
||||
|
||||
heads_per_group = self.n_local_heads // self.n_local_groups
|
||||
o_inv = o_inv.view(
|
||||
num_tokens, self.n_local_groups, heads_per_group * self.head_dim
|
||||
).permute(1, 0, 2)
|
||||
wo_a_w = self.wo_a.weight.view(
|
||||
self.n_local_groups, -1, heads_per_group * self.head_dim
|
||||
)
|
||||
z = torch.bmm(o_inv, wo_a_w.transpose(1, 2))
|
||||
z = z.permute(1, 0, 2)
|
||||
if self.wo_a.gather_output and self.wo_a.tp_size > 1:
|
||||
z = tensor_model_parallel_all_gather(z)
|
||||
z = z.reshape(num_tokens, self.n_local_groups * self.o_lora_rank)
|
||||
|
||||
return self.wo_b(z)
|
||||
|
||||
def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
|
||||
aux_streams = self.aux_stream_list
|
||||
if aux_streams is not None:
|
||||
@@ -551,6 +444,15 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
|
||||
positions: torch.Tensor,
|
||||
out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place
|
||||
) -> None:
|
||||
# ── Blackwell (SM100+) path ──────────────────────────────────
|
||||
# FlashMLA and fused CUDA kernels don't work on SM100.
|
||||
# Use CSA/SDPA attention with pure PyTorch instead.
|
||||
cap = current_platform.get_device_capability()
|
||||
if cap is not None and cap.major >= 10:
|
||||
self._attention_impl_blackwell(hidden_states, positions, out)
|
||||
return
|
||||
|
||||
# ── Original path (SM90 and below) ───────────────────────────
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
|
||||
@@ -652,6 +554,101 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
|
||||
# ([num_tokens, padded_heads, head_dim]).
|
||||
self.mla_attn(q, kv, positions, output=out)
|
||||
|
||||
def _attention_impl_blackwell(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
) -> None:
|
||||
"""Blackwell (SM100+) attention: pure PyTorch, no FlashMLA.
|
||||
|
||||
Same projection flow as the original path, but replaces:
|
||||
1. torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
|
||||
→ pure PyTorch RoPE + KV cache insert
|
||||
2. self.mla_attn (FlashMLA) → PyTorch SDPA
|
||||
"""
|
||||
from vllm.model_executor.layers.csa_attention import (
|
||||
fused_qnorm_rope_kv_insert_py,
|
||||
full_sdpa_attention,
|
||||
)
|
||||
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
|
||||
qr_kv, kv_score, indexer_kv_score, indexer_weights = (
|
||||
self.attn_gemm_parallel_execute(hidden_states)
|
||||
)
|
||||
qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
|
||||
qr, kv = fused_q_kv_rmsnorm(
|
||||
qr, kv,
|
||||
self.q_norm.weight.data,
|
||||
self.kv_norm.weight.data,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
# wq_b
|
||||
q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
|
||||
|
||||
# Run compressor on default stream (it's Triton-based, should work)
|
||||
if self.compressor is not None:
|
||||
self.compressor(kv_score, positions, self.rotary_emb)
|
||||
|
||||
# Run indexer if present
|
||||
if self.indexer is not None:
|
||||
self.indexer(
|
||||
hidden_states, qr, indexer_kv_score, indexer_weights,
|
||||
positions, self.indexer_rotary_emb,
|
||||
)
|
||||
|
||||
# RoPE on Q + KV cache insert (pure PyTorch)
|
||||
if isinstance(attn_metadata, dict):
|
||||
from vllm.v1.attention.backends.mla.sparse_swa import (
|
||||
DeepseekSparseSWAMetadata,
|
||||
)
|
||||
swa_metadata = cast(
|
||||
"DeepseekSparseSWAMetadata | None",
|
||||
attn_metadata.get(self.swa_cache_layer.prefix),
|
||||
)
|
||||
if swa_metadata is not None:
|
||||
swa_kv_cache = self.swa_cache_layer.kv_cache
|
||||
swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
|
||||
fused_qnorm_rope_kv_insert_py(
|
||||
q, kv, swa_kv_cache_2d,
|
||||
swa_metadata.slot_mapping,
|
||||
positions.to(torch.int64),
|
||||
self.rotary_emb.cos_sin_cache,
|
||||
self.eps,
|
||||
swa_metadata.block_size,
|
||||
nope_dim=self.nope_head_dim,
|
||||
rope_dim=self.rope_head_dim,
|
||||
)
|
||||
else:
|
||||
self._apply_rope_q(q, positions)
|
||||
else:
|
||||
# Dummy run
|
||||
self._apply_rope_q(q, positions)
|
||||
out.zero_()
|
||||
return
|
||||
|
||||
# Attention using PyTorch SDPA (works on Blackwell)
|
||||
o = full_sdpa_attention(q, kv, self.softmax_scale)
|
||||
|
||||
# Write into the output buffer (same shape as original path)
|
||||
if self.n_local_heads < self.padded_heads:
|
||||
out[:, :self.n_local_heads, :] = o
|
||||
out[:, self.n_local_heads:, :] = 0
|
||||
else:
|
||||
out.copy_(o)
|
||||
|
||||
def _apply_rope_q(self, q, positions):
|
||||
"""Apply GPT-J RoPE to Q in-place (fallback when no SWA metadata)."""
|
||||
half = self.rope_head_dim // 2
|
||||
cos_q = self.rotary_emb.cos_sin_cache[positions, :half].unsqueeze(1).to(q.dtype)
|
||||
sin_q = self.rotary_emb.cos_sin_cache[positions, half:].unsqueeze(1).to(q.dtype)
|
||||
q_rope = q[:, :, self.nope_head_dim:].clone()
|
||||
q[:, :, self.nope_head_dim:][:, :, 0::2] = q_rope[:, :, 0::2] * cos_q - q_rope[:, :, 1::2] * sin_q
|
||||
q[:, :, self.nope_head_dim:][:, :, 1::2] = q_rope[:, :, 0::2] * sin_q + q_rope[:, :, 1::2] * cos_q
|
||||
|
||||
def _fused_qnorm_rope_kv_insert(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
|
||||
Reference in New Issue
Block a user