fix: make fill_next_token_bitmask stride-aware

cpp/grammar_matcher.cc

@@ -189,7 +189,29 @@ 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<int32_t*>(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];
+      // 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;
+    }
+  }
+
+  return reinterpret_cast<int32_t*>(token_bitmask.data) + index * row_stride;
 }
 
 void _DebugGetMaskedTokensFromBitmask(
@@ -225,8 +247,25 @@ void ApplyMask32Bits(
       logits->ndim == 2
           ? std::make_pair(static_cast<int>(logits->shape[0]), static_cast<int>(logits->shape[1]))
           : std::make_pair(1, static_cast<int>(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<int>(logits->strides[0]) : logits_shape.second;
+  int bitmask_stride0 = bitmask.strides != nullptr
+                            ? static_cast<int>(bitmask.strides[0])
+                            : static_cast<int>(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<int>(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<uint32_t*>(bitmask.data) + idx * bitmask_stride0;
@@ -274,8 +313,25 @@ void ApplyMask16Bits(
       logits->ndim == 2
           ? std::make_pair(static_cast<int>(logits->shape[0]), static_cast<int>(logits->shape[1]))
           : std::make_pair(1, static_cast<int>(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<int>(logits->strides[0]) : logits_shape.second;
+  int bitmask_stride0 = bitmask.strides != nullptr
+                            ? static_cast<int>(bitmask.strides[0])
+                            : static_cast<int>(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<int>(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<uint32_t*>(bitmask.data) + idx * bitmask_stride0;

tests/python/test_token_bitmask_operations.py

@@ -125,6 +125,121 @@ 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)] + ["<eos>"], 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 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