Skip to content

[bugfix] preserve QSA MRoPE batch dimension - #179

Open
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:fix/qsa-mrope-batch-dimension
Open

[bugfix] preserve QSA MRoPE batch dimension#179
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:fix/qsa-mrope-batch-dimension

Conversation

@taking-lying-flat

Copy link
Copy Markdown
Contributor

Summary

  • preserve the MRoPE batch dimension while materializing QSA cos/sin tensors
  • align batch-dependent RoPE tensors with QSA's [B, S, ...] query and block-key layouts
  • add a regression test proving that changing one sample's MRoPE frequencies does not affect another sample

Problem

MRoPE frequencies have shape [S, B, 1, D_rope]. Flattening them with reshape(S, -1) merges B into the rotary feature dimension. For B=2, this silently mixes both samples' position encodings; for larger batches, the resulting rotary dimension can exceed the QSA head dimension.

Tests

  • python -m pytest -q tests/test_qsa_indexer.py
  • pre-commit run --files src/mcore_bridge/model/modules/qsa_indexer.py tests/test_qsa_indexer.py

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.

2 participants