This document is relevant for: Trn2, Trn3

Structured Outputs and Tool Calling#

Overview#

Structured outputs and tool calling enable constrained generation on the vLLM Neuron backend using JSON schemas, regex patterns, choice enums, and BNF grammars. The model output is guaranteed to conform to the specified constraint by applying a grammar-derived bitmask to logits before sampling. Tool calling (function calling) builds directly on this same infrastructure vLLM internally converts tool schemas into structured output constraints, so no additional code changes were needed.

This feature is implemented across three layers:

Scheduler (scheduler.py): Computes the grammar bitmask on CPU via vLLM’s StructuredOutputManager and attaches it to SchedulerOutput
Model Runner (neuron_model_runner.py): Unpacks the packed int32 bitmask into a boolean tensor, reorders it to match the batch layout, and passes it into the model forward call
vLLM Neuron Sampler (on-device): Applies logits.masked_fill(~logit_mask, -inf) before argmax so only grammar-valid tokens can be selected

For scheduler background, see neuron-scheduler. For overall vLLM integration, see vllm-integration-design-reference.

Problem Statement#

With on-device sampling enabled, the model forward pass is compiled into a NEFF (Neuron Executable File Format). In this mode, logits are never returned to the CPU. Instead, the device runs the full pipeline from logits through argmax and returns only the sampled token IDs. (Without on-device sampling, logits are returned to CPU for sampling there.)

In standard vLLM (v1, GPU), the bitmask is NOT passed into the model forward call. The model executes inside a CUDA graph and returns logits; then, in a separate step outside the graph, a Triton kernel applies the grammar bitmask on-GPU (logits.masked_fill → -inf) before the sampler runs. The key distinction is that logits are accessible between the model forward and sampling steps.

In our Neuron implementation, structured outputs use the on-device sampling path. Since logits are never materialized outside the compiled NEFF, we cannot intercept them between the forward pass and sampling. The mask must be applied ON device before sampling, so we thread it through model.forward() to get it to the sampler:

Scheduler: attach bitmask to SchedulerOutput
                    ↓
Model Runner: extract bitmask, pass as logit_mask parameter
                    ↓
Model.forward(logit_mask=...) → On-device sampler applies mask

We added the logit_mask parameter because vLLM Neuron does on-device sampling and the mask must be applied ON device before sampling. We had to thread it through model.forward() to get it to the sampler.

Solution#

Data Flow#

The end-to-end flow spans the CPU host and the Neuron device, crossing the NEFF boundary once per decode step:

┌─────────────────────────────────────────────────────────────────┐
│  CPU (Host)                                                     │
├─────────────────────────────────────────────────────────────────┤
│  SCHEDULER (scheduler.py)                                       │
│  1. Request arrives with structured_output (JSON/regex/etc.)    │
│  2. grammar_init(request)  initialize FSM from schema          │
│  3. Each decode step: grammar_bitmask() → packed int32 tensor   │
│  4. Attach bitmask to scheduler_output._grammar_bitmask         │
├─────────────────────────────────────────────────────────────────┤
│  MODEL RUNNER (neuron_model_runner.py)                           │
│  5. _get_grammar_bitmask() extracts from scheduler_output       │
│  6. Unpack: [batch, packed_vocab] int32 → [batch, vocab] bool   │
│  7. Pass logit_mask to model forward                            │
└────────────────────────────┬────────────────────────────────────┘
                             │  Inputs transferred to device:
                             │  - input_ids, positions
                             │  - logit_mask [batch, vocab_size]
                             ▼
═══════════════════════════════════════════════════════════════════
CPU EXECUTION ↑  |  ↓ ON-DEVICE EXECUTION
═══════════════════════════════════════════════════════════════════
                             │
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│  Neuron Device (NEFF)                                           │
├─────────────────────────────────────────────────────────────────┤
│  MODEL (e.g. llama3/model.py)                                   │
│  8. Forward pass computes logits                                │
│  9. Pass logit_mask to sampler                                  │
├─────────────────────────────────────────────────────────────────┤
│  SAMPLER (nn/sampler.py)                                        │
│  10. logits.masked_fill(~logit_mask, -inf)                      │
│  11. argmax → only grammar-valid tokens can be selected         │
└────────────────────────────┬────────────────────────────────────┘
                             │  Output: token_ids [batch]
                             ▼
═══════════════════════════════════════════════════════════════════
ON-DEVICE EXECUTION ↑  |  ↓ CPU EXECUTION
═══════════════════════════════════════════════════════════════════
                             │
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│  CPU (Host)                                                     │
│  12. Update grammar FSM state with selected token               │
│  13. Next iteration: goto step 3 (new bitmask for new state)    │
└─────────────────────────────────────────────────────────────────┘

Per decode step the cycle is: forward pass → get bitmask → apply mask → sample → update FSM state → repeat.

Bitmask Format#

vLLM constructs a packed bitmask to minimize memory transfer between CPU and device.

Shape: [batch_size, packed_vocab_size] where packed_vocab_size = ceil(vocab_size / 32)
Dtype: int32
Each int32 element encodes 32 vocabulary tokens. Bit = 1 means token is ALLOWED, bit = 0 means DISALLOWED (little-endian bit order).

For Llama-3 with vocab_size = 128,256, this yields packed_vocab_size = 4,008 int32 values per request roughly 16 KB per row.

Scheduler Integration#

In ContinuousBatchingNeuronScheduler.schedule(), after padding is applied (Step 6), the grammar bitmask is computed and attached:

# Step 7: Attach grammar bitmask for structured outputs
grammar_output = self.get_grammar_bitmask(scheduler_output)
if grammar_output is not None and grammar_output.grammar_bitmask is not None:
    scheduler_output._grammar_bitmask = grammar_output.grammar_bitmask
    scheduler_output._structured_output_request_ids = (
        grammar_output.structured_output_request_ids
    )

get_grammar_bitmask() is inherited from vLLM’s base Scheduler class. It internally calls StructuredOutputManager to compute the FSM-based bitmask. The bitmask is kept as a numpy array at this stage conversion to torch happens in the model runner.

Model Runner Integration#

_get_grammar_bitmask() in neuron_model_runner.py performs three operations:

1. Reordering The compact bitmask from the scheduler only contains rows for requests with structured output constraints. The model runner expands this to match the full batch layout (including non-SO requests and speculative decode token offsets). Non-SO rows are filled with -1 (all bits set = all tokens allowed).

# Build mapping: req_id -> logit_index (accounting for spec tokens)
for batch_index, req_id in enumerate(self.input_batch.req_ids):
    logit_index = batch_index + cumulative_offset
    cumulative_offset += len(spec_tokens.get(req_id, ()))
    if req_id in struct_out_req_ids_set:
        struct_out_req_batch_indices[req_id] = logit_index

# Full bitmask: -1 (all allowed) for non-SO rows
sorted_bitmask = torch.full(
    (num_logit_rows, packed_vocab_size), fill_value=-1, dtype=torch.int32
)

2. Unpacking The packed int32 bitmask is unpacked to a boolean tensor of shape [num_logit_rows, vocab_size]:

packed_uint = sorted_bitmask.view(num_logit_rows, packed_vocab_size, 1)
bit_positions = torch.arange(32, dtype=torch.int32)
unpacked = ((packed_uint >> bit_positions) & 1).view(num_logit_rows, -1)
unpacked = unpacked[:, :self.vocab_size]
bitmask = unpacked.bool()

3. Device Transfer The boolean mask is moved to the Neuron device and passed as logit_mask to the model forward call:

bitmask = bitmask.to(device=self.device)

# In execute_model():
model_kwargs["logit_mask"] = logit_mask
model_output = self.model(**model_kwargs)

Note: TP sharding of the mask is handled inside the vLLM Neuron sampler, matching the lm_head sharding strategy. The model runner passes the full-vocab mask.

Benchmarking#

Benchmarked on Llama-3-8B with TP8, max_model_len=256, on-device greedy sampling (all_greedy=true), using the xgrammar backend. Results averaged over 10 runs.

Per-Token Overhead#

Configuration	Per-Token Latency	vs Baseline
Simple JSON SO	759.7 ms/token	+1.8%
Complex JSON SO	739.0 ms/token	−1.0%
No SO (baseline)	746.4 ms/token

The ~1-2% variation is within measurement noise. Structured output adds virtually no per-token overhead.

Component Breakdown (Warm, Steady-State)#

Component	Latency
Scheduler: grammar bitmask computation	~0.12 ms
Model runner: bitmask unpacking + transfer	~0.24 ms
Forward pass delta (with vs without mask)	~0 ms (within variance)
Total per-token SO overhead	~0.4 ms

Forward Pass Execution#

Phase	With Mask	Without Mask	Delta
1st (NEFF compile)	8,395 ms	8,117 ms	+278 ms
2nd (warmup)	5,372 ms	5,361 ms	+11 ms
3rd+ (steady)	~750 ms	~762 ms	−12 ms

The first structured output request triggers a NEFF recompilation (~8.4s one-time cost) because the model graph changes to accept the logit_mask input tensor. After warmup, the forward pass with mask is within noise of the baseline.

Tool Calling#

Tool calling (function calling) works on top of the structured outputs infrastructure. No additional code changes were needed vLLM internally converts tool schemas into structured output constraints.

How It Works#

Tool calling in vLLM has two distinct paths depending on tool_choice:

tool_choice
    |
+------------+------------+
|                         |
auto / none              required / named
|                         |
No structured outputs       adjust_request() injects
Model generates free text   JSON schema into request
|                         |
Llama3JsonToolParser        xgrammar compiles grammar FSM
extracts tool calls         logit_mask constrains each token
via regex post-gen                    |
|               vLLM Neuron sampler applies:
|               logits.masked_fill(~mask, -inf)
|                         |
+------------+------------+
    |
tool_calls returned

auto/none: No structured outputs involved. The model generates free text, and Llama3JsonToolParser extracts tool calls post-generation using regex + JSON parsing.
required/named: vLLM’s adjust_request() converts the tool schema into StructuredOutputsParams(json=tool_array_schema). From that point, it flows through the same grammar bitmask pipeline as any structured output request no modification needed.

This works because vLLM internally treats tool_choice=required identically to a structured output request. adjust_request() sets request.structured_outputs = StructuredOutputsParams(json=tool_array_schema).

Why On-Device Sampling Is Required for required/named#

The vLLM Neuron model’s forward() has two code paths:

# Path A: No on-device sampling → returns raw logits
if self.on_device_sampling_config is None:
    return logits  # logit_mask is IGNORED here

# Path B: On-device sampling → sampler applies mask
sampled_tokens = self.sampler(logits, sampling_params, logit_mask=logit_mask)

Without on-device sampling, the model returns raw logits before reaching the sampler. The logit_mask is passed into forward() but never applied unconstrained tokens violate the grammar FSM and cause Failed to advance FSM errors.

Mode	CPU Sampling	On-Device Sampling	Why
`tool_choice="auto"`	✅	✅	No SO needed
`tool_choice="none"`	✅	✅	No SO needed
`tool_choice="required"`	❌	✅	Needs `logit_mask` applied in sampler
Named tool choice	❌	✅	Needs `logit_mask` applied in sampler

Avoiding Recompilation in SO-Enabled Mode#

When structured-output support is enabled, logit_mask is always passed as a tensor to the model forward call, never None. This ensures torch.compile traces a single code path during warmup, and no recompilation is needed when the first structured output request arrives at serving time.

When structured-output support is enabled, warmup passes an all-True mask (all tokens allowed):

# In warmup_prefill() and warmup_decode():
dummy_logit_mask = torch.ones(
    num_reqs, self.vocab_size, dtype=torch.bool, device=self.device
)
_ = self.model(..., logit_mask=dummy_logit_mask)

By default, vLLM Neuron optimizes for SO-off traffic and does not create a dummy logit mask. This avoids the no-op mask allocation / transfer / sampling overhead for deployments that do not serve structured-output requests.

For servers that need to accept structured-output requests, enable SO support explicitly:

{
  "neuron_config": {
    "enable_structured_outputs": true
  }
}

In that mode, if no structured output requests are active, an all-True mask is created as a no-op:

# In execute_model():
logit_mask = self._get_grammar_bitmask(scheduler_output)
if logit_mask is None:
    logit_mask = torch.ones(
        num_logit_rows, self.vocab_size, dtype=torch.bool, device=self.device
    )

An all-True mask is mathematically a no-op: logits.masked_fill(~True, -inf) fills nothing. This avoids any if logit_mask is not None branches that would cause torch.compile to trace a different graph.

In SO-enabled mixed-traffic mode, the vLLM Neuron sampler always receives a mask tensor, so masked_fill stays on the traced path:

# In Sampler.forward():
logits = logits.masked_fill(~logit_mask, float('-inf'))

This “always-pass, never-None” behavior keeps a mixed structured-output / non-structured-output server on one stable tensor path and avoids recompilation when the first structured-output request arrives.

For SO-off-only deployments, this no-op mask is skipped by default, or can be made explicit by setting:

{
  "neuron_config": {
    "enable_structured_outputs": false
  }
}

This is a performance mode for servers that will not accept structured-output requests. If a request does include structured_outputs while this mode is enabled, vLLM Neuron rejects it with an actionable error asking the user to restart the server with enable_structured_outputs=true. It does not silently ignore the schema request, because returning an unconstrained answer would violate the client’s requested contract.

Limitations#

On-device sampling required Structured outputs are only supported with on-device sampling enabled. Attempting to use them with CPU sampling raises NotImplementedError.

FAQ#

Q: How does vLLM pass the bitmask? Is it through logit_mask?

In standard vLLM, the bitmask is NOT passed into the model. It’s applied after the model returns logits, during CPU sampling. The grammar/bitmask logic lives in the sampler layer, not the model layer. We added the logit_mask parameter because vLLM Neuron does on-device sampling and the mask must be applied ON device before sampling. We had to thread it through model.forward() to get it to the sampler.

Q: Why is the scheduler attaching the bitmask?

The scheduler has access to StructuredOutputManager, which stores the grammar and computes the bitmask. We attach _grammar_bitmask to SchedulerOutput because that’s the data structure that flows from scheduler to model_runner. The model runner does not have access to StructuredOutputManager.

Q: Why not use SamplingMetadata or create a NeuronSamplingMetadata instead?

SamplingMetadata is created inside vLLM’s InputBatch, which we don’t subclass. To add grammar_bitmask there, we’d need to subclass InputBatch or patch it after creation. SchedulerOutput is simpler since we already override the scheduler and it’s the known data channel to model_runner. Even if we defined NeuronSamplingMetadata, we’d still need all the places that read metadata to understand it (or at least to carry an extra field through). That’s a core vLLM change, not just plugin code.

Q: But runner and scheduler are the same process why serialize through SchedulerOutput?

They are in the same process, but the scheduler and model runner are separate components with a well-defined interface (SchedulerOutput). Using this existing data channel keeps the design clean and avoids coupling the model runner to the StructuredOutputManager internals.

Structured Outputs and Tool Calling

Contents