Guancheng Du

[Guancheng Du @ducatiqwq]

(7/7) Takeaway.

Natural language is not periodic, but RoPE keeps rotating and creates ghost sub-diagonals that distract attention. Keeping only the first period removes them, improves long-context performance with or without training, and comes with practical speedups.

WavePrune: One Period Is Often Enough for RoPE
Paper link: https://t.co/Mbuy6HNNrv

[Guancheng Du @ducatiqwq]

RoPE encodes positional information by rotating query and key vectors based on their positions, with each RoPE channel contributing a cosine function of the relative position.

But is this periodicity really beneficial?

Introducing our new paper — WavePrune: One Period Is Often Enough for RoPE!

[Guancheng Du @ducatiqwq]

(6/7) An efficient kernel: FlashWavePrune

Since each channel has its own cutoff, QKᵀ is no longer a plain GEMM and maps poorly onto Tensor Cores. Our fix: split the head dimensions into groups of 16, matching the K size of one hardware MMA instruction, and give each group a single shared cutoff (the largest wavelength in it). Each group then runs as a standard MMA, followed by one mask on its output. Note that all results above use this approximation.

WavePrune lets us skip work. For each (query block, key block) tile, we compare each group’s cutoff with the tile’s range of distances |i−j|:
All distances exceed the cutoff → skip the group
All distances are within the cutoff → compute without masking

Otherwise → compute, then mask
Skipped groups save both FLOPs and memory reads.
vs. FlashAttention-2 (A100, head dim 128, rope_theta 1e6):

[Guancheng Du @ducatiqwq]

(5/7) Pretraining with WavePrune.

We pretrain GPT-2 base, medium, and large at 2K context and evaluate at 4K. Training loss is unchanged, while extrapolation loss drops substantially at every scale. Pretraining with WavePrune also outperforms applying it only at inference.

[Guancheng Du @ducatiqwq]

(4/7) Training-free gains on HELMET.

Applied to off-the-shelf models without any training, WavePrune improves 4 of 5 models (64K/128K average):
• Qwen3-8B: 35.7 → 40.0
• Qwen3-32B: 55.1 → 56.3
• Gemma3-12B: 38.2 → 38.7
• Ministral3-3B: 39.1 → 39.7

Llama3.1-8B regresses (38.2 → 35.2); we discuss this in Appendix E.

[Guancheng Du @ducatiqwq]

(3/7) The fix: keep a single period.

RoPE splits each head into two-dimensional channels. Channel c rotates by f_c per position, so one full period spans w_c = 2π/f_c positions.

WavePrune counts channel c in the attention logit only when the query–key distance stays within that period:

L_ij = (1/√d_h) · Σ_c φ_i^(c)(q_i)ᵀ φ_j^(c)(k_j) · 𝟙[|i−j| ≤ w_c]
Here φ^(c) is the standard RoPE rotation of channel c. The only thing we add is the indicator, 𝟙[|i−j| ≤ w_c].
With this single change, per-channel oscillation and ghost sub-diagonals both disappear.

[Guancheng Du @ducatiqwq]

(2/7) Periodicity causes position aliasing, hence ghost sub-diagonals.

In a RoPE channel with frequency f, a QK pair at distance d contributes ∝ cos(fd + ω), where ω is the initial angle. Because cos is 2π-periodic, distances with the same (fd mod 2π) look identical to the model.

In Qwen3-8B, high-frequency channels flip sign rapidly as d grows. After summation and softmax, these flips appear as spurious stripes in the attention map, which we call ghost sub-diagonals. We prove they are inherent to RoPE and distract attention, so training alone cannot remove them.

[Guancheng Du @ducatiqwq]

(1/7) RoPE does not need to be periodic.

Attention retrieves information by matching queries to keys, so a relevant key should be reachable at any distance. RoPE instead rotates q and k with position, so their score oscillates with distance and becomes unreliable at long range. Methods like YaRN rescale the frequencies but keep the periodic rotation.

[Guancheng Du @ducatiqwq]

(7/7) Takeaway.

Natural language is not periodic, but RoPE keeps rotating. That periodicity creates ghost sub-diagonals that distract attention. Keeping only the first period removes them, improves long-context performance with or without training, and comes with practical speedups.

[Guancheng Du @ducatiqwq]

(6/7) An efficient kernel.

Since each channel has its own cutoff, QKᵀ is no longer a plain GEMM and maps poorly onto Tensor Cores. Our fix: split the head dimensions into groups of 16, matching the K size of one hardware MMA instruction, and give each group a single shared cutoff (the largest wavelength in it). Each group then runs as a standard MMA, followed by one mask on its output. Note that all results above use this approximation.

WavePrune lets us skip work. For each (query block, key block) tile, we compare each group’s cutoff with the tile’s range of distances |i−j|:
All distances exceed the cutoff → skip the group
All distances are within the cutoff → compute without masking

Otherwise → compute, then mask
Skipped groups save both FLOPs and memory reads.
vs. FlashAttention-2 (A100, head dim 128, rope_theta 1e6):

32K: 1.15× prefill, 1.24× decode

128K: 1.12× prefill, 1.28× decode

[Guancheng Du @ducatiqwq]

(5/7) Pretraining with WavePrune.

We pretrain GPT-2 base, medium, and large at 2K context and evaluate at 4K. Training loss is unchanged, while extrapolation loss drops substantially at every scale. Pretraining with WavePrune also outperforms applying it only at inference.

[Guancheng Du @ducatiqwq]

(4/7) Training-free gains on HELMET.

Applied to off-the-shelf models without any training, WavePrune improves 4 of 5 models (64K/128K average):
• Qwen3-8B: 35.7 → 40.0
• Qwen3-32B: 55.1 → 56.3
• Gemma3-12B: 38.2 → 38.7
• Ministral3-3B: 39.1 → 39.7

Llama3.1-8B regresses (38.2 → 35.2); we discuss this in Appendix E.

[Guancheng Du @ducatiqwq]

(3/7) The fix: keep a single period.

RoPE splits each head into two-dimensional channels. Channel c rotates by f_c per position, so one full period spans w_c = 2π/f_c positions.

WavePrune counts channel c in the attention logit only when the query–key distance stays within that period:

L_ij = (1/√d_h) · Σ_c φ_i^(c)(q_i)ᵀ φ_j^(c)(k_j) · 𝟙[|i−j| ≤ w_c]
Here φ^(c) is the standard RoPE rotation of channel c. The only thing we add is the indicator, 𝟙[|i−j| ≤ w_c].
With this single change, per-channel oscillation and ghost sub-diagonals both disappear.

[Guancheng Du @ducatiqwq]

(2/7) Periodicity causes position aliasing, hence ghost sub-diagonals.

In a RoPE channel with frequency f, a QK pair at distance d contributes ∝ cos(fd + ω), where ω is the initial angle. Because cos is 2π-periodic, distances with the same (fd mod 2π) look identical to the model.

In Qwen3-8B, high-frequency channels flip sign rapidly as d grows. After summation and softmax, these flips appear as spurious stripes in the attention map, which we call ghost sub-diagonals. We prove they are inherent to RoPE and distracts attention, so training alone cannot remove them.

[Guancheng Du @ducatiqwq]

(1/7) RoPE does not need to be periodic.

Attention retrieves information by matching queries to keys, so a relevant key should be reachable at any distance. RoPE instead rotates q and k with position, so their score oscillates with distance and becomes unreliable at long range. Methods like YaRN rescale the frequencies but keep the periodic rotation.

[Guancheng Du @ducatiqwq]

1) Training-free: improves HELMET on 4 of 5 LLMs, with 35.7 → 40.0 on Qwen3-8B.
2) Pretraining: consistently improves length extrapolation by large margins
3) Speedup: FlashWavePrune achieves 1.15× prefill and 1.24× decode speedups at 32K context.