Skip to content

Commit af98e4d

Browse files
authored
Add GPU bitmask expansion for constrained sampling (apple#141)
* 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 (apple#114). * Fix Swift format
1 parent b941c3c commit af98e4d

2 files changed

Lines changed: 652 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)