xgrammar¶

Classes:

`CompiledGrammar`()	This is the primary object to store compiled grammar.
`Grammar`()	This class represents a grammar object in XGrammar, and can be used later in the grammar-guided generation.
`GrammarCompiler`(tokenizer_info, *[, ...])	The compiler for grammars.
`GrammarMatcher`(compiled_grammar, *[, ...])	Match the output of the LLM to the specified grammar, then generate the mask for the next token.
`StructuralTagItem`(*, begin, schema, end)	A structural tag item.
`TokenizerInfo`(encoded_vocab[, vocab_type, ...])	The tokenizer info contains the vocabulary, the type of the vocabulary, and necessary information for the grammar-guided generation.
`VocabType`(value[, names, module, qualname, ...])	The type of the vocabulary.

Functions:

`allocate_token_bitmask`(batch_size, vocab_size)	Allocate the bitmask for the next token prediction.
`apply_token_bitmask_inplace`(logits, bitmask, *)	Apply the bitmask to the logits in-place.
`get_bitmask_shape`(batch_size, vocab_size)	Return the shape of the bitmask: (batch_size, ceil(vocab_size / 32)).
`reset_token_bitmask`(bitmask)	Reset the bitmask to the full mask.

class xgrammar.CompiledGrammar¶

This is the primary object to store compiled grammar.

A CompiledGrammar can be used to construct GrammarMatcher to generate token masks efficiently.

Note

Do not construct this class directly, instead use GrammarCompiler to construct the object.

Attributes:

`grammar`	The original grammar.
`memory_size_bytes`	The approximate memory usage of the compiled grammar in bytes.
`tokenizer_info`	The tokenizer info associated with the compiled grammar.

property grammar: Grammar¶: The original grammar.

property memory_size_bytes: int¶: The approximate memory usage of the compiled grammar in bytes.

property tokenizer_info: TokenizerInfo¶: The tokenizer info associated with the compiled grammar.

class xgrammar.Grammar¶

This class represents a grammar object in XGrammar, and can be used later in the grammar-guided generation.

The Grammar object supports context-free grammar (CFG). EBNF (extended Backus-Naur Form) is used as the format of the grammar. There are many specifications for EBNF in the literature, and we follow the specification of GBNF (GGML BNF) in https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md.

When printed, the grammar will be converted to GBNF format.

Methods:

`builtin_json_grammar`()	Get the grammar of standard JSON.
`concat`(*grammars)	Create a grammar that matches the concatenation of the grammars in the list.
`from_ebnf`(ebnf_string, *[, root_rule_name])	Construct a grammar from EBNF string.
`from_json_schema`(schema, *[, ...])	Construct a grammar from JSON schema.
`from_regex`(regex_string, *[, ...])	Create a grammar from a regular expression string.
`from_structural_tag`(tags, triggers)	Create a grammar from structural tags.
`union`(*grammars)	Create a grammar that matches any of the grammars in the list.

static builtin_json_grammar() → Grammar¶

Get the grammar of standard JSON. This is compatible with the official JSON grammar specification in https://www.json.org/json-en.html.

Returns:: grammar – The JSON grammar.
Return type:: Grammar

static concat(*grammars: Grammar) → Grammar¶

Create a grammar that matches the concatenation of the grammars in the list. That is equivalent to using the + operator to concatenate the grammars in the list.

Parameters:: grammars (List[Grammar]) – The grammars to create the concatenation of.
Returns:: grammar – The concatenation of the grammars.
Return type:: Grammar

static from_ebnf(ebnf_string: str, *, root_rule_name: str = 'root') → Grammar¶

Construct a grammar from EBNF string. The EBNF string should follow the format in https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md.

Parameters:

ebnf_string (str) – The grammar string in EBNF format.
root_rule_name (str, default: "root") – The name of the root rule in the grammar.

Raises:

RuntimeError – When converting the regex pattern fails, with details about the parsing error.

static from_json_schema(schema: Union[str, Type[BaseModel], Dict[str, Any]], *, any_whitespace: bool = True, indent: Optional[int] = None, separators: Optional[Tuple[str, str]] = None, strict_mode: bool = True, print_converted_ebnf: bool = False) → Grammar¶

Construct a grammar from JSON schema. Pydantic model or JSON schema string can be used to specify the schema.

It allows any whitespace by default. If user want to specify the format of the JSON, set any_whitespace to False and use the indent and separators parameters. The meaning and the default values of the parameters follows the convention in json.dumps().

It internally converts the JSON schema to a EBNF grammar.

Parameters:

schema (Union[str, Type[BaseModel], Dict[str, Any]]) – The schema string or Pydantic model or JSON schema dict.
any_whitespace (bool, default: True) – Whether to use any whitespace. If True, the generated grammar will ignore the indent and separators parameters, and allow any whitespace.
indent (Optional[int], default: None) –
The number of spaces for indentation. If None, the output will be in one line.

Note that specifying the indentation means forcing the LLM to generate JSON strings strictly formatted. However, some models may tend to generate JSON strings that are not strictly formatted. In this case, forcing the LLM to generate strictly formatted JSON strings may degrade the generation quality. See <https://github.com/sgl-project/sglang/issues/2216#issuecomment-2516192009> for more details.
separators (Optional[Tuple[str, str]], default: None) – Two separators used in the schema: comma and colon. Examples: (“,”, “:”), (”, “, “: “). If None, the default separators will be used: (“,”, “: “) when the indent is not None, and (”, “, “: “) otherwise.
strict_mode (bool, default: True) –
Whether to use strict mode. In strict mode, the generated grammar will not allow properties and items that is not specified in the schema. This is equivalent to setting unevaluatedProperties and unevaluatedItems to false.

This helps LLM to generate accurate output in the grammar-guided generation with JSON schema.
print_converted_ebnf (bool, default: False) – If True, the converted EBNF string will be printed. For debugging purposes.

Returns:

grammar – The constructed grammar.

Return type:

Grammar

Raises:

RuntimeError – When converting the json schema fails, with details about the parsing error.

static from_regex(regex_string: str, *, print_converted_ebnf: bool = False) → Grammar¶

Create a grammar from a regular expression string.

Parameters:

regex_string (str) – The regular expression pattern to create the grammar from.
print_converted_ebnf (bool, default: False) – This method will convert the regex pattern to EBNF first. If this is true, the converted EBNF string will be printed. For debugging purposes. Default: False.

Returns:

grammar – The constructed grammar from the regex pattern.

Return type:

Grammar

Raises:

RuntimeError – When parsing the regex pattern fails, with details about the parsing error.

static from_structural_tag(tags: List[StructuralTagItem], triggers: List[str]) → Grammar¶

Create a grammar from structural tags. The structural tag handles the dispatching of different grammars based on the tags and triggers: it initially allows any output, until a trigger is encountered, then dispatch to the corresponding tag; when the end tag is encountered, the grammar will allow any following output, until the next trigger is encountered.

The tags parameter is used to specify the output pattern. It is especially useful for LLM function calling, where the pattern is: <function=func_name>{“arg1”: …, “arg2”: …}</function>. This pattern consists of three parts: a begin tag (<function=func_name>), a parameter list according to some schema ({“arg1”: …, “arg2”: …}), and an end tag (</function>). This pattern can be described in a StructuralTagItem with a begin tag, a schema, and an end tag. The structural tag is able to handle multiple such patterns by passing them into multiple tags.

The triggers parameter is used to trigger the dispatching of different grammars. The trigger should be a prefix of a provided begin tag. When the trigger is encountered, the corresponding tag should be used to constrain the following output. There can be multiple tags matching the same trigger. Then if the trigger is encountered, the following output should match one of the tags. For example, in function calling, the triggers can be [“<function=”]. Then if “<function=” is encountered, the following output must match one of the tags (e.g. <function=get_weather>{“city”: “Beijing”}</function>).

The corrrespondence of tags and triggers is automatically determined: all tags with the same trigger will be grouped together. User should make sure any trigger is not a prefix of another trigger: then the corrrespondence of tags and triggers will be ambiguous.

To use this grammar in grammar-guided generation, the GrammarMatcher constructed from structural tag will generate a mask for each token. When the trigger is not encountered, the mask will likely be all-1 and not have to be used (fill_next_token_bitmask returns False, meaning no token is masked). When a trigger is encountered, the mask should be enforced (fill_next_token_bitmask will return True, meaning some token is masked) to the output logits.

The benefit of this method is the token boundary between tags and triggers is automatically handled. The user does not need to worry about the token boundary.

Parameters:

tags (List[StructuralTagItem]) – The structural tags.
triggers (List[str]) – The triggers.

Examples

>>> class Schema1(BaseModel):
>>>     arg1: str
>>>     arg2: int
>>> class Schema2(BaseModel):
>>>     arg3: float
>>>     arg4: List[str]
>>> tags = [
>>>     StructuralTagItem(begin="<function=f>", schema=Schema1, end="</function>"),
>>>     StructuralTagItem(begin="<function=g>", schema=Schema2, end="</function>"),
>>> ]
>>> triggers = ["<function="]
>>> grammar = Grammar.from_structural_tag(tags, triggers)

static union(*grammars: Grammar) → Grammar¶

Create a grammar that matches any of the grammars in the list. That is equivalent to using the | operator to concatenate the grammars in the list.

Parameters:: grammars (List[Grammar]) – The grammars to create the union of.
Returns:: grammar – The union of the grammars.
Return type:: Grammar

class xgrammar.GrammarCompiler(tokenizer_info: TokenizerInfo, *, max_threads: int = 8, cache_enabled: bool = True, cache_limit_bytes: int = -1)¶

The compiler for grammars. It is associated with a certain tokenizer info, and compiles grammars into CompiledGrammar with the tokenizer info. It allows parallel compilation with multiple threads, and has a cache to store the compilation result, avoiding compiling the same grammar multiple times.

Parameters:

tokenizer_info (TokenizerInfo) – The tokenizer info.
max_threads (int, default: 8) – The maximum number of threads used to compile the grammar.
cache_enabled (bool, default: True) – Whether to enable the cache.
cache_limit_bytes (int, default: -1) – The maximum memory usage for the cache in the specified unit. Note that the actual memory usage may slightly exceed this value.

Attributes:

cache_limit_bytes

The maximum memory usage for the cache in bytes.

Methods:

`clear_cache`()	Clear all cached compiled grammars.
`compile_builtin_json_grammar`()	Get CompiledGrammar from the standard JSON.
`compile_json_schema`(schema, *[, ...])	Get CompiledGrammar from the specified JSON schema and format.
`compile_regex`(regex)	Get CompiledGrammar from the specified regex.
`compile_structural_tag`(tags, triggers)	Compile a grammar from structural tags.
`get_cache_size_bytes`()	The approximate memory usage of the cache in bytes.

property cache_limit_bytes: int¶: The maximum memory usage for the cache in bytes. Returns -1 if the cache has no memory limit.

clear_cache() → None¶: Clear all cached compiled grammars.

compile_builtin_json_grammar() → CompiledGrammar¶

Get CompiledGrammar from the standard JSON.

Returns:: compiled_grammar – The compiled grammar.
Return type:: CompiledGrammar

compile_json_schema(schema: Union[str, Type[BaseModel], Dict[str, Any]], *, any_whitespace: bool = True, indent: Optional[int] = None, separators: Optional[Tuple[str, str]] = None, strict_mode: bool = True) → CompiledGrammar¶

Get CompiledGrammar from the specified JSON schema and format. The indent and separators parameters follow the same convention as in json.dumps().

Parameters:

schema (Union[str, Type[BaseModel], Dict[str, Any]]) – The schema string or Pydantic model or JSON schema dict.
indent (Optional[int], default: None) – The number of spaces for indentation. If None, the output will be in one line.
separators (Optional[Tuple[str, str]], default: None) – Two separators used in the schema: comma and colon. Examples: (“,”, “:”), (”, “, “: “). If None, the default separators will be used: (“,”, “: “) when the indent is not None, and (”, “, “: “) otherwise.
strict_mode (bool, default: True) –
Whether to use strict mode. In strict mode, the generated grammar will not allow properties and items that is not specified in the schema. This is equivalent to setting unevaluatedProperties and unevaluatedItems to false.

This helps LLM to generate accurate output in the grammar-guided generation with JSON schema.

Returns:

compiled_grammar – The compiled grammar.

Return type:

CompiledGrammar

compile_regex(regex: str) → CompiledGrammar¶

Get CompiledGrammar from the specified regex.

Parameters:: regex (str) – The regex string.
Returns:: compiled_grammar – The compiled grammar.
Return type:: CompiledGrammar

compile_structural_tag(tags: List[StructuralTagItem], triggers: List[str]) → CompiledGrammar¶

Compile a grammar from structural tags. See Grammar.from_structural_tag() for more details.

Parameters:

tags (List[StructuralTagItem]) – The structural tags.
triggers (List[str]) – The triggers.

Returns:

compiled_grammar – The compiled grammar.

Return type:

CompiledGrammar

get_cache_size_bytes() → int¶: The approximate memory usage of the cache in bytes.

class xgrammar.GrammarMatcher(compiled_grammar: CompiledGrammar, *, override_stop_tokens: Optional[Union[List[int], int]] = None, terminate_without_stop_token: bool = False, max_rollback_tokens: int = 0)¶

Match the output of the LLM to the specified grammar, then generate the mask for the next token. This is the core class in the grammar-guided generation.

This class maintains a stateful matcher that can accept tokens and strings, then match them to the specified grammar. The matcher can provide a bitmask for the next token prediction, so that the output of the LLM follows the specified grammar. Its state can be reset and rolled back by tokens. It also provides utilities for jump-forward decoding.

After matching the whole grammar, the matcher will accept a stop token. The token mask at this time will only allow stop tokens. After accepting the stop token, the matcher will terminate, then it cannot accept any new token or generate a new token mask, meaning the generation is finished.

Under the hood, it utilizes a pushdown automaton with backtracking to match the grammar, with optimizations specific to LLM token mask generation.

Parameters:

compiled_grammar (CompiledGrammar) – The initialization context for the grammar matcher.
override_stop_tokens (Optional[Union[int, List[int]]], default: None) – If not None, the stop tokens to override the ones in the grammar.
terminate_without_stop_token (bool, default: False) – Whether to terminate the matcher without accepting a stop token.
max_rollback_tokens (int, default: 0) – The maximum number of rollback tokens allowed. The rollback operation is useful for jump-forward decoding and speculative decoding.

Methods:

`accept_token`(token_id, *[, debug_print])	Accept one token and update the state of the matcher.
`fill_next_token_bitmask`(bitmask[, index, ...])	Fill the bitmask for the next token prediction.
`find_jump_forward_string`()	Find the jump-forward string for jump-forward decoding.
`is_terminated`()	Check if the matcher has terminated.
`reset`()	Reset the matcher to the initial state.
`rollback`([num_tokens])	Rollback the matcher to a previous state by several tokens.

Attributes:

`max_rollback_tokens`	Get the maximum number of rollback tokens allowed.
`stop_token_ids`	The ids of the stop tokens used in the matcher.

accept_token(token_id: int, *, debug_print: bool = False) → bool¶

Accept one token and update the state of the matcher.

Parameters:

token_id (int) – The id of the token to accept.
debug_print (bool, default: False) – Whether to print information about the internal state of the matcher. Helpful for debugging.

Returns:

accepted – Whether the token is accepted.

Return type:

bool

fill_next_token_bitmask(bitmask: Tensor, index: int = 0, *, debug_print: bool = False) → bool¶

Fill the bitmask for the next token prediction. The input bitmask can be generated by allocate_token_bitmask, and must be on CPU. bitmask[index] will be filled with the next token bitmask.

This method does not change the matcher state.

Parameters:

bitmask (torch.Tensor) – The bitmask for the next token prediction.
index (int, default: 0) – The batch id of the bitmask.
debug_print (bool, default: False) – Whether to print information about generated bitmask. Helpful for debugging.

Returns:

need_apply – Whether the bitmask need to be applied (not all-true). An optimization: if False, this means the bitmask is already all-true, so no need to apply it.

Return type:

bool

find_jump_forward_string() → str¶

Find the jump-forward string for jump-forward decoding. This is the longest string that certainly conforms with the current grammar from the current matcher state. This string can become the output of the LLM without requiring LLM decoding.

This method does not change the matcher state.

Returns:: jump_forward_string – The jump-forward string.
Return type:: str

is_terminated() → bool¶

Check if the matcher has terminated. If terminate_without_stop_token is False, the matcher will terminate if it has accepted the stop token. Otherwise, the matcher will terminate after matching the whole grammar.

Returns:: terminated – Whether the matcher has terminated.
Return type:: bool

property max_rollback_tokens: int¶

Get the maximum number of rollback tokens allowed.

Returns:: max_rollback_tokens – The maximum number of rollback tokens.
Return type:: int

reset() → None¶: Reset the matcher to the initial state.

rollback(num_tokens: int = 1) → None¶

Rollback the matcher to a previous state by several tokens.

Parameters:: num_tokens (int, default: 1) – The number of tokens to rollback. It cannot exceed the current number of steps, nor can it exceed the specified maximum number of rollback tokens.

property stop_token_ids: List[int]¶

The ids of the stop tokens used in the matcher. If specified, the provided stop tokens will be used. Otherwise, the stop tokens will be detected from the vocabulary.

Returns:: stop_token_ids – The ids of the stop tokens.
Return type:: List[int]

class xgrammar.StructuralTagItem(*, begin: str, schema: Union[str, Type[BaseModel], Dict[str, Any]], end: str)¶

A structural tag item. See Grammar.from_structural_tag() for more details.

begin¶

The begin tag.

Type:: str

schema_¶

The schema.

Type:: Union[str, Type[BaseModel]]

end¶

The end tag.

Type:: str

Attributes:

model_config

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_config: ClassVar[ConfigDict] = {}¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class xgrammar.TokenizerInfo(encoded_vocab: Union[List[bytes], List[str]], vocab_type: VocabType = VocabType.RAW, *, vocab_size: Optional[int] = None, stop_token_ids: Optional[Union[List[int], int]] = None, add_prefix_space: bool = False)¶

The tokenizer info contains the vocabulary, the type of the vocabulary, and necessary information for the grammar-guided generation.

Note that although some tokenizers will encode the tokens in a special format, e.g. “<0x1B>” for “” in the ByteFallback tokenizer, and “Ġ” for ” ” in the Byte-Level BPE tokenizer, TokenizerInfo always decodes the vocabulary to the original format (e.g. “” and ” “).

Also note that some models (e.g. Phi-3 and Deepseek-V2) may pad the vocabulary to a multiple of 32. In this case, the model’s vocab_size is larger than the tokenizer’s vocabulary size. Please pass the model’s vocab_size to the vocab_size parameter in the constructor, because this information is used to determine the size of the token mask.

Parameters:

encoded_vocab (Union[List[bytes], List[str]]) – The encoded vocabulary of the tokenizer.
vocab_type (VocabType, default: VocabType.RAW) – The type of the vocabulary. See also VocabType.
vocab_size (Optional[int], default: None) – The size of the vocabulary. If not provided, the vocabulary size will be len(encoded_vocab).
stop_token_ids (Optional[List[int]], default: None) – The stop token ids. If not provided, the stop token ids will be auto detected (but may not be correct).
add_prefix_space (bool, default: False) – Whether the tokenizer will prepend a space before the text in the tokenization process.

Attributes:

`add_prefix_space`	Whether the tokenizer will prepend a space before the text in the tokenization process.
`decoded_vocab`	The decoded vocabulary of the tokenizer.
`prepend_space_in_tokenization`	Whether the tokenizer will prepend a space before the text in the tokenization process.
`special_token_ids`	The special token ids.
`stop_token_ids`	The stop token ids.
`vocab_size`	The size of the vocabulary.
`vocab_type`	The type of the vocabulary.

Methods:

`dump_metadata`()	Dump the metadata of the tokenizer to a json string.
`from_huggingface`(tokenizer, *[, vocab_size, ...])	Construct the tokenizer info from the huggingface tokenizer.
`from_vocab_and_metadata`(encoded_vocab, metadata)	Construct the tokenizer info from the vocabulary and the metadata string in json format.

property add_prefix_space: bool¶: Whether the tokenizer will prepend a space before the text in the tokenization process.

property decoded_vocab: List[bytes]¶: The decoded vocabulary of the tokenizer. This converts the tokens in the LLM’s vocabulary back to the original format of the input text. E.g. for type ByteFallback, the token <0x1B> is converted back to “”.

dump_metadata() → str¶: Dump the metadata of the tokenizer to a json string. It can be used to construct the tokenizer info from the vocabulary and the metadata string.

static from_huggingface(tokenizer: PreTrainedTokenizerBase, *, vocab_size: Optional[int] = None, stop_token_ids: Optional[Union[List[int], int]] = None) → TokenizerInfo¶

Construct the tokenizer info from the huggingface tokenizer. This constructor supports various tokenizer backends, including the huggingface fast tokenizer and tiktoken tokenizer. Necessary information is automatically detected from the tokenizer.

The vocab_size parameter is introduced to handle the misalignment between the model’s vocab_size and the tokenizer’s vocabulary size. User should pass the model’s vocab_size (could be defined in the model config) here. See docs of vocab_size for more details.

The stop token ids is by default the eos_token_id of the tokenizer. If there are other stop tokens, you can specify them manually.

Parameters:

tokenizer (PreTrainedTokenizerBase) – The huggingface tokenizer.
vocab_size (Optional[int], default: None) –
The vocabulary size defined by the model (not the tokenizer). This equals to the vocab dimention of the model’s lm_head. This is the size of the token mask.

It can be: 1. the same as the tokenizer’s vocabulary size. This is the most common case. 2. larger than the tokenizer’s vocabulary size. This happens when the model has padding

to lm_head, possibly due to aligning lm_head to the power of 2. E.g. Phi-3 and Deepseek-V2.
1. smaller than the tokenizer’s vocabulary size. This happens when the tokenizer has some added tokens that will not supported by the model. E.g. Llama-3.2 Vision and Molmo-72B-0924 has padded <|image|> tokens, but they will not be considered in lm_head or generated by the model.
model_vocab_size need to be provided for case 2 and 3. If not provided, it will be set to the tokenizer’s vocabulary size.
stop_token_ids (Optional[List[int]], default: None) – The stop token ids. If not provided, the eos_token_id of the tokenizer will be used.

Returns:

tokenizer_info – The tokenizer info.

Return type:

TokenizerInfo

static from_vocab_and_metadata(encoded_vocab: List[Union[bytes, str]], metadata: str) → TokenizerInfo¶

Construct the tokenizer info from the vocabulary and the metadata string in json format.

Parameters:

encoded_vocab (List[Union[bytes, str]]) – The encoded vocabulary of the tokenizer.
metadata (str) – The metadata string in json format.

property prepend_space_in_tokenization: bool¶

Whether the tokenizer will prepend a space before the text in the tokenization process.

This property is deprecated. Use add_prefix_space instead.

property special_token_ids: List[int]¶: The special token ids. Special tokens include control tokens, reserved tokens, padded tokens, etc. Now it is automatically detected from the vocabulary.

property stop_token_ids: List[int]¶: The stop token ids.

property vocab_size: int¶: The size of the vocabulary.

property vocab_type: VocabType¶: The type of the vocabulary.

class xgrammar.VocabType(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)¶

The type of the vocabulary. Used in TokenizerInfo. XGrammar supports three types of vocabularies:

RAW

The vocabulary is in the raw format. The tokens in the vocabulary are kept in their original form without any processing. This kind of tokenizer includes the tiktoken tokenizer, e.g. microsoft/Phi-3-small-8k-instruct, Qwen/Qwen-7B-Chat, etc.

BYTE_FALLBACK

The vocabulary used in the byte fallback BPE tokenizer. The tokens are encoded through the byte-fallback conversion. E.g. “” -> “<0x1B>”, ” apple” -> “▁apple”. This kind of tokenizer includes meta-llama/Llama-2-7b-chat, microsoft/Phi-3.5-mini-instruct, etc.

BYTE_LEVEL

The vocabulary used in the byte level BPE tokenizer. The tokens are encoded through the byte-to-unicode conversion, as in https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59

This kind of tokenizer includes meta-llama/Meta-Llama-3-8B-Instruct, meta-llama/Meta-Llama-3.1-8B-Instruct, etc.

xgrammar.allocate_token_bitmask(batch_size: int, vocab_size: int) → Tensor¶

Allocate the bitmask for the next token prediction. The bitmask is an int32 tensor on CPU with shape (batch_size, ceil(vocab_size / 32)). Users who have their own needs to manage CUDA memory can construct the tensor with get_bitmask_shape and bitmask_dtype themselves.

The reason why we use int32 instead of uint32 is that old versions of PyTorch do not support uint32.

Parameters:

batch_size (int) – The batch size of the bitmask.
vocab_size (int) – The size of the vocabulary.

Returns:

bitmask – The shape of the bitmask.

Return type:

torch.Tensor

xgrammar.apply_token_bitmask_inplace(logits: Tensor, bitmask: Tensor, *, vocab_size: Optional[int] = None, indices: Optional[List[int]] = None) → None¶

Apply the bitmask to the logits in-place. The bitmask is a 01 bitwise compressed tensor, where 0 means the token is masked and 1 means the token is not masked. It can be generated by allocate_token_bitmask and filled by fill_next_token_bitmask. After applying the bitmask, the masked logits will be set to -inf.

The shape of logits and bitmask should be (batch_size, vocab_size) and (batch_size, bitmask_size) respectively. bitmask_size = ceil(vocab_size / 32). The operation is:

for i in range(batch_size):
    for j in range(vocab_size):
        if get_bitmask_value(bitmask, i, j) == 0:
            logits[i, j] = -inf

get_bitmask_value(bitmask, i, j) gets the j-th bit of the i-th row of the bitmask.

## Padding

This method allows additional padding on the vocabulary dimension of logits or bitmask. If padding exists, provide the real vocab size to the vocab_size parameter, and the operation will be applied to logits[…, :vocab_size] and bitmask[…, :ceil(vocab_size / 32)].

If vocab_size is not provided, the vocab size will be detected as min(logits.shape[-1], bitmask.shape[-1] * 32).

## Indices

Indices can be used to specify which logits in the batch to apply the bitmask to. It is especially useful when there are structured requests and unstructured requests mixed in the same batch by skipping masking the logits in the unstructured requests. When specified, the operation will be

for batch_id in indices:
    for j in range(vocab_size):
        if get_bitmask_value(bitmask, batch_id, j) == 0:
            logits[batch_id, j] = -inf

When indices is specified, the batch sizes of logits and bitmask do not need to be the same. As long as the indices are valid, the operation will be performed.

## Device

The logits and bitmask should be on the same device. If both them are on GPU, we launch a GPU kernel to apply bitmask. If both them are on CPU, we use a CPU implementation. The GPU kernel is optimized and should be preferred.

In practice, the bitmask is allocated on CPU, and the logits is usually on GPU, so users should manually copy the bitmask to GPU before calling this function.

Parameters:

logits (torch.Tensor) – The tensor to apply the bitmask to.
bitmask (torch.Tensor) – The bitmask to apply.
vocab_size (Optional[int], default: None) – The size of the vocabulary. If not provided, the vocab size will be detected as min(logits.shape[-1], bitmask.shape[-1] * 32).
indices (Optional[List[int]], default: None) – A list of indices to specify which logits in the batch to apply the bitmask to. Should be unique. If None, apply the bitmask to all logits in the batch.

xgrammar.get_bitmask_shape(batch_size: int, vocab_size: int) → Tuple[int, int]¶: Return the shape of the bitmask: (batch_size, ceil(vocab_size / 32)).

xgrammar.reset_token_bitmask(bitmask: Tensor) → None¶: Reset the bitmask to the full mask.