Skip to content

fix(rotary): preserve fp32 inputs in mrope torch fallbacks - #641

Open
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/mrope-fp32-alias
Open

taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/mrope-fp32-alias

Conversation

@taking-lying-flat

Copy link
Copy Markdown
Contributor

For FP32 Q/K, the .float() slices in the Torch mRoPE fallbacks still alias the input. Writing the low half first therefore changes the input used to compute the high half. Rotating (1, 2) by 45 degrees produces (-0.707107, 0.914214) instead of (-0.707107, 2.121320).

Compute both results before either in-place write in _mrope_torch and apply_mrope_torch_fallback. This preserves the existing interface without cloning all of Q/K. The change is limited to the two Torch fallbacks; the Triton kernels are unchanged.

Validation on upstream 0781324:

  • Local regressions cover both fallbacks on CPU/CUDA with FP32, FP16 and BF16, including Q/K with different head counts and unchanged dimensions outside the rotary region. Before the fix: 4 FP32 cases failed and 8 dtype controls passed. After the fix: all 19 mRoPE cases passed (12 local regressions plus 7 existing cases).
  • git diff --check passed.
  • Hardware: Intel Core i7-12800HX; NVIDIA RTX A1000 Laptop GPU, 4 GiB, driver 595.91.07. Linux x86_64; PyTorch 2.13.0+cu130, Triton 3.7.1.
  • Exact local command: /home/cherry-cloud/miniconda3/envs/cherry/bin/python /tmp/run-freetoken-mrope-pr-tests.py --tb=short. The runner loads the real source modules while bypassing unrelated package initializers and the parent quantization-reset fixture, then runs tests/kernels/test_mrope.py through pytest. Normal package import was blocked by a missing msgpack dependency; this was not a full supported-environment test run.

This PR contains only production code. The regression additions and correction to the local test reference remain outside the commit. No checkpoint, full-model inference, or end-to-end performance validation was run; this does not claim an improvement to the normal BF16/Triton inference path.

KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 10, 2026
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 10, 2026
…in both mrope torch fallbacks

Assisted-by: Claude Opus 5.5

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant