View source for Implementation:Mlc ai Web llm Grammar Matcher Decoding

{{PageInfo|type=Implementation|title=Mlc_ai_Web_llm_Grammar_Matcher_Decoding}}
{{Metadata
| prefix = Mlc_ai_Web_llm
| type = Implementation
| name = Grammar_Matcher_Decoding
| repo = mlc-ai/web-llm
| version = 0.2.80
| status = Active
| package = @mlc-ai/web-llm
| source = src/llm_chat.ts
| lines = L616-L681 (grammar init), L1207-L1254 (bitmask application), L1472-L1484 (token acceptance)
| external_dependency = @mlc-ai/web-xgrammar 0.1.27
| relates_to = Mlc_ai_Web_llm_Grammar_Constrained_Decoding, Mlc_ai_Web_llm_Response_Format
}}

== Overview ==

'''Grammar Matcher Decoding''' is the concrete implementation of grammar-constrained token masking during LLM decoding in '''@mlc-ai/web-llm'''. It is implemented within the <code>LLMChatPipeline</code> class in <code>src/llm_chat.ts</code>, integrating the '''@mlc-ai/web-xgrammar''' library's <code>GrammarMatcher</code>, <code>GrammarCompiler</code>, and <code>TokenizerInfo</code> into the autoregressive decode loop. The implementation handles grammar compilation during prefill, bitmask computation and GPU application during each decode step, and parse state updates after each sampled token.

== Description ==

The implementation spans three key locations in <code>src/llm_chat.ts</code>:

=== 1. Grammar Initialization in prefillStep() (Lines 616-681) ===

During the prefill phase, the engine initializes the grammar matcher concurrently with prompt processing:

* If <code>response_format.type</code> is <code>"json_object"</code>, <code>"grammar"</code>, or <code>"structural_tag"</code>, a grammar initialization promise is created.
* '''Cache check:''' If the response format key matches the cached key and a <code>GrammarMatcher</code> already exists, the matcher is simply reset via <code>grammarMatcher.reset()</code> -- no recompilation needed.
* '''Fresh initialization:''' Otherwise, the existing matcher is disposed and a new one is created:
** <code>TokenizerInfo</code> is created from the tokenizer's raw token table (only once, then cached).
** <code>GrammarCompiler</code> is created from the <code>TokenizerInfo</code> (only once, then cached).
** The schema/grammar is compiled into a <code>CompiledGrammar</code> using the appropriate compiler method.
** A <code>GrammarMatcher</code> is created from the <code>CompiledGrammar</code>.
** The <code>CompiledGrammar</code> is disposed after the matcher is created.
* Compilation time is recorded in <code>curRoundGrammarInitTotalTime</code>.

=== 2. Bitmask Application in sampleTokenFromLogits() (Lines 1207-1254) ===

During each decode step, the grammar bitmask is applied to logits on GPU:

* A boolean <code>grammarConstrained</code> is computed from <code>response_format.type</code>.
* <code>grammarMatcher.getNextTokenBitmask()</code> returns an <code>Int32Array</code> representing the valid token bitmask.
* The bitmask is copied to GPU as an <code>int32</code> tensor.
* <code>fapplyBitmask(logitsOnGPU, seqIdsArray, bitMaskOnGPU)</code> applies the mask in place, setting invalid token logits to <code>-infinity</code>.
* Per-token grammar time is accumulated in <code>curRoundGrammarPerTokenTotalTime</code>.

=== 3. Token Acceptance (Lines 1472-1484) ===

After sampling, the grammar matcher state is updated:

* <code>grammarMatcher.acceptToken(sampledToken)</code> advances the parse state.
* If the token is rejected (returns <code>false</code>), an error is thrown -- this should never happen because the bitmask already excluded invalid tokens.
* Accept time is accumulated in <code>curRoundGrammarPerTokenTotalTime</code>.

== Code Reference ==

=== Grammar Initialization (prefillStep) ===

''Source: <code>src/llm_chat.ts</code>, lines 616-681''

<syntaxhighlight lang="typescript">
// -1. Instantiate grammar matcher according to generation config.
// This step is overlapped with prefilling the prompt to hide overhead.
let grammarMatcherInitPromise: Promise<void> | undefined = undefined;
const responseFormat = genConfig?.response_format;
if (
  responseFormat?.type === "json_object" ||
  responseFormat?.type === "grammar" ||
  responseFormat?.type === "structural_tag"
) {
  const curResponseFormatKey = this.getResponseFormatKey(responseFormat);
  if (
    curResponseFormatKey === this.responseFormatCacheKey &&
    this.grammarMatcher
  ) {
    // Reuse cached grammar matcher -- just reset its state
    const tGrammarInitStart = performance.now();
    log.info("Reuse grammar matcher.");
    this.grammarMatcher.reset();
    this.curRoundGrammarInitTotalTime =
      (performance.now() - tGrammarInitStart) / 1e3;
  } else {
    // Dispose current, reinitialize, and update cache key
    grammarMatcherInitPromise = new Promise(async (resolve) => {
      const tGrammarInitStart = performance.now();
      log.info("Initialize new grammar matcher.");
      if (this.grammarMatcher) {
        this.grammarMatcher.dispose();
      }
      if (this.xgTokenizerInfo === undefined) {
        log.info("Initialize token table.");
        const rawTokenTable = getTokenTableFromTokenizer(this.tokenizer);
        this.xgTokenizerInfo = await xgr.TokenizerInfo.createTokenizerInfo(
          rawTokenTable,
          this.token_postproc_method,
          this.prepend_space_in_encode,
          this.fullVocabSize,
          this.stopTokens,
        );
        this.grammarCompiler =
          await xgr.GrammarCompiler.createGrammarCompiler(
            this.xgTokenizerInfo,
          );
      }
      const grammar: xgr.CompiledGrammar =
        responseFormat.type === undefined
          ? await this.grammarCompiler!.compileBuiltinJSONGrammar()
          : responseFormat.type === "json_object"
            ? await this.grammarCompiler!.compileJSONSchema(
                responseFormat.schema!,
              )
            : responseFormat.type === "grammar"
              ? await this.grammarCompiler!.compileGrammar(
                  responseFormat.grammar!,
                )
              : await this.grammarCompiler!.compileStructuralTag(
                  responseFormat.structural_tag!,
                );
      this.grammarMatcher =
        await xgr.GrammarMatcher.createGrammarMatcher(grammar);
      grammar.dispose();
      this.responseFormatCacheKey = curResponseFormatKey;
      this.curRoundGrammarInitTotalTime =
        (performance.now() - tGrammarInitStart) / 1e3;
      resolve();
    });
  }
}
</syntaxhighlight>

=== Bitmask Application (sampleTokenFromLogits) ===

''Source: <code>src/llm_chat.ts</code>, lines 1207-1254''

<syntaxhighlight lang="typescript">
const grammarConstrained =
  response_format?.type === "json_object" ||
  response_format?.type === "grammar" ||
  response_format?.type === "structural_tag";

// 0. Update logitsOnGPU with on-GPU grammar bitmasking
if (grammarConstrained) {
  const grammarBitmaskBegin = performance.now();

  this.tvm.beginScope();
  if (this.grammarMatcher === undefined) {
    throw Error("Expect grammar matcher to be initialized.");
  }

  const tBitmaskStart = performance.now();
  const bitMaskOnCPU: Int32Array =
    await this.grammarMatcher.getNextTokenBitmask();
  this.curRoundGrammarPerTokenTotalTime +=
    (performance.now() - tBitmaskStart) / 1e3;

  if (bitMaskOnCPU.length !== this.bitmaskSize) {
    throw new Error(
      `InternalError: Expect grammar bitmask to be ` +
        `size ${this.bitmaskSize}, but got ${bitMaskOnCPU.length}.`,
    );
  }
  const bitMaskOnGPU = this.tvm
    .empty([1, this.bitmaskSize], "int32", this.device)
    .copyFrom(bitMaskOnCPU);
  const seqIdsArray = this.tvm
    .empty([1], "int32", this.device)
    .copyFrom([0]);
  this.fapplyBitmask(
    logitsOnGPU.view([1, this.fullVocabSize]),
    seqIdsArray,
    bitMaskOnGPU,
  );
  this.tvm.endScope();
}
</syntaxhighlight>

=== Token Acceptance ===

''Source: <code>src/llm_chat.ts</code>, lines 1472-1484''

<syntaxhighlight lang="typescript">
// 6. Update grammar matcher with new token
if (grammarConstrained) {
  if (this.grammarMatcher === undefined) {
    throw Error("Expect grammar matcher to be initialized.");
  }
  const tAcceptStart = performance.now();
  const accepted = this.grammarMatcher.acceptToken(sampledToken);
  this.curRoundGrammarPerTokenTotalTime +=
    (performance.now() - tAcceptStart) / 1e3;
  if (!accepted) {
    throw Error("Grammar matcher rejected the newly sampled token.");
  }
}
</syntaxhighlight>

== I/O Contract ==

{| class="wikitable"
|-
! Direction !! Type !! Description
|-
| '''Input''' || <code>ChatCompletionRequest</code> with <code>response_format</code> || A request containing a schema, grammar, or structural tag definition
|-
| '''Output''' || <code>ChatCompletion</code> || A completion response whose <code>choices[].message.content</code> is guaranteed to conform to the specified grammar
|}

=== Internal Data Flow ===

{| class="wikitable"
|-
! Step !! Operation !! Data Type !! Location
|-
| 1 || Schema compilation || <code>ResponseFormat</code> -> <code>CompiledGrammar</code> -> <code>GrammarMatcher</code> || <code>prefillStep()</code>
|-
| 2 || Bitmask computation || <code>GrammarMatcher</code> -> <code>Int32Array</code> || <code>sampleTokenFromLogits()</code>
|-
| 3 || GPU bitmask application || <code>Int32Array</code> -> <code>tvmjs.Tensor</code> -> applied to logits || <code>sampleTokenFromLogits()</code>
|-
| 4 || Token sampling || Masked logits -> <code>number</code> (token ID) || <code>sampleTokenFromLogits()</code>
|-
| 5 || State update || <code>number</code> -> <code>grammarMatcher.acceptToken()</code> || <code>sampleTokenFromLogits()</code>
|}

=== External Dependencies ===

{| class="wikitable"
|-
! Package !! Version !! Components Used
|-
| <code>@mlc-ai/web-xgrammar</code> || 0.1.27 || <code>GrammarMatcher</code>, <code>GrammarCompiler</code>, <code>TokenizerInfo</code>, <code>CompiledGrammar</code>
|-
| <code>@mlc-ai/web-runtime</code> || (bundled) || TVM runtime for GPU tensor operations (<code>tvm.empty()</code>, <code>fapplyBitmask()</code>)
|}

=== Performance Metrics ===

The implementation records two performance metrics available in <code>usage.extra</code> of the response:

* '''<code>grammar_init_s</code>''' -- Total seconds spent initializing the grammar matcher (compilation + matcher creation). If <code>n > 1</code> (multiple choices), this is the sum across all choices.
* '''<code>grammar_per_token_s</code>''' -- Average seconds per token spent on bitmask generation and token acceptance. If <code>n > 1</code>, this is the average across all choices.

== Usage Examples ==

=== Observing Grammar Metrics in a Non-Streaming Request ===

<syntaxhighlight lang="typescript">
import * as webllm from "@mlc-ai/web-llm";

const engine = await webllm.CreateMLCEngine(
  "Llama-3.2-3B-Instruct-q4f16_1-MLC",
  { logLevel: "INFO" },
);

const schema = JSON.stringify({
  type: "object",
  properties: {
    name: { type: "string" },
    house: {
      type: "string",
      enum: ["Gryffindor", "Hufflepuff", "Ravenclaw", "Slytherin"],
    },
    alive: { type: "boolean" },
  },
  required: ["name", "house", "alive"],
});

const reply = await engine.chat.completions.create({
  stream: false,
  messages: [
    {
      role: "user",
      content: "Fill in info about Harry Potter in JSON format.",
    },
  ],
  max_tokens: 128,
  response_format: {
    type: "json_object",
    schema: schema,
  } as webllm.ResponseFormat,
});

// The output is guaranteed to be valid JSON conforming to the schema
console.log("Content:", reply.choices[0].message.content);

// Grammar-specific performance metrics
const extra = reply.usage?.extra;
if (extra) {
  console.log("Grammar init time:", extra.grammar_init_s, "seconds");
  console.log("Grammar per-token time:", extra.grammar_per_token_s, "seconds");
}
</syntaxhighlight>

=== Streaming with Grammar Constraints ===

Grammar-constrained decoding also works with streaming responses. The grammar bitmask is applied at each decode step regardless of whether output is streamed:

<syntaxhighlight lang="typescript">
import * as webllm from "@mlc-ai/web-llm";

const engine = await webllm.CreateMLCEngine("Phi-3.5-mini-instruct-q4f16_1-MLC");

const schema = JSON.stringify({
  type: "object",
  properties: {
    city: { type: "string" },
    temperature: { type: "number" },
    conditions: { type: "string" },
  },
  required: ["city", "temperature", "conditions"],
});

const chunks = await engine.chat.completions.create({
  stream: true,
  messages: [
    { role: "user", content: "Give me weather info for Tokyo in JSON." },
  ],
  max_tokens: 128,
  response_format: {
    type: "json_object",
    schema: schema,
  } as webllm.ResponseFormat,
});

let fullContent = "";
for await (const chunk of chunks) {
  const delta = chunk.choices[0]?.delta?.content ?? "";
  fullContent += delta;
  process.stdout.write(delta);
}

// fullContent is guaranteed to be valid JSON matching the schema
const parsed = JSON.parse(fullContent);
console.log("\nParsed:", parsed);
</syntaxhighlight>

== Related Pages ==

* [[Mlc_ai_Web_llm_Grammar_Constrained_Decoding|Principle: Grammar-Constrained Decoding]] -- [[implements::Principle:Mlc_ai_Web_llm_Grammar_Constrained_Decoding]]
* [[Mlc_ai_Web_llm_Response_Format|Implementation: Response Format]] -- The interface that specifies which grammar to compile
* [[Mlc_ai_Web_llm_JSON_Parse_Output|Implementation: JSON Parse Output]] -- Pattern for parsing the output produced by this decoding process
* [[requires_env::Environment:Mlc_ai_Web_llm_WebGPU_Browser_Runtime]]
* [[uses_heuristic::Heuristic:Mlc_ai_Web_llm_Grammar_Matcher_Reuse]]
* [[uses_heuristic::Heuristic:Mlc_ai_Web_llm_Tokenizer_JSON_Preference]]

[[Category:Implementations]]