Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:InternLM Lmdeploy GuidedDecoding

From Leeroopedia


Knowledge Sources
Domains Text Generation, Structured Output
Last Updated 2026-02-07 15:00 GMT

Overview

Implements grammar-guided decoding using xgrammar bitmasks to constrain token generation to valid outputs according to a formal grammar (e.g., JSON schema, regex).

Description

The GuidedDecoding class integrates the xgrammar library into TurboMind's generation pipeline to enforce structured output constraints during token sampling. It inherits from BaseGenerationParam and operates in four distinct phases:

Setup: Collects grammar matchers from the current batch's request caches. Each request may optionally have an xgrammar::GrammarMatcher attached. The active flag is set if any request in the batch uses guided decoding.

FillMask: When active, generates a token bitmask using xgrammar's FillNextTokenBitmask for each non-terminated matcher. This bitmask indicates which tokens are valid according to the grammar's current state. Only rank 0 in the tensor-parallel group fills the mask; for terminated or absent matchers, the bitmask is zeroed out.

ApplyMask: Broadcasts the bitmask across tensor-parallel ranks (if multi-rank), copies it to device memory, and applies it to the logits tensor in-place using ApplyTokenBitmaskInplace. This effectively sets invalid token logits to negative infinity.

Update: After sampling, copies the selected output token IDs back to CPU and feeds them to each active matcher via AcceptToken, advancing the grammar state for the next step.

The class maintains per-phase Data structs containing bitmask tensors and matcher references, plus shared pinned-memory buffers for the bitmask and output IDs.

Usage

Used within the Generation module's pipeline. Activated automatically when requests have grammar matchers attached (set via ModelRequest::setGrammar()). Supports JSON schema and regex-based output constraints.

Code Reference

Source Location

Signature

class GuidedDecoding: public BaseGenerationParam {
public:
    explicit GuidedDecoding(const BaseGenerationParam& base, const comm::HostComm& tp_group, int phases);

    void Setup(int phase, TensorMap& env);

    void FillMask(int phase, TensorMap& env);

    void ApplyMask(int phase, TensorMap& env);

    void Update(int phase, TensorMap& env);

private:
    comm::HostComm tp_group_;
    struct Data;
    std::vector<std::shared_ptr<Data>> data_;
    Tensor_<int32_t> bitmask_buf_;
    Buffer_<int>     output_ids_buf_;
};

Import

#include "src/turbomind/generation/guided_decoding.h"

I/O Contract

Inputs

Name Type Required Description
base BaseGenerationParam Yes Base parameters (max_batch_size, vocab_size, vocab_size_padded, session_len)
tp_group const comm::HostComm& Yes Tensor-parallel communication group for bitmask broadcast
phases int Yes Number of pipeline phases
env["batch"] BatchData* Yes Current batch data with request caches containing grammar matchers
env["logits"] Tensor Yes (ApplyMask) Logits tensor to mask in-place
env["output_ids"] Tensor Yes (Update) Sampled token IDs to feed back to grammar matchers

Outputs

Name Type Description
env["logits"] (modified) Tensor Logits with invalid tokens masked to negative infinity

Usage Examples

// Construction (done inside Generation module)
GuidedDecoding guided(base_param, tp_group, phases);

// During batch lifecycle:
guided.Setup(phase, env);      // Collect matchers from batch
guided.FillMask(phase, env);   // Generate token bitmasks
guided.ApplyMask(phase, env);  // Apply bitmasks to logits
// ... sampling happens ...
guided.Update(phase, env);     // Feed sampled tokens to grammar matchers

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment