Commit af98e4d
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 format1 parent b941c3c commit af98e4d
2 files changed
Lines changed: 652 additions & 0 deletions
File tree
- swift
- Sources/CoreAILanguageModels/Samplers
- Tests/LanguageModelsTests
0 commit comments