RotaryEmbed: support full-width cos/sin cache for 2D / vision RoPE

[Acture]

What

RotaryEmbed's non-interleaved path read only the first embed_dim/2 values of the cos/sin
cache and reused them for both halves of each head vector. That is correct for standard 1-D
RoPE (cos = concat(freqs, freqs), identical halves), but fuse_convert_rotaryembed lowers the
general expression x*cos + rotate_half(x)*sin — with a full-width cos — into this op without
checking that the halves are equal. For 2-D / vision RoPE where the two halves differ
(cos = concat(cos_h, cos_w), e.g. the Penguin-VL / Qwen-VL-style vision encoder) the second half
was rotated with the wrong frequencies, silently producing wrong output.

Fixes #6833.

Fix

Read the second half's cos/sin for the second output half. Backward-compatible via a
cache-width guard:

• full-width cache (cos_cache.w == embed_dim) → the two halves are independent (2-D rope);
• half-width cache (embed_dim/2, the standard-rope form the unit test supplies) → the second
  half reuses the first, so existing half-width callers are bit-identical.

Applied to every backend (naive / arm / x86 / loongarch / mips / vulkan) and their bf16/fp16 paths.
The vulkan shaders take the cache row width as a new push constant.

Test

tests/test_rotaryembed.cpp previously only used a half-width cache, which is trivially
identical-halves and never exercised the differing-halves path. Added a full-width
differing-halves case for the non-interleaved op.

Verification

• test_rotaryembed and test_rotaryembed_oom pass locally on the naive + ARM backends
  (macOS Apple Silicon), covering both the existing half-width and the new full-width cases.
• End-to-end: the Penguin-VL-2B vision encoder (28-layer bidirectional, 2-D RoPE) converted with
  pnnx now runs in-graph at max|diff| = 6.2e-4 vs the PyTorch reference (fp32); before the fix
  the same fused graph gave max|diff| = 3.39.
• Standard-rope regression: patched vs unpatched build produce byte-identical output on a
  fused graph fed an identical-halves cache.
• x86 / loongarch / mips / vulkan carry the same mechanical change but were not built locally
  (no toolchain on this host) — relying on CI to validate them.

[tencent-adm]

All committers have signed the CLA.

[Acture]

Thank you for your submission, we really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.You have signed the CLA already but the status is still pending? Let us recheck it.

Got Internal Server Error when I tried to sign this.

Signed.

[Acture]

@codex review

[Acture]

@dosu any concerns with this fix?

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: 5967676ea5

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[dosubot]

The fix looks sound overall. The cache-width guard (cos_cache.w == embed_dim vs embed_dim/2) is a clean way to distinguish 2D vision RoPE from standard 1D RoPE, with zero overhead in the hot loop [1]. The backward compatibility claim holds up — when the cache is half-width, cw = 0 makes the code path identical to the old behavior, producing byte-identical output.

The mechanical consistency across all 11 backend files (naive, ARM NEON/fp16, x86 AVX/SSE/bf16, LoongArch, MIPS, Vulkan) is good — same pattern everywhere reduces divergence risk [1].

A few minor observations (none blocking):

1. No dimension validation — if cos_cache.w is neither embed_dim nor embed_dim/2 (e.g. a malformed cache), the behavior is silently wrong or could read out of bounds. An assertion or early error return would add safety. In practice PNNX generates the caches so this is unlikely, but PNNX's fuse_convert_rotaryembed doesn't explicitly validate cache dimensions either [2].

2. Interleaved path unchanged — the interleaved code path still reads only embed_dim/2 cos/sin values [3]. If a future model uses interleaved 2D RoPE, it would hit the same bug. Worth a comment noting this limitation.

3. Vulkan testing gap — the Vulkan shader changes (new push constant for cache width) are structurally correct but were validated only on naive + ARM [1]. CI should cover this, but worth confirming the Vulkan compute tests pass.

4. Test coverage — the new full-width differing-halves test case covers the core fix well, and existing half-width tests provide regression coverage. A test for interleaved + full-width cache (even as a negative test) would be a nice addition.

None of these are blockers — the fix is correct, well-tested for its scope, and backward compatible.

---