Skip to content

Add GPU-based constrained sampling for pipelined engine#114

Draft
stikves wants to merge 1 commit into
apple:mainfrom
stikves:feature/gpu-constrained-sampling
Draft

Add GPU-based constrained sampling for pipelined engine#114
stikves wants to merge 1 commit into
apple:mainfrom
stikves:feature/gpu-constrained-sampling

Conversation

@stikves

@stikves stikves commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Enables grammar-constrained generation (JSON schema) on the pipelined engine by applying the xgrammar bitmask inside the MPSGraph sampler on GPU, eliminating the forced fallback to the slower sequential engine.

Changes:

  • Add bitmask expansion graph to MPSGraphArgmaxSampler and MPSGraphCompositeSampler (bitwise AND + notEqual + cast, zero measurable latency overhead)
  • Add ConstrainedGenerationSession.fillBitmask(into:) for zero-copy write into shared MTLBuffer
  • Add CoreAIPipelinedEngine.generateConstrained() with semi-pipelined loop (inference overlaps bitmask computation)
  • Add PipelinedConstrainedDecodingStrategy for the GPU path
  • Route pipelined engine to GPU strategy in CoreAILanguageModel

Performance:

Bitmask expansion adds 0.65ms per token (measured on 32K vocab). On a Qwen3-VL-2B model, constrained generation on pipelined vs sequential produces identical outputs with comparable throughput

Tests:

  • Added new unit tests, and verified existing ones
  • Ran end-to-end inference
  • Greedy outputs are identical between pipelined and sequential

Enables grammar-constrained generation (JSON schema) on the pipelined
engine by applying the xgrammar bitmask inside the MPSGraph sampler on
GPU, eliminating the forced fallback to the slower sequential engine.

- Add bitmask expansion graph to MPSGraphArgmaxSampler and
  MPSGraphCompositeSampler (bitwise AND + notEqual + cast, zero
  measurable latency overhead)
- Add ConstrainedGenerationSession.fillBitmask(into:) for zero-copy
  write into shared MTLBuffer
- Add CoreAIPipelinedEngine.generateConstrained() with semi-pipelined
  loop (inference overlaps bitmask computation)
- Add PipelinedConstrainedDecodingStrategy for the GPU path
- Route pipelined engine to GPU strategy in CoreAILanguageModel
@stikves
stikves force-pushed the feature/gpu-constrained-sampling branch from 6170f93 to 0e2c429 Compare July 17, 2026 17:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant