[EXPERIMENT][WIP][OpenVINO] Fix DeepSeek-V3 YaRN attention-scale and RoPE-interleave bugs

docs/source/openvino/models.mdx

@@ -50,7 +50,7 @@ Here is the list of the supported architectures :
 - DeiT
 - DeepSeek
 - DeepSeek-V2
-- DeepSeek-V3
+- DeepSeek-V3 (YaRN `rope_scaling` MLA attention-scale/RoPE-interleave fix for `deepseek_v3`/`deepseek_v2`)
 - DistilBERT
 - ERNIE 4.5
 - ELECTRA

optimum/exporters/openvino/model_patcher.py

@@ -62,7 +62,6 @@
     is_transformers_version,
 )
 
-
 if is_transformers_version(">=", "4.53"):
     from transformers.masking_utils import (
         ALL_MASK_ATTENTION_FUNCTIONS,
@@ -3554,6 +3553,15 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
         sin = sin[position_ids].unsqueeze(unsqueeze_dim)  # [bs, 1, seq_len, dim]
         q_fp32 = q.to(dtype=torch.float32, device=q.device)
         k_fp32 = k.to(dtype=torch.float32, device=k.device)
+        # DeepSeek-V3's apply_rotary_pos_emb (see e.g. modeling_deepseek.py) first re-interleaves the last
+        # dimension into (d // 2, 2) pairs and transposes them before applying rotate_half. This differs from
+        # the plain Llama-style rotate_half used elsewhere, and is required for correct MLA rope application
+        # (q_pe/k_pe). Omitting this step silently applies the wrong element pairing to the rotation and
+        # produces incoherent generation regardless of weight precision (FP16/INT8/INT4 alike).
+        b, h, s, d = q_fp32.shape
+        q_fp32 = q_fp32.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+        b, h, s, d = k_fp32.shape
+        k_fp32 = k_fp32.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
         q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)
         k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)
         return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)
@@ -3634,6 +3642,12 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
         dropout_p=self.attention_dropout if self.training else 0.0,
         # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
         is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        # Deepseek MLA does not use the default 1/sqrt(head_dim) attention scale: with YaRN rope_scaling
+        # (mscale_all_dim set), self.softmax_scale additionally applies a `mscale ** 2` correction factor
+        # (see modeling_deepseek.py). Omitting `scale` here silently falls back to SDPA's default scale and
+        # produces systematically wrong (uncalibrated) attention logits for any Deepseek checkpoint that uses
+        # YaRN scaling, degrading generation quality regardless of weight precision (FP16/INT8/INT4 alike).
+        scale=self.softmax_scale,
     )
 
     attn_output = attn_output.transpose(1, 2).contiguous()
@@ -3757,6 +3771,10 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
         dropout_p=self.attention_dropout if self.training else 0.0,
         # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
         is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        # See identical fix/comment in deepseek_v3_attn_forward above: SDPA's default scale omits the
+        # YaRN `mscale ** 2` correction folded into self.softmax_scale, causing systematically wrong
+        # attention logits for Deepseek checkpoints using YaRN rope_scaling.
+        scale=self.softmax_scale,
     )
     attn_output = attn_output.transpose(1, 2).contiguous()
 

tests/openvino/utils_tests.py

@@ -188,6 +188,14 @@ def _create_tiny_kokoro_model():
     "deberta": "optimum-intel-internal-testing/tiny-random-deberta",
     "deberta-v2": "optimum-intel-internal-testing/tiny-random-DebertaV2Model",
     "decilm": "optimum-intel-internal-testing/tiny-random-decilm",
+    # This fixture's config already sets rope_scaling={"type": "yarn", "mscale_all_dim": 0.707, ...},
+    # exercising the DeepseekPatcher.deepseek_v3_attn_forward YaRN softmax_scale/RoPE-interleave code
+    # path fixed in the "Fix deepseek_v3_attn_forward" commit. Note: its qk_rope_head_dim=2 makes the
+    # interleaved-pair RoPE rearrangement a mathematical no-op (a 1x2 transpose+reshape is the identity),
+    # so test_compare_to_transformers's logits-allclose assertion does not by itself regression-guard
+    # that part of the fix at this fixture size — see agent-results/optimum-intel in the OMEGA repo for
+    # a standalone isolated-math reproduction. Bumping qk_rope_head_dim/qk_nope_head_dim (e.g. to 4-8)
+    # would let this fixture catch that regression automatically in the future.
     "deepseek": "optimum-intel-internal-testing/tiny-random-deepseek-v3",
     "deit": "optimum-intel-internal-testing/tiny-random-DeiTModel",
     "convnext": "optimum-intel-internal-testing/tiny-random-convnext",