Skip to content

Commit 6921969

Browse files
committed
Add GPU bitmask expansion for constrained sampling
MPSGraph subgraph that expands packed Int32 xgrammar bitmasks into Float16 logits masks entirely on GPU. Blocked tokens get -65504 (Float16 min), ensuring they are never selected by argmax or topK. Both MPSGraphArgmaxSampler and MPSGraphCompositeSampler gain: - bitmaskBuffer (lazy): shared-memory MTLBuffer, allocated on first use - constrainedExecutable (lazy): compiled graph with bitmask input, built on first applyBitmask: true call - encode(..., applyBitmask:) overloads on all encode paths Zero overhead for unconstrained generation — no bitmask buffer or constrained graph is compiled until the first constrained call. Existing protocol methods and behavior unchanged. Part 2 of 4 for GPU-based constrained sampling (#114).
1 parent b66a122 commit 6921969

2 files changed

Lines changed: 640 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)