Support incontiguous logits and bitmask. - #359
Merged
Merged
Conversation
Ubospica
approved these changes
Jul 12, 2025
Ubospica
left a comment
Ubospica
left a comment
Collaborator
There was a problem hiding this comment.
LGTM! We can add a test and merge it
Collaborator
|
We can merge it for now and add tests later |
Seven-Streams
pushed a commit
to CaiJohn/xgrammar
that referenced
this pull request
Aug 25, 2026
fill_next_token_bitmask addressed bitmask rows as `data + index * buffer_size`, assuming a compact layout, while apply_token_bitmask_inplace addresses them as `data + idx * strides[0]`. For a non-contiguous bitmask the two disagree: the mask is written to one row and read from another, leaving the requested logits row entirely -inf. Nothing raises, and the shape checks cannot catch it because a strided view has exactly the same shape as a contiguous one. apply became stride-aware in mlc-ai#359 (CUDA) and mlc-ai#390 (CPU/Triton); fill never followed. Before mlc-ai#390 both sides assumed the compact layout and therefore agreed, so non-contiguous bitmasks worked end to end -- this is a regression first released in v0.1.23. Address rows by strides[0] when strides are present, matching apply. All three callers of CheckAndGetBitmaskPtr are covered by the one change. Since a row is read as a single packed DynamicBitset, the vocabulary dimension must remain unit-stride; reject anything else rather than silently misreading it, matching the constraint mlc-ai#359 already documents for CUDA. Also fall back to the compact layout when strides is NULL in ApplyMask32Bits and ApplyMask16Bits, which dereferenced it unconditionally. The vendored DLPack is v1.0, where NULL strides is legal and means compact, and that header ships in the package's include dir.
Seven-Streams
pushed a commit
to CaiJohn/xgrammar
that referenced
this pull request
Sep 1, 2026
fill_next_token_bitmask addressed bitmask rows as `data + index * buffer_size`, assuming a compact layout, while apply_token_bitmask_inplace addresses them as `data + idx * strides[0]`. For a non-contiguous bitmask the two disagree: the mask is written to one row and read from another, leaving the requested logits row entirely -inf. Nothing raises, and the shape checks cannot catch it because a strided view has exactly the same shape as a contiguous one. apply became stride-aware in mlc-ai#359 (CUDA) and mlc-ai#390 (CPU/Triton); fill never followed. Before mlc-ai#390 both sides assumed the compact layout and therefore agreed, so non-contiguous bitmasks worked end to end -- this is a regression first released in v0.1.23. Address rows by strides[0] when strides are present, matching apply. All three callers of CheckAndGetBitmaskPtr are covered by the one change. Since a row is read as a single packed DynamicBitset, the vocabulary dimension must remain unit-stride; reject anything else rather than silently misreading it, matching the constraint mlc-ai#359 already documents for CUDA. Also fall back to the compact layout when strides is NULL in ApplyMask32Bits and ApplyMask16Bits, which dereferenced it unconditionally. The vendored DLPack is v1.0, where NULL strides is legal and means compact, and that header ships in the package's include dir.
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
2 participants
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.
This PR supports incontiguous logits and bitmasks when applying token bitmasks (The vocab stride still requires to be 1)