Add GPU-based constrained sampling for pipelined engine - #114
Conversation
0e2c429 to
d4edbfb
Compare
905f81f to
cdef085
Compare
43f429e to
1dbb338
Compare
fe2a707 to
3c6b90d
Compare
f8cdbd8 to
0b42177
Compare
carinapeng
left a comment
carinapeng
left a comment
There was a problem hiding this comment.
Thanks for looking into this. Getting this working is tricky because constrained decoding has some data dependency that vanilla decoding doesn't. From the grammar side, the accepted token to fill bit masking runs on CPU and the bit mask for token n+1 would depend on the sampled value of token n so this GPU-CPU roundtrip might need more design change. Happy to chat more on this and this work would probably benefit from breaking this into smaller pieces!
Masking happens on the GPU sampler: Yes, there is a CPU work for the xgrammar, however unlike the sequential engine, it does not "stop world" to generate masks. The bitmask buffer is filled on CPU by xgrammar (which does need the prior token), hence the unavoidable synchronization. Then the sampler executable directly uses this buffer. This is explained in the PR description, along with the relative performance metrics. |
|
This 2k line change is hard to review with confidence. A proposal is to land the GPU-masking capability and tests first, then the jump-forward + session cache as a follow-up |
Thanks — agreed the masking itself is on-GPU and correct, and the "no stop-the-world" framing is right: we avoid the full logit readback + CPU softmax the sequential engine pays -- this can be a PR in itself to land first To make sure we're on the same page on the pipelining, though.. There's still a per-token sync point — the sampled token n has to be read back to CPU so acceptToken(n) can run before fillBitmask(n+1). That's the "unavoidable synchronization" you mention, and this might be why the sampled positions can't overlap. So the real speedups here are (1) GPU masking and (2) jump-forward batching the deterministic runs — not pipelining of the constrained steps Proposal to move forward:
|
Sure, I'll send individual PRs and rebase this as they get merged. This will stay as the last piece where everything is eventually wired. |
…direct bitmask fill - rollback(_:) — unwind grammar state by N tokens (budget: 64) - findJumpForwardString() — peek at deterministic continuations - fillBitmask(into:) → BitmaskResult — write bitmask directly into a caller-provided pointer (e.g. GPU-visible MTLBuffer), avoiding array allocation per token - BitmaskResult enum: .terminated, .unconstrained, .constrained - C bridge additions: XGrammarRollback, XGrammarFindJumpForwardString, XGrammarFillNextTokenBitmask Part 1 of 4 for GPU-based constrained sampling (apple#114).
Both MPSGraphArgmaxSampler and MPSGraphCompositeSampler now accept an optional bitmask MTLBuffer. When provided, the graph applies the mask before sampling — tokens with bit=0 are suppressed by adding -inf to their logits. The bitmask is a packed Int32 array: bit (tokenID % 32) of word (tokenID / 32). This matches xgrammar's output format so the mask can be filled directly by ConstrainedGenerationSession.fillBitmask(). Part 2 of 4 for GPU-based constrained sampling (apple#114).
…direct bitmask fill - rollback(_:) — unwind grammar state by N tokens (budget: 64) - findJumpForwardString() — peek at deterministic continuations - fillBitmask(into:) → BitmaskResult — write bitmask directly into a caller-provided pointer (e.g. GPU-visible MTLBuffer), avoiding array allocation per token - BitmaskResult enum: .terminated, .unconstrained, .constrained - C bridge additions: XGrammarRollback, XGrammarFindJumpForwardString, XGrammarFillNextTokenBitmask Part 1 of 4 for GPU-based constrained sampling (apple#114).
…direct bitmask fill (#131) * Extend ConstrainedGenerationSession with rollback, jump-forward, and direct bitmask fill - rollback(_:) — unwind grammar state by N tokens (budget: 64) - findJumpForwardString() — peek at deterministic continuations - fillBitmask(into:) → BitmaskResult — write bitmask directly into a caller-provided pointer (e.g. GPU-visible MTLBuffer), avoiding array allocation per token - BitmaskResult enum: .terminated, .unconstrained, .constrained - C bridge additions: XGrammarRollback, XGrammarFindJumpForwardString, XGrammarFillNextTokenBitmask Part 1 of 4 for GPU-based constrained sampling (#114). * Add tests for rollback and findJumpForwardString (carinapeng)
* 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). * Fix Swift format
Enable grammar-constrained generation (JSON schema enforcement) on the pipelined engine by applying xgrammar bitmasks directly in the MPSGraph GPU sampler, eliminating the per-token logit transfer to CPU. Key components: - ConstrainedGenerationCapable protocol: capability signal for routing constrained generation to engines that support GPU-side bitmask application. Replaces concrete type checks with protocol-based dispatch. - ConstrainedSessionHandle: encapsulated, narrow-API handle to the grammar session. Private session storage with only the 5 operations the engine loop needs exposed. Single-writer safety guaranteed by the engine acquire/release gate. - PipelinedConstrainedDecodingStrategy: AsyncSequence-based strategy that drives the GPU constrained loop. SingleUseFlag iteration guard, proper error propagation, Task.checkCancellation, early-stop on consumer drop. - Session cache (checkout/checkin): sessions reused across calls with the same schema. Engine Task defer block owns return-to-cache, ensuring the handle is never accessed concurrently. - Bitmask expansion in MPSGraph samplers: bitwise AND + notEqual + cast graph with zero measurable overhead. Both ArgmaxSampler and CompositeSampler support constrained path via applyBitmask parameter. - ConstrainedGenerationSession.fillBitmask(into:): zero-copy write of grammar bitmask directly into GPU-visible shared MTLBuffer. - Integration tests via MockConstrainedEngine conforming to the protocol, exercising session cache, schema invalidation, stop sequences, error propagation, and maxTokens boundary without Metal hardware.
…safety - precondition assert in generateConstrained ensures no prior generation is in flight when called directly (the public generate() path already drains, but protocol callers could skip it) - One-line comment on logits buffer binding clarifying sequential safety
…llback, masking - Fix logits buffer overflow on prefill: size to prefillTokens.count, not 1 - Fix jump-forward dropping lastToken from model context: prepend it to jumpTokens so the model sees the full sequence - Enable rollback: set maxRollbackTokens=64, make C bridge return bool, guard tokenizeJumpForward on rollback budget - Clear history after constrained reset to prevent KV corruption on next call - Fix capabilities: check ConstrainedGenerationCapable conformance, not just supportsLogits - Replace precondition with throw for in-flight check (recoverable error) - NaN-safe masking: use graph.select instead of multiply-add (also faster, eliminates 3 graph ops)
- Elevate multi-token stop sequence warning to level=0 (always visible) - Set innerIterator = nil on finish/error to release the stream and trigger onTermination (prevents session leak) - Remove outdated LLMAsset reference from error message
4828fe0 to
209aebd
Compare
| } | ||
|
|
||
| /// Whether the loaded engine supports GPU-side constrained generation, or `nil` when unloaded. | ||
| var loadedEngineIsConstrainedCapable: Bool? { |
There was a problem hiding this comment.
what about staticShapeInferenceEngine? does this affect that?
4 participants
Summary
Performance (Qwen3-4B, relative to unconstrained pipelined baseline)
Changes
ConstrainedGenerationCapable— Protocol for engines that support GPU-side bitmask applicationConstrainedSessionHandle— Encapsulated handle with narrow API, stores tokenizer for JFPipelinedConstrainedDecodingStrategy— AsyncSequence-based strategy with single-use guard, early-stop on consumer dropMPSGraphSamplers— Bitmask expansion graph,applyBitmaskon encode pathsConstrainedGenerationSession—fillBitmask(into:)withBitmaskResultenum,rollback,findJumpForwardStringrollback,findJumpForwardString,isCompletedllm-runner— Routes to pipelined strategy, performance metrics for constrained pathTest plan