Skip to content

RotaryEmbed drops the second half of cos/sin — wrong output for 2D / vision RoPE #6833

Description

@Acture

Summary

RotaryEmbed's non-interleaved path reads only the first embed_dim/2 values of the
cos/sin cache and reuses them for both halves of each head vector:

// src/layer/rotaryembed.cpp (naive), same in arm/x86/vulkan/loongarch/mips
for (int j = 0; j < embed_dim / 2; j++)
{
    const float cos_val = *cos_ptr++;   // only reads cos[0 .. embed_dim/2)
    const float sin_val = *sin_ptr++;
    *outptr0++ = x0 * cos_val - x1 * sin_val;   // out[j]
    *outptr1++ = x0 * sin_val + x1 * cos_val;   // out[j+half] reuses the SAME cos/sin
}

This is correct for standard 1-D RoPE, where cos = concat(freqs, freqs) and the two
halves are identical. But fuse_convert_rotaryembed lowers the general HF expression
x*cos + rotate_half(x)*sin (with a full-width cos) into this op without checking
that the two halves are equal
. For models whose halves differ — 2-D / vision RoPE, where
cos = concat(cos_h, cos_w) — the second output half is rotated with the wrong (height)
frequencies instead of the width frequencies, silently producing wrong output.

The correct general form (directly matching the fused expression) is:

out[j]      = x[j]*cos[j]        - x[j+half]*sin[j]
out[j+half] = x[j+half]*cos[j+half] + x[j]*sin[j+half]   // <-- second-half cos/sin

Found this in trying #6788.

Reproduction

Vision encoder of tencent/Penguin-VL-2B (Qwen3-0.6B-based, bidirectional, 2-D RoPE;
cos = concat(cos_h, cos_w), the two 64-wide halves differ by up to 2.0). Converted with
pnnx, run in fp32 blobs, compared against the PyTorch reference encoder output (1024 tokens
× 1024 dims):

ncnn RotaryEmbed max-diff vs PyTorch per-token cosine sim (mean)
current (reads half) 3.39 0.45
patched (reads full width) 6.2e-4 ~1.0

The 3.39 is not a precision artifact — it is deterministic (identical for fp16 and fp32
weights) and catastrophic (mean per-token cosine similarity 0.45; all 1024 tokens have >10% relative error). The PyTorch-side wrapper matches the reference encoder to 5.7e-5, so
the divergence is entirely in the op. It is masked in end-to-end tests because greedy
argmax decoding tolerates a badly wrong vision encoding for simple inputs.

Proposed fix

Read the second-half cos/sin for the second output half. Backward-compatible:

  • For standard RoPE the two halves are identical, so the result is bit-identical
    (verified: patched vs current build byte-for-byte on the same graph with identical-halves
    cos).
  • Guard on cos_cache.w: use the second half only when the cache is full width
    (w == embed_dim); when it is embed_dim/2 (as test_rotaryembed.cpp supplies today),
    fall back to the current duplicated behavior so existing half-width callers keep working.
  • Add a differing-halves test case so the path is covered.

Happy to send a PR covering all backends (naive / arm / x86 / vulkan / loongarch / mips)
plus the test. This unblocks in-graph conversion of 2-D vision RoPE (Penguin-VL, and any
Qwen-VL-style vision encoder that concatenates per-axis frequencies), which currently
forces ports to compute RoPE outside the graph.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions