fix: make fill_next_token_bitmask stride-aware

[CaiJohn]

Summary

fill_next_token_bitmask addresses bitmask rows as data + index * buffer_size (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, and
the requested logits row comes out entirely -inf. Nothing raises — it is silent. The shape check
can't catch it: bm[::2] and a contiguous bitmask have identical shapes; only the strides differ.

Why this matters: non-contiguous bitmasks worked end to end before #390 — this is a regression,
not a new feature. The fix makes fill use strides[0] too, so both sides agree on where a row
lives, as they did before.

Trigger: bitmask.stride(0) != GetBitmaskSize(vocab_size). Row slices bm[:n] are contiguous and
safe.

Reproduction (pure Python, public API, CPU — on main / 0.2.3)

First two rows are negative controls:

import torch, xgrammar as xgr

tok = xgr.TokenizerInfo([f"t{i}" for i in range(40)] + ["<eos>"], stop_token_ids=[40])
ctx = xgr.GrammarCompiler(tok).compile_grammar(xgr.Grammar.from_ebnf('root ::= "t0"'))
bs = xgr.get_bitmask_shape(1, tok.vocab_size)[-1]

def probe(bm, tag):
    m = xgr.GrammarMatcher(ctx)
    logits = torch.zeros(bm.shape[0], tok.vocab_size)
    m.fill_next_token_bitmask(bm, 1)
    xgr.apply_token_bitmask_inplace(logits, bm)
    # `root ::= "t0"` accepts only token 0, so exactly that logit must stay finite.
    print(f"{tag:22s} stride={str(bm.stride()):10s} -> "
          f"surviving={torch.isfinite(logits[1]).nonzero().flatten().tolist()}")

probe(xgr.allocate_token_bitmask(2, tok.vocab_size),      "official allocate")
probe(torch.zeros(4, bs, dtype=torch.int32)[:2],          "row slice bm[:2]")
probe(torch.zeros(4, bs, dtype=torch.int32)[::2],         "strided bm[::2]")
probe(torch.zeros(2, bs + 3, dtype=torch.int32)[:, :bs],  "col-cropped (padded)")

official allocate     stride=(2, 1)  -> surviving=[0]     OK
row slice bm[:2]      stride=(2, 1)  -> surviving=[0]     OK
strided bm[::2]       stride=(4, 1)  -> surviving=[]      ALL -inf  <-- BUG
col-cropped (padded)  stride=(5, 1)  -> surviving=[]      ALL -inf  <-- BUG

This is a regression from #390 (first in v0.1.23), not a revert

Before #390, apply addressed rows as idx * shape[1], and fill enforces shape[1] == buffer_size
(grammar_matcher.cc:167), so both sides computed the same offset. #390 changed apply to
strides[0] (correctly, to fix logits stride, vLLM
#19493) but left fill unchanged, so the two
diverged. Building both versions and running the identical script confirms it:

xgrammar 0.1.22                      xgrammar 0.2.3
strided bm[::2]     -> [0] OK        strided bm[::2]     -> []  ALL -inf
col-cropped padded  -> [0] OK        col-cropped padded  -> []  ALL -inf

The old version computes the correct mask (surviving [0], not merely "not all -inf"). This is
not a revert of #390 — its logits-stride fix is correct and stays; the defect is that fill is
still not stride-aware, consistent with the direction #359 and #390 took (apply supports
non-contiguous bitmasks).

Changes (cpp/grammar_matcher.cc)

1. CheckAndGetBitmaskPtr — row offset becomes index * row_stride, with
   row_stride = strides ? strides[0] : buffer_size. All three callers
   (fill_next_token_bitmask, _DebugGetMaskedTokensFromBitmask, _IsSingleTokenBitmask) go
   through this one function.
2. Vocab-dim unit-stride check — a row is read as one packed DynamicBitset, so having
   trusted strides[0] the vocab dim must stay unit-stride; reject bm[:, ::2] rather than
   silently misread it (same constraint Support incontiguous logits and bitmask. #359 documents for CUDA). Verified to fire.
3. ApplyMask32Bits/ApplyMask16Bits — fall back to compact layout when strides is NULL
   instead of dereferencing it. Vendored DLPack is v1.0, where NULL means compact, and that header
   ships in the package's include dir, so a C++ caller can follow it, pass NULL, and segfault.
4. Overlapping-row guard (fill and apply) — a row spans shape[-1] elements, so with more
   than one row a stride smaller than that overlaps adjacent rows and lets a write run past the
   buffer (the stride(0) < buffer_size OOB case). Reject it with a clear error rather than
   corrupt memory. Gated on row count > 1: a single row cannot overlap and legitimately carries an
   unconstrained stride (e.g. a (1, vocab) logits view with stride0 == 1).

Tests

tests/python/test_token_bitmask_operations.py, next to #390's own stride test. Four parametrized
layouts (allocate/row_slice contiguous controls; strided/col_cropped non-contiguous), plus
guard tests for bm[:, ::2] (vocab stride) and an overlapping-row layout. The load-bearing one is
test_fill_apply_round_trip_row_addressing: it passes pre-#390 and fails on main, which is what
makes the regression checkable.

Results (CPU, macOS): the whole file is 55 passed / 0 failed (57 skipped, CUDA/Triton/MLX);
full tests/python/ 2456 passed / 0 failed.

Scope / honesty

• Latent, not active. Enumerated every xgrammar fill call site in vLLM (7aab6e2) and SGLang
  (947a14d), reproduced each in standalone torch: 13 paths, 0 trigger — they pass the whole
  tensor + a row index, or a step-1 slice. But it's out of reach, not defended: a
  fill(bitmask[i:i+1]) refactor makes it active. (Secondary: with stride(0) < buffer_size, e.g.
  expand, it's an out-of-bounds write — contrived, raises the severity ceiling not the
  probability; change 4 now rejects it with a clear error.)
• Alternative: have fill reject non-contiguous bitmasks — smaller, but runs against Support incontiguous logits and bitmask. #359/Fix apply_bitmask logit for both CPU and triton versions when shape and stride doesn't match #390
  and kills usages that worked pre-Fix apply_bitmask logit for both CPU and triton versions when shape and stride doesn't match #390. Happy to switch if maintainers prefer that contract.
• The NULL-strides fallback (change 3) is there because the vendored DLPack is v1.0, where NULL
  strides is legal and means compact; apply dereferenced it unconditionally. It's a defensive
  change, reachable only from C++. Happy to split it out.
• Deferred: per-index bounds-checking on the indices path (raised in review) is pre-existing
  behaviour this PR doesn't touch, so I kept it out of a focused stride fix — happy to do it in a
  follow-up.
• Not tested on CUDA (no local GPU); the change is in the shared CPU path, and CUDA already
  handles non-contiguous bitmasks per Support incontiguous logits and bitmask. #359.

Related, none duplicate

#	relation
#390	introduced this; fixed logits stride, also changed bitmask addressing
#359	CUDA apply supports non-contiguous bitmask — the intended direction
#220	logits vocab-dim padding — realistic source of the column-crop case

[gemini-code-assist]

Code Review

This pull request updates the grammar matcher to support non-contiguous bitmasks and logits by correctly respecting strides instead of assuming a compact layout. It also adds comprehensive tests to verify row addressing across various memory layouts. The review feedback highlights potential security and memory safety issues, suggesting the addition of bounds checks on indices and validation of row strides to prevent overlapping memory writes or out-of-bounds accesses.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Seven-Streams]

Maybe we can use xgrammar_log(fatal) here? Since it's a user's input error.

[Seven-Streams]

ditto

[Seven-Streams]

ditto. xgrammar_check is used to check the correctness of inner logic, while xgammar_log(fatal) is used to check the users' inputs.

[CaiJohn]

Thanks — and no objection to the change itself: XGRAMMAR_LOG(FATAL) does read better for an input error, since it drops the Check failed: (cond) is false: prefix and leaves just the message. I have the conversion done locally for all six guards this PR adds.

Before I push it, though, I want to check one thing, because the header documents the mapping the other way round. In cpp/support/logging.h:

• XGRAMMAR_CHECK (~:207) — documented as being for checking the correctness of user inputs
• XGRAMMAR_ICHECK (~:213) — documented as being for internal conditions
• XGRAMMAR_DCHECK (~:220) — same, debug-only, and preferred over ICHECK

Those comments came in with #354, and internal assertions in cpp/ mostly go through DCHECK rather than CHECK.

So I'd rather not guess. Which do you mean?

1. The doc comments are stale and XGRAMMAR_LOG(FATAL) is now the intended form for input validation — I push the conversion as is, and can update the header comments in the same PR or a follow-up if you want.
2. The guidance is really about the message format (no Check failed: prefix for user-facing errors), in which case XGRAMMAR_LOG(FATAL) is right here and the header could use a note.
3. XGRAMMAR_CHECK is fine per the header and these should stay as they are.

Happy with any of them — just say which and I'll push accordingly.