[Bugfix][ROCm][CI] Restore the DeepSeek-V4 input GEMM override point (#51821)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Aug 13, 2026 at 03:33 UTC
b369f10d5c5dc3af29e5a2f9213a462b8faadae9
1 file changed
+8
-1
vllm/models/deepseek_v4/attention.py
+8
-1
@@ -443,6 +443,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
443
# Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
444
return self._o_proj(o, positions)
445
446
+ def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
447
+ # Override point: the ROCm layer preshuffles this weight in place, so
448
+ # it cannot go through fused_wqa_wkv directly.
449
+ # MergedColumnParallelLinear returns (output, bias); bias is None.
450
+ qr_kv, _ = self.fused_wqa_wkv(hidden_states)
451
+ return qr_kv
452
+
453
def _run_parallel_input_projections(
454
self, hidden_states: torch.Tensor
455
) -> tuple[
@@ -494,7 +501,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
501
aux_fns[2] = indexer_compressor_kv_score
502
503
qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel(
497
- lambda: self.fused_wqa_wkv(hidden_states)[0],
504
+ lambda: self._fused_wqa_wkv_gemm(hidden_states),
505
aux_fns,
506
self.ln_events[0],
507
self.ln_events[1:4],