[Bugfix] Fix xgrammar dtype mismatch on macOS CPU inference

[karanb192]

Summary

• Fixes guided decoding with xgrammar failing on macOS CPU inference due to dtype mismatch
• The xgrammar CPU kernel requires float32 logits, but macOS CPU inference uses float16/bfloat16
• Adds dtype conversion in apply_grammar_bitmask() to convert to float32 before applying the bitmask, then copy back to original dtype

Credits

Fix approach based on @Inokinoki's suggestion in #31901.

Test Plan

Reproduction script from issue:

from pydantic import BaseModel
from vllm import LLM, SamplingParams
from vllm.sampling_params import StructuredOutputsParams

llm = LLM(model="Qwen/Qwen3-0.6B", max_model_len=512)

class OutputJsonFormat(BaseModel):
    python_code: str

base_schema = OutputJsonFormat.model_json_schema()

prompt = "Write a simple hello world in Python:"

sampling_params = SamplingParams(
    temperature=0.0,
    structured_outputs=StructuredOutputsParams(json=base_schema),
)

outputs = llm.generate(prompt, sampling_params)
print(outputs[0].outputs[0].text)

Before fix: ValueError: logits must be of type float32
After fix: Successfully generates structured output

Related Issue

Fixes #31901

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request addresses a dtype mismatch bug that occurs during guided decoding with xgrammar on macOS with CPU inference. The fix correctly converts logits to float32 before applying the bitmask to accommodate the CPU kernel's requirement. My review includes a suggestion to make this fix more targeted by adding a device check, which will prevent unnecessary data type conversions and potential performance overhead during GPU inference.

[karanb192]

Test Results on macOS M3 Pro (36GB RAM)

Environment

• Hardware: MacBook Pro M3 Pro, 36GB RAM
• macOS: 15.x (arm64)
• Python: 3.12.11
• PyTorch: 2.9.1 (MPS available, CUDA not available)
• xgrammar: 0.1.29
• vLLM: Built from this branch

Test Results

Test	Status
Unit test (dtype conversion logic)	✅ Passed
xgrammar library test (float32, float16, bfloat16)	✅ Passed
Full integration test with Qwen/Qwen3-0.6B	✅ Passed

Integration Test Output

Prompt: Write a simple hello world in Python:
Output: {"python_code": "print('Hello, World!')"}

The model successfully generated structured JSON output using xgrammar on macOS CPU with float16 dtype.

---

Response to Gemini Code Review

✅ Accepted Gemini's suggestion - Added logits.device.type == "cpu" check to make the fix more targeted.

Analysis

I investigated the xgrammar library and found:

1. xgrammar 0.1.29 (current) already supports float16/bfloat16 on CPU natively
2. The original bug existed in older xgrammar versions that only supported float32
3. The CPU kernel at line 22 in older versions raised ValueError: logits must be of type float32

Why keep this fix?

• Provides backward compatibility with older xgrammar versions
• Defensive coding in case of future regressions
• Minimal overhead with the CPU device check

The updated code now includes the CPU device check as suggested:

if logits.device.type == "cpu" and logits.dtype != torch.float32:

This ensures:

• No unnecessary dtype conversions on GPU inference
• Fix only applies where needed (CPU with non-float32 dtype)

[karanb192]

Hi @DarkLight1337 @ywang96 - could you please review this PR when you have a chance?

This fixes issue #31901 - guided decoding with xgrammar failing on macOS CPU inference due to dtype mismatch. The fix has been tested locally on macOS M3 Pro and all CI checks are passing.

Thank you! 🙏

[LucasWilkinson]

cc @aarnphm @khluu for structured ouputs

[Sberm]

You stole somebody’s code without any acknowledgement, and you belong in the circus, cause you’s a freaking clown.

[Inokinoki]

You stole somebody’s code without any acknowledgement, and you belong in the circus, cause you’s a freaking clown.

Don't take it too hard ;) There is a link in the comment sourcing the issue.

I think it's better to wait for the confirmation and the suggestions from the maintainers, in order to make sure that the fix can generalise.

[karanb192]

Thanks @Inokinoki for clarifying. To be explicit: the fix approach was based on @Inokinoki's suggestion in #31901, which is linked in the PR description ("Fixes #31901"). I've added explicit credit in the PR description to make this clearer.

Happy to discuss the technical merits of the fix with maintainers.

[saattrupdan]

@DarkLight1337 @ywang96, could you please have a look at this short PR? It would enable MacOS support for all our vLLM projects using structured generation, including the EuroEval framework.

[Inokinoki]

Thanks @Inokinoki for clarifying. To be explicit: the fix approach was based on @Inokinoki's suggestion in #31901, which is linked in the PR description ("Fixes #31901"). I've added explicit credit in the PR description to make this clearer.

Happy to discuss the technical merits of the fix with maintainers.

Congrats that this is approved!

I would appreciate it if the maintainer @DarkLight1337 /you would like to add me as a "co-authored-by" in the commit message.
But if this blocks the merge, plz do not do that - just keep me informed, if this quick fix brings any issues.

[DarkLight1337]

I did that already!

[saattrupdan]

@DarkLight1337 With this PR being approved, are the failing tests the reason for not merging? @karanb192 or @Inokinoki, are you able to look at them, if it's needed?

[DarkLight1337]

Rebased the branch, see if the tests pass now