Repository navigation
Conversation
Optional position_ids runtime input, shape (batch, seq_len) int64: RoPE cos/sin are gathered per token instead of the contiguous start_pos slice, and rows may use different positions, so packed documents can restart positions at zero. apply_rope broadcasts the batched (batch, seq_len, head_dim // 2) frequency form over heads. Dtype, shape, and bounds are validated with clear errors before compute; kv_caches and MoT combinations are rejected. Default None keeps the contiguous slice bit-exact and no state-dict keys change. Co-Authored-By: Timothy Ngotiaoco <timothy_ngotiaoco@harvard.edu>
The shape and bounds checks for position_ids ran after the token embedding, the prefix concat and the RoPE table device transfer, so an invalid position_ids wasted that compute before failing. All position_ids validation now runs at the top of forward: dtype, rank, shape against the residual length (text from tokens or inputs_embeds plus any prefix_embeds), and bounds against the RoPE table length. Error messages are unchanged and the default None path is untouched. Tests assert the token embedding is not called for a bad shape, an out-of-range id, and text-length ids with prefix_embeds set. They also cover the inputs_embeds route and a MoMa forward with explicit contiguous positions. Co-Authored-By: Timothy Ngotiaoco <timothy_ngotiaoco@harvard.edu>
A model materialized with to_empty() has no RoPE tables until init_weights_and_freqs() runs. In that state the explicit position_ids bounds check was skipped, so the token embedding and the prefix concat ran before the later RoPE assertion fired. The top-of-forward position_ids validation now rejects the call first, with the same message. The position_ids=None path is unchanged. The test builds the model on the meta device, asserts through a forward pre-hook spy that the token embedding never runs, and checks as a control that it does once init_weights_and_freqs() has been called. Co-Authored-By: Timothy Ngotiaoco <timothy_ngotiaoco@harvard.edu>
Codecov Report✅ All modified and coverable lines are covered by tests.
🚀 New features to boost your workflow:
|
Only the per-token gather is new; the transfer and the contiguous slice behave exactly as before. Co-Authored-By: Timothy Ngotiaoco <timothy_ngotiaoco@harvard.edu>
amazloumi
marked this pull request as draft
October 16, 2026 14:38
amazloumi
marked this pull request as ready for review
October 16, 2026 14:42
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
position_idsinput toTransformer.forward: a(batch, seq_len)int64 tensor of per-token RoPE positions. Rows may differ, so documents in a packed sequence can restart at position zero.apply_ropeaccepts the batched(batch, seq_len, head_dim // 2)form alongside the shared one.modality.prefix_embeds) and the RoPE table range at the top offorward, before any compute.kv_cachesand MoT are rejected.position_ids=None) is unchanged.Testing
uv run ruff checkanduv run ruff format --checkpass on the touched filesuv run pyright kempnerforge/passes (0 errors)uv run pytest tests/unit/: 2281 passed, 3 skippedposition_idsreproduceposition_ids=Nonebitwise (dense and MoMa), and the state dict gains no entriesCloses #259