From f9208cae50f6ee41c21343e820b8866222844aca Mon Sep 17 00:00:00 2001 From: john Date: Wed, 15 Jul 2026 18:03:16 +0800 Subject: [PATCH 1/3] fix: make fill_next_token_bitmask stride-aware 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 #359 (CUDA) and #390 (CPU/Triton); fill never followed. Before #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 #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. --- cpp/grammar_matcher.cc | 34 ++++++- tests/python/test_token_bitmask_operations.py | 99 +++++++++++++++++++ 2 files changed, 128 insertions(+), 5 deletions(-) diff --git a/cpp/grammar_matcher.cc b/cpp/grammar_matcher.cc index d6622d700..1e0126123 100644 --- a/cpp/grammar_matcher.cc +++ b/cpp/grammar_matcher.cc @@ -189,7 +189,23 @@ int32_t* CheckAndGetBitmaskPtr(const DLTensor& token_bitmask, int vocab_size, in token_bitmask.device.device_type == kDLROCMHost ) << "The provided bitmask's device is not valid: should be CPU"; - return reinterpret_cast(token_bitmask.data) + index * buffer_size; + // The bitmask may be a non-contiguous view, so the row offset must follow strides[0] instead + // of assuming a compact layout: ApplyMask32Bits addresses rows the same way, and filling and + // applying must agree on where a row lives. Null strides means compact (DLPack). The vocab + // dimension must stay unit-stride, since a row is read as one contiguous DynamicBitset. + int64_t row_stride = buffer_size; + if (token_bitmask.strides != nullptr) { + int64_t vocab_stride = token_bitmask.strides[token_bitmask.ndim - 1]; + XGRAMMAR_CHECK(vocab_stride == 1) + << "The provided bitmask must be contiguous along the vocabulary dimension, but got " + "stride " + << vocab_stride; + if (token_bitmask.ndim == 2) { + row_stride = token_bitmask.strides[0]; + } + } + + return reinterpret_cast(token_bitmask.data) + index * row_stride; } void _DebugGetMaskedTokensFromBitmask( @@ -225,8 +241,12 @@ void ApplyMask32Bits( logits->ndim == 2 ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) : std::make_pair(1, static_cast(logits->shape[0])); - int logits_stride0 = logits->strides[0]; - int bitmask_stride0 = bitmask.strides[0]; + // Null strides means compact (DLPack), in which case a row is one shape[-1]-long span. + int logits_stride0 = + logits->strides != nullptr ? static_cast(logits->strides[0]) : logits_shape.second; + int bitmask_stride0 = bitmask.strides != nullptr + ? static_cast(bitmask.strides[0]) + : static_cast(bitmask.shape[bitmask.ndim - 1]); if (indices.has_value()) { for (auto idx : indices.value()) { uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; @@ -274,8 +294,12 @@ void ApplyMask16Bits( logits->ndim == 2 ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) : std::make_pair(1, static_cast(logits->shape[0])); - int logits_stride0 = logits->strides[0]; - int bitmask_stride0 = bitmask.strides[0]; + // Null strides means compact (DLPack), in which case a row is one shape[-1]-long span. + int logits_stride0 = + logits->strides != nullptr ? static_cast(logits->strides[0]) : logits_shape.second; + int bitmask_stride0 = bitmask.strides != nullptr + ? static_cast(bitmask.strides[0]) + : static_cast(bitmask.shape[bitmask.ndim - 1]); if (indices.has_value()) { for (auto idx : indices.value()) { uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; diff --git a/tests/python/test_token_bitmask_operations.py b/tests/python/test_token_bitmask_operations.py index 11137d3fd..cb89f12c4 100644 --- a/tests/python/test_token_bitmask_operations.py +++ b/tests/python/test_token_bitmask_operations.py @@ -125,6 +125,105 @@ def test_apply_token_bitmask_inplace_shape_stride_mismatch(device: str): torch.testing.assert_close(logits, expected) +def _bitmask_buffer_size(vocab_size: int) -> int: + return xgr.get_bitmask_shape(1, vocab_size)[-1] + + +def _single_token_grammar() -> Tuple[xgr.TokenizerInfo, xgr.CompiledGrammar]: + """A grammar accepting exactly one token, so the expected mask is unambiguous.""" + tokenizer_info = xgr.TokenizerInfo([f"t{i}" for i in range(40)] + [""], stop_token_ids=[40]) + compiled_grammar = xgr.GrammarCompiler(tokenizer_info).compile_grammar( + xgr.Grammar.from_ebnf('root ::= "t0"') + ) + return tokenizer_info, compiled_grammar + + +def _bitmask_allocated(batch_size: int, vocab_size: int) -> torch.Tensor: + return xgr.allocate_token_bitmask(batch_size, vocab_size) + + +def _bitmask_row_sliced(batch_size: int, vocab_size: int) -> torch.Tensor: + buffer_size = _bitmask_buffer_size(vocab_size) + return torch.zeros(batch_size * 2, buffer_size, dtype=torch.int32)[:batch_size] + + +def _bitmask_strided(batch_size: int, vocab_size: int) -> torch.Tensor: + buffer_size = _bitmask_buffer_size(vocab_size) + return torch.zeros(batch_size * 2, buffer_size, dtype=torch.int32)[::2] + + +def _bitmask_col_cropped(batch_size: int, vocab_size: int) -> torch.Tensor: + """A buffer preallocated on a padded vocab, then cropped to the real buffer size.""" + buffer_size = _bitmask_buffer_size(vocab_size) + return torch.zeros(batch_size, buffer_size + 3, dtype=torch.int32)[:, :buffer_size] + + +# `allocate` and `row_slice` are contiguous and act as negative controls: they must pass +# both before and after any stride fix. `strided` and `col_cropped` have +# stride(0) != buffer_size, which is what these tests are about. +bitmask_layouts = { + "allocate": _bitmask_allocated, + "row_slice": _bitmask_row_sliced, + "strided": _bitmask_strided, + "col_cropped": _bitmask_col_cropped, +} + + +@pytest.mark.parametrize("layout", list(bitmask_layouts)) +def test_fill_next_token_bitmask_row_addressing(layout: str): + """fill_next_token_bitmask must write into the row the caller asked for, regardless of + the bitmask's memory layout. + + It addresses rows as `data + index * buffer_size`, ignoring stride(0), so for a + non-contiguous bitmask it writes outside the requested row. + """ + tokenizer_info, compiled_grammar = _single_token_grammar() + batch_size, index = 2, 1 + + expected = xgr.allocate_token_bitmask(batch_size, tokenizer_info.vocab_size) + xgr.GrammarMatcher(compiled_grammar).fill_next_token_bitmask(expected, index) + + bitmask = bitmask_layouts[layout](batch_size, tokenizer_info.vocab_size) + xgr.GrammarMatcher(compiled_grammar).fill_next_token_bitmask(bitmask, index) + + torch.testing.assert_close(bitmask[index], expected[index]) + + +@pytest.mark.parametrize("layout", list(bitmask_layouts)) +def test_fill_apply_round_trip_row_addressing(layout: str): + """fill_next_token_bitmask -> apply_token_bitmask_inplace must mask the logits row + corresponding to the bitmask row that was filled. + + fill addresses rows by buffer_size and apply by stride(0). They agree only when the + bitmask is contiguous; otherwise the mask is written to one row and read from another, + leaving the victim row all -inf. Silent: nothing raises. + """ + tokenizer_info, compiled_grammar = _single_token_grammar() + batch_size, index = 2, 1 + vocab_size = tokenizer_info.vocab_size + + bitmask = bitmask_layouts[layout](batch_size, vocab_size) + xgr.GrammarMatcher(compiled_grammar).fill_next_token_bitmask(bitmask, index) + + logits = torch.zeros(batch_size, vocab_size) + xgr.apply_token_bitmask_inplace(logits, bitmask) + + # `root ::= "t0"` accepts only token 0, so exactly that logit must stay finite. + surviving = torch.isfinite(logits[index]).nonzero().flatten().tolist() + assert surviving == [0] + + +def test_fill_next_token_bitmask_rejects_strided_vocab_dim(): + """A row of the bitmask is read as one packed bitset, so a bitmask whose vocabulary + dimension is not unit-stride must be rejected rather than silently misread.""" + tokenizer_info, compiled_grammar = _single_token_grammar() + buffer_size = _bitmask_buffer_size(tokenizer_info.vocab_size) + bitmask = torch.zeros(2, buffer_size * 2, dtype=torch.int32)[:, ::2] + + with pytest.raises(RuntimeError, match="contiguous along the vocabulary dimension"): + xgr.GrammarMatcher(compiled_grammar).fill_next_token_bitmask(bitmask, 1) + + def get_apply_token_bitmask_kernel(impl: str) -> Callable: if impl == "cpu": from xgrammar.kernels.apply_token_bitmask_inplace_cpu import apply_token_bitmask_inplace_cpu From b7dd98cf0f220ae1f87621217fc3759adb248f59 Mon Sep 17 00:00:00 2001 From: john Date: Fri, 17 Jul 2026 11:30:34 +0800 Subject: [PATCH 2/3] style: black-format test file --- tests/python/test_token_bitmask_operations.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/python/test_token_bitmask_operations.py b/tests/python/test_token_bitmask_operations.py index cb89f12c4..780742e3a 100644 --- a/tests/python/test_token_bitmask_operations.py +++ b/tests/python/test_token_bitmask_operations.py @@ -131,7 +131,9 @@ def _bitmask_buffer_size(vocab_size: int) -> int: def _single_token_grammar() -> Tuple[xgr.TokenizerInfo, xgr.CompiledGrammar]: """A grammar accepting exactly one token, so the expected mask is unambiguous.""" - tokenizer_info = xgr.TokenizerInfo([f"t{i}" for i in range(40)] + [""], stop_token_ids=[40]) + tokenizer_info = xgr.TokenizerInfo( + [f"t{i}" for i in range(40)] + [""], stop_token_ids=[40] + ) compiled_grammar = xgr.GrammarCompiler(tokenizer_info).compile_grammar( xgr.Grammar.from_ebnf('root ::= "t0"') ) From 8eb2e0b665fa69026011af82a924564d3382d59b Mon Sep 17 00:00:00 2001 From: john Date: Fri, 17 Jul 2026 11:45:41 +0800 Subject: [PATCH 3/3] fix: reject overlapping bitmask/logits rows in fill and apply Address code-review feedback: a row spans shape[-1] elements, so a row stride smaller than that overlaps adjacent rows and lets a write run past the buffer (the OOB case noted for stride(0) < buffer_size). Reject such layouts with a clear error in CheckAndGetBitmaskPtr (fill) and, symmetrically, in ApplyMask32Bits/ApplyMask16Bits (apply), rather than corrupting memory. Only guard when there is more than one row: a single row cannot overlap and legitimately carries an unconstrained stride (e.g. a (1, vocab) logits view with stride0 == 1), which a naive ndim==2 guard would wrongly reject. --- cpp/grammar_matcher.cc | 32 +++++++++++++++++++ tests/python/test_token_bitmask_operations.py | 14 ++++++++ 2 files changed, 46 insertions(+) diff --git a/cpp/grammar_matcher.cc b/cpp/grammar_matcher.cc index 1e0126123..aa20adb27 100644 --- a/cpp/grammar_matcher.cc +++ b/cpp/grammar_matcher.cc @@ -202,6 +202,12 @@ int32_t* CheckAndGetBitmaskPtr(const DLTensor& token_bitmask, int vocab_size, in << vocab_stride; if (token_bitmask.ndim == 2) { row_stride = token_bitmask.strides[0]; + // A row spans buffer_size int32s; with more than one row a smaller stride would overlap + // adjacent rows and let a write run past the buffer. A single row cannot overlap, whatever + // its stride, so only guard when rows > 1. Reject rather than corrupt memory. + XGRAMMAR_CHECK(token_bitmask.shape[0] <= 1 || row_stride >= buffer_size) + << "The row stride of the bitmask must be at least " << buffer_size + << " so rows do not overlap, but got " << row_stride; } } @@ -247,6 +253,19 @@ void ApplyMask32Bits( int bitmask_stride0 = bitmask.strides != nullptr ? static_cast(bitmask.strides[0]) : static_cast(bitmask.shape[bitmask.ndim - 1]); + // Mirror the fill-side check: with more than one row, a row stride smaller than the per-row + // span (shape[-1]) overlaps adjacent rows -- corrupting on the logits -inf write. A single row + // cannot overlap and legitimately carries an unconstrained stride, so only guard when rows > 1. + if (logits_shape.first > 1) { + XGRAMMAR_CHECK(logits_stride0 >= logits_shape.second) + << "The row stride of the logits must be at least " << logits_shape.second + << " so rows do not overlap, but got " << logits_stride0; + } + if (bitmask.ndim == 2 && bitmask.shape[0] > 1) { + XGRAMMAR_CHECK(bitmask_stride0 >= static_cast(bitmask.shape[1])) + << "The row stride of the bitmask must be at least " << bitmask.shape[1] + << " so rows do not overlap, but got " << bitmask_stride0; + } if (indices.has_value()) { for (auto idx : indices.value()) { uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; @@ -300,6 +319,19 @@ void ApplyMask16Bits( int bitmask_stride0 = bitmask.strides != nullptr ? static_cast(bitmask.strides[0]) : static_cast(bitmask.shape[bitmask.ndim - 1]); + // Mirror the fill-side check: with more than one row, a row stride smaller than the per-row + // span (shape[-1]) overlaps adjacent rows -- corrupting on the logits -inf write. A single row + // cannot overlap and legitimately carries an unconstrained stride, so only guard when rows > 1. + if (logits_shape.first > 1) { + XGRAMMAR_CHECK(logits_stride0 >= logits_shape.second) + << "The row stride of the logits must be at least " << logits_shape.second + << " so rows do not overlap, but got " << logits_stride0; + } + if (bitmask.ndim == 2 && bitmask.shape[0] > 1) { + XGRAMMAR_CHECK(bitmask_stride0 >= static_cast(bitmask.shape[1])) + << "The row stride of the bitmask must be at least " << bitmask.shape[1] + << " so rows do not overlap, but got " << bitmask_stride0; + } if (indices.has_value()) { for (auto idx : indices.value()) { uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; diff --git a/tests/python/test_token_bitmask_operations.py b/tests/python/test_token_bitmask_operations.py index 780742e3a..94de759ea 100644 --- a/tests/python/test_token_bitmask_operations.py +++ b/tests/python/test_token_bitmask_operations.py @@ -226,6 +226,20 @@ def test_fill_next_token_bitmask_rejects_strided_vocab_dim(): xgr.GrammarMatcher(compiled_grammar).fill_next_token_bitmask(bitmask, 1) +def test_fill_next_token_bitmask_rejects_overlapping_rows(): + """A row spans buffer_size int32s, so a row stride smaller than that overlaps adjacent rows + and a write would run past the buffer. Such a layout must be rejected, not silently corrupt + memory.""" + tokenizer_info, compiled_grammar = _single_token_grammar() + buffer_size = _bitmask_buffer_size(tokenizer_info.vocab_size) + assert buffer_size >= 2 # otherwise there is no room for an overlapping row stride + # Two rows with row stride 1 < buffer_size, vocab dim still unit-stride: rows overlap. + bitmask = torch.zeros(buffer_size + 1, dtype=torch.int32).as_strided((2, buffer_size), (1, 1)) + + with pytest.raises(RuntimeError, match="rows do not overlap"): + xgr.GrammarMatcher(compiled_grammar).fill_next_token_bitmask(bitmask, 1) + + def get_apply_token_bitmask_kernel(impl: str) -> Callable: if impl == "cpu": from xgrammar.kernels.apply_token_bitmask_inplace_cpu import apply_token_bitmask_inplace_cpu