PR #37654: [Feature] Expose xgrammar bitmask backend selection in StructuredOutputsConfig

• Repository: vllm-project/vllm
• Author: chaunceyjiang
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/37654

Description (problem / solution / changelog)

Fix https://github.com/vllm-project/vllm/issues/37650

Purpose

Test Plan

Test Result

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

Changed files

• vllm/config/structured_outputs.py (modified, +8/-0)
• vllm/v1/structured_output/utils.py (modified, +10/-2)
• vllm/v1/worker/gpu_model_runner.py (modified, +6/-1)

PR #1: feat(structured-outputs): expose xgrammar bitmask backend selection via platform interface

• Repository: xyDong0223/vllm
• Author: Copilot
• State: closed | merged: True
• Link: https://github.com/xyDong0223/vllm/pull/1

Description (problem / solution / changelog)

xgr.apply_token_bitmask_inplace supports multiple backends ("auto", "cpu", "cuda", "triton", "torch_compile", "torch_native") but the backend was hardcoded to "auto" with no user control. This adds a platform-aware backend selection mechanism following the existing get_vit_attn_backend pattern.

Changes

• vllm/config/structured_outputs.py: Add BitmaskBackend type alias and bitmask_backend: BitmaskBackend = "auto" field to StructuredOutputsConfig

• vllm/platforms/interface.py: Add get_supported_bitmask_backends() and get_bitmask_backend() to Platform base class; base implementation supports ["auto", "cpu"] and returns "auto" unchanged

• Platform overrides:

  • cuda.py: supports all backends; "auto" delegates to xgrammar
  • cpu.py: supports ["auto", "cpu"]; resolves "auto" → "cpu"
  • rocm.py: supports ["auto", "cpu", "cuda", "torch_compile", "torch_native"]; "auto" delegates to xgrammar
  • xpu.py: supports ["auto", "cpu"]; resolves "auto" → "cpu"

• vllm/v1/structured_output/utils.py: Add bitmask_backend: str = "auto" param to apply_grammar_bitmask(), thread it into both xgr.apply_token_bitmask_inplace() calls

• vllm/v1/worker/gpu_model_runner.py: Resolve backend once at init via current_platform.get_bitmask_backend(config.bitmask_backend), cache as self._bitmask_backend, pass on each call

Usage

# Via Python API
from vllm import LLM, SamplingParams
from vllm.config import StructuredOutputsConfig

llm = LLM(model="...", structured_outputs_config=StructuredOutputsConfig(bitmask_backend="triton"))

# Via CLI
vllm serve ... --structured-outputs-config.bitmask_backend=triton

Note: This only affects the legacy GPU model runner path that calls xgr.apply_token_bitmask_inplace. The new GPU model runner (vllm/v1/worker/gpu/model_runner.py) uses a custom Triton kernel directly and is unaffected.

Feature: Allow user selection of structured output (xgrammar) backend for bitmask application

Implements: https://github.com/vllm-project/vllm/issues/37650

Background

In vLLM, when structured output tasks are enabled, the function xgr.apply_token_bitmask_inplace from xgrammar supports different backends ("auto", "cpu", "cuda", "triton", "torch_compile", "torch_native"), but currently the backend is hard-coded as "auto" and not exposed to the user.

Design

Following @shen-shanshan's suggestion in the issue comments (https://github.com/vllm-project/vllm/issues/37650#issuecomment-4096210982), the implementation adds a Platform interface for various hardware platforms to customize their xgrammar backend selection logic, similar to the existing get_vit_attn_backend pattern in vllm/platforms/interface.py#L243-L278.

Implementation Plan

1. Add bitmask_backend field to StructuredOutputsConfig (vllm/config/structured_outputs.py)

Add a new type alias and field:

BitmaskBackend = Literal[
    "auto", "cpu", "cuda", "triton", "torch_compile", "torch_native"
]

Add to the StructuredOutputsConfig class:

bitmask_backend: BitmaskBackend = "auto"
"""Select the backend for applying the structured output token bitmask
via xgrammar's apply_token_bitmask_inplace. Options: "auto", "cpu",
"cuda", "triton", "torch_compile", "torch_native".
Default "auto" lets the platform decide the best backend."""

This exposes it via CLI as --structured-outputs-config.bitmask_backend=triton and via Python API as StructuredOutputsConfig(bitmask_backend="triton").

2. Add Platform interface methods (vllm/platforms/interface.py)

Add two new classmethods to the Platform base class, following the get_supported_vit_attn_backends / get_vit_attn_backend pattern:

@classmethod
def get_supported_bitmask_backends(cls) -> list[str]:
    """Return the list of supported bitmask backends on this platform."""
    return ["auto", "cpu"]

@classmethod
def get_bitmask_backend(cls, backend: str = "auto") -> str:
    """
    Get the bitmask backend for structured output on this platform.
    If user specifies a backend explicitly, validate it's supported and use it.
    If "auto", let the platform choose the best default.
    """
    if backend != "auto":
        supported = cls.get_supported_bitmask_backends()
        if backend not in supported:
            raise ValueError(
                f"Bitmask backend '{backend}' is not supported on "
                f"{cls.device_name}. Supported: {supported}"
            )
        return backend
    return "auto"

3. Override in platform subclasses

• CudaPlatform (vllm/platforms/cuda.py): Support all backends ["auto", "cpu", "cuda", "triton", "torch_compile", "torch_native"]. Default "auto" lets xgrammar choose.
• CpuPlatform (vllm/platforms/cpu.py): Support ["auto", "cpu"]. When "auto", resolve to "cpu" since only CPU kernels are available.
• RocmPlatform (vllm/platforms/rocm.py): Support ["auto", "cpu", "cuda", "torch_compile", "torch_native"]. Default "auto".
• XPUPlatform (vllm/platforms/xpu.py): Support ["auto", "cpu"]. When "auto", resolve to "cpu".

4. Modify apply_grammar_bitmask in vllm/v1/structured_output/utils.py

Add a bitmask_backend: str = "auto" parameter and pass it through to xgr.apply_token_bitmask_inplace:

def apply_grammar_bitmask(
    scheduler_output: SchedulerOutput,
    grammar_output: GrammarOutput,
    input_batch: InputBatch,
    logits: torch.Tensor,
    bitmask_backend: str = "auto",  # NEW
) -> None:
    # ... existing code ...
    
    if logits.device.type == "cpu" and logits.dtype != torch.float32:
        logits_float32 = logits.to(torch.float32)
        xgr.apply_token_bitmask_inplace(
            logits_float32, grammar_bitmask, indices=index_tensor,
            backend=bitmask_backend,  # NEW
        )
        logits.copy_(logits_float32.to(logits.dtype))
    else:
        xgr.apply_token_bitmask_inplace(
            logits, grammar_bitmask, indices=index_tensor,
            backend=bitmask_backend,  # NEW
        )

5. Modify callers to resolve and pass the backend

In vllm/v1/worker/gpu_model_runner.py (the legacy GPU model runner that calls apply_grammar_bitmask from utils.py):

• During initialization, resolve the bitmask backend using the platform interface:

  from vllm.platforms import current_platform
  self._bitmask_backend = current_platform.get_bitmask_backend(
      self.vllm_config.structured_outputs_config.bitmask_backend
  )

• Pass it when calling:

  if grammar_output is not None:
      apply_grammar_bitmask(
          scheduler_output, grammar_output, self.input_batch, logits,
          bitmask_backend=self._bitmask_backend,
      )

Note: The new GPU model runner path (vllm/v1/worker/gpu/model_runner.py) uses its own...

This pull request was created from Copilot chat.

💬 Send tasks to Copilot coding agent from Slack and Teams to turn conversations into code. Copilot posts an update in your thread when it's finished.

Changed files

• vllm/config/structured_outputs.py (modified, +9/-0)
• vllm/platforms/cpu.py (modified, +20/-0)
• vllm/platforms/cuda.py (modified, +19/-0)
• vllm/platforms/interface.py (modified, +35/-0)
• vllm/platforms/rocm.py (modified, +19/-0)
• vllm/platforms/xpu.py (modified, +20/-0)
• vllm/v1/structured_output/utils.py (modified, +12/-2)
• vllm/v1/worker/gpu_model_runner.py (modified, +7/-1)