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

[mlukasze]

⚠️ AUTOMATICALLY GENERATED BY OMEGA AGENT — REQUIRES HUMAN REVIEW ⚠️
This PR was created by an AI agent as part of automated model enablement.
A human maintainer must review and approve it before it can be considered for merge.
Do NOT merge without human review and sign-off.

What does this PR do?

Fixes two independent correctness bugs in DeepseekPatcher.deepseek_v3_attn_forward
(optimum/exporters/openvino/model_patcher.py) that silently corrupt MLA attention for any
deepseek_v3-architecture checkpoint that uses YaRN rope_scaling with mscale_all_dim set
(e.g. real DeepSeek-V3, and derivatives such as ByteDance-Seed/academic-ds-9B).

Provenance: this fix was found and verified as part of an automated model-enablement pipeline
run for ByteDance-Seed/academic-ds-9B (tracked internally at
openvinotoolkit/omega#67 for reference — no external access needed to review this PR).
This is an experimental, agent-authored draft PR per that pipeline's standing rules; see the
banner above.

The two bugs

1. Missing YaRN softmax_scale. torch.nn.functional.scaled_dot_product_attention was
   called without an explicit scale= argument, silently falling back to SDPA's default
   1/sqrt(head_dim) instead of the model's own self.softmax_scale, which folds in
   DeepSeek-V3's YaRN mscale-squared correction. For ByteDance-Seed/academic-ds-9B's YaRN
   config (factor=40, mscale_all_dim=1.0) this under-scales attention logits by ~1.87x.
2. Missing interleaved-pair RoPE rearrangement. The local apply_rotary_pos_emb helper
   omitted the view into (d/2, 2) pairs + transpose + reshape that DeepSeek-V3's own
   modeling_deepseek.py applies before rotate_half. Without it, q_pe/k_pe rotary
   components use the wrong element pairing, corrupting MLA's positional attention component.

deepseek_v2_attn_forward (the sibling function for DeepSeek-V2-architecture models) already
had both of these correct — bug (1) is also fixed there for consistency/correctness, since it is
present verbatim in that function too. This brings deepseek_v3_attn_forward in line with its
sibling, which is strong internal-consistency evidence this was an oversight rather than an
intentional difference.

Verification

• Eager PyTorch (no OpenVINO/tracing involved): monkey-patched deepseek_v3_attn_forward
  onto the live HF AutoModelForCausalLM for ByteDance-Seed/academic-ds-9B (a real, public
  deepseek_v3-architecture checkpoint). Before the fix: patched-attention logits differed from
  the original unpatched model by up to ~11.8 max absolute difference on a single forward pass.
  After the fix: matches to ~1e-5 (float precision noise).
• OpenVINO IR: re-exported FP16/INT8/INT4 OV IR with the fix. Before the fix, the exported
  FP16 IR did not predict " Paris" anywhere in its top-10 tokens for "The capital of France is"
  (HF correctly predicts it top-1). After the fix, the OV IR matches HF exactly (top-1 " Paris",
  logit 34.15).
• Generation quality: 5 spot-check prompts (logistic regression code, capital of France,
  story completion, business/ERP passage, AI/technology passage) — all three precisions
  (FP16/INT8/INT4) produce coherent, on-topic, semantically-correct text matching the raw HF
  PyTorch reference after the fix; before the fix, all three produced garbled word-salad or
  empty output on most prompts.
• Isolated math reproduction (no model download / no deps beyond torch): a standalone script
  reproducing both bugs in isolation with realistic head dims (qk_rope_head_dim=64,
  qk_nope_head_dim=128, YaRN factor=40/mscale_all_dim=1.0 matching this checkpoint's actual
  config) confirms a ~1.87x softmax-scale ratio and a large (>8.0 max-abs) RoPE divergence. This
  script is available on request as CI-independent evidence for reviewers (kept out of the repo
  diff per this pipeline's tiny-model/test conventions — it needs no fixture or network access).

Tests

tests/openvino/test_decoder.py::OVModelForCausalLMIntegrationTest::test_compare_to_transformers_5_deepseek
already runs the "deepseek" (deepseek_v3) architecture through this exact code path — its
fixture (optimum-intel-internal-testing/tiny-random-deepseek-v3) already sets
rope_scaling={"type": "yarn", "mscale_all_dim": 0.707, "factor": 40, ...} and asserts
torch.allclose(ov_outputs.logits, transformers_outputs.logits, atol=1e-4). With this fix
applied, that assertion passes (confirmed locally). However, we found and want to be fully
transparent that this existing test does not reliably regression-guard bug (2): the fixture's
qk_rope_head_dim=2 makes the interleaved-pair rearrangement a mathematical no-op (a 1x2
transpose+reshape is the identity permutation), and reverting the fix locally
(git revert --no-commit) and re-running the same test still passes at this fixture size. We
therefore did not add a new pytest unit test tightly coupled to specific tensor shapes/seeds (per
this pipeline's testing conventions, that would be a brittle, narrow addition); instead:

• We documented this limitation directly in tests/openvino/utils_tests.py next to the
  "deepseek" fixture entry, with a recommendation to bump qk_rope_head_dim/qk_nope_head_dim
  (e.g. to 4-8) so this fixture would catch a future regression of bug (2) automatically. We do
  not have upload access to optimum-intel-internal-testing to make that fixture change
  ourselves.
• We ran the isolated-math reproduction script described above as the CI-independent regression
  check for both bugs, at realistic (non-degenerate) head dims.
• We ran test_compare_to_transformers for minicpm3 (a sibling architecture sharing this same
  model_patcher.py file, though a different, untouched attention-forward function) as a
  regression check per this pipeline's "shared file" convention — it passes, confirming no
  collateral regression from this change.
• Added a one-line note next to the DeepSeek-V3 entry in docs/source/openvino/models.mdx.

WWB accuracy (context, not blocking this specific fix)

Precision	Device	Similarity (post-fix)	Similarity (pre-fix)	Δ	vs 0.90 threshold
INT8	CPU	0.7701	0.7516	+0.0185	❌ below
INT4	CPU	0.5666	0.4419	+0.1246	❌ below
INT8	GPU	0.8388	0.8850	-0.0462	❌ below
INT4	GPU	0.3374	0.4917	-0.1543	❌ below

This fix improved CPU accuracy on both precisions but did not, by itself, bring
ByteDance-Seed/academic-ds-9B above the pipeline's 0.90 similarity gate on either device — a
separate GPU-plugin-numerics issue and a possible INT8/INT4 quantization-calibration/MoE-routing
gap remain and are being tracked/handled independently in the enablement pipeline. That said,
this fix is independently correct and valuable regardless of those remaining issues: it corrects
attention/RoPE math that was objectively wrong for any YaRN-scaled deepseek_v3 (and, for bug 1,
deepseek_v2) checkpoint, irrespective of quantization or device.

Known warnings (acknowledged per this pipeline's quality checklist)

• mutation: check_pr_quality.py reports 27 pre-existing in-place config mutations elsewhere
  in model_patcher.py (lines 1034-1064, 2480-2581) — none in the code touched by this PR.
• base_class: a recent upstream commit (e6f612a7, "Remove onnx dependency") touched
  model_configs.py; not applicable here since this PR does not add or change any
  OpenVINOConfig/base-class code.

Installation instructions

pip install git+https://github.com/mlukasze/optimum-intel.git@enable/ByteDance-Seed-academic-ds-9B
pip install --pre -U openvino openvino-tokenizers nncf --extra-index-url https://storage.openvinotoolkit.org/simple/wheels/nightly

Exporting cmd-line

optimum-cli export openvino -m ByteDance-Seed/academic-ds-9B ov_model --weight-format fp16

Inference script

import openvino_genai as ov_genai

pipe = ov_genai.LLMPipeline("ov_model", "CPU")
result = pipe.generate("The capital of France is", max_new_tokens=5)
print(result)  # After the fix: correctly predicts " Paris" (matches HF top-1, logit 34.15)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?