feat: GBNF grammar-constrained sampling

[hvico]

Motivation

FastFlowLM excels at fast inference on AMD NPU hardware, but currently offers no way to constrain the output format of generated text. When using FLM for structured extraction tasks (JSON, function calls, form-filling), the model can emit malformed output — missing keys, wrong types, trailing commentary, markdown wrappers — that downstream parsers must handle defensively or reject entirely.

llama.cpp solved this years ago with GBNF grammars: a formal grammar (EBNF-like) is compiled into a pushdown automaton that masks logits at each decoding step, guaranteeing the output conforms to the grammar. This is the de facto standard for constrained decoding in the GGUF ecosystem.

This PR brings that same capability to FastFlowLM, so AMD NPU users get the same structured-output guarantees that GPU users have enjoyed via llama.cpp.

What this PR does

Ports the GBNF grammar parser and enforcer from llama.cpp (llama-grammar.h / llama-grammar.cpp, MIT-licensed) into FastFlowLM as a self-contained module, and wires it into the existing sampling pipeline and REST API.

New files

File	Lines	Description
src/include/grammar/grammar.hpp	129	Grammar data structures, GBNF parser interface, and FlmGrammar class declaration
src/common/grammar/grammar.cpp	798	Full implementation: GBNF parser, pushdown automaton, logit filtering, state advancement

Modified files

File	Change	Description
src/include/modules/sampler.hpp	+13	Added FlmGrammar* member and set_grammar() method to Sampler class
src/common/modules/sampler.cpp	+13	Grammar filter applied before top-k; grammar state advanced after sampling
src/include/AutoModel/automodel.hpp	+17	Added set_grammar() / clear_grammar() public API and grammar state members
src/common/AutoModel/automodel.cpp	+38	Token-piece cache construction, grammar lifecycle management, auto-cleanup after generation
src/server/rest_handler.cpp	+6	Parse grammar and grammar_root fields from API requests

How it works

Architecture

API request with "grammar" field
        │
        ▼
  RestHandler::configure_chat_engine_parameters()
        │  parses grammar string, calls auto_chat_engine->set_grammar()
        ▼
  AutoModel::set_grammar()
        │  1. Builds token-piece lookup table (cached across requests)
        │  2. Parses GBNF string into rules via FlmGrammar::create()
        │  3. Initializes pushdown automaton stacks
        │  4. Wires grammar pointer into Sampler
        ▼
  Sampler::sample()  (called once per generated token)
        │  1. Copy logits from NPU output
        │  2. Apply repetition/frequency penalties
        │  3. ★ grammar->apply(logits, vocab_size)     ← NEW
        │  4. Top-k filtering
        │  5. Softmax + top-p + min-p + temperature
        │  6. Multinomial sampling
        │  7. ★ grammar->accept(sampled_token)          ← NEW
        │  8. Ring buffer update
        ▼
  AutoModel::_shared_generate()
        │  after generation completes:
        │  ★ clear_grammar()                            ← auto-cleanup
        ▼
  Response with grammar-conformant text

Grammar engine (ported from llama.cpp)

The grammar engine is a character-level pushdown automaton:

1. Parse phase: The GBNF string is parsed into a set of production rules. Each rule is a sequence of elements: character matches ([a-z], "literal", .), character ranges ([^abc]), rule references, and alternates (|).

2. Apply phase (FlmGrammar::apply): Before top-k filtering, the grammar iterates over all vocabulary tokens. Each token is decoded to UTF-8 codepoints via a pre-built lookup table (token_pieces_cache_). Tokens whose decoded text would violate the current grammar state get their logits set to -inf, preventing them from being sampled.

3. Accept phase (FlmGrammar::accept): After a token is sampled, the grammar advances its internal stacks by consuming the token's codepoints character by character, producing the set of valid next states.

What was ported vs. stripped

Ported from llama.cpp (fully functional):

• Complete GBNF parser (literals, char ranges, negation, alternates, rule references, grouping, repetitions *+?{m,n}, any-char .)
• Pushdown automaton with multi-stack state tracking
• Character-level matching with partial UTF-8 sequence handling
• Left-recursion detection
• Grammar cloning (for parallel/batched inference)

Intentionally excluded (to minimize scope and dependencies):

• <token> syntax (requires tokenizer integration at parse time — only <[numeric_id]> is supported)
• Lazy grammars and trigger patterns (advanced llama.cpp feature for tool-use routing)

The grammar module has zero dependencies on llama.cpp runtime — it's entirely self-contained within FastFlowLM.

API usage

Request format

The grammar parameter accepts a GBNF grammar string. The optional grammar_root parameter specifies the start rule (defaults to "root").

Works with all three endpoints:

• POST /v1/chat/completions (OpenAI-compatible)
• POST /v1/completions (OpenAI-compatible)
• POST /api/chat (Ollama-compatible)

Example: Force digit-only output

curl http://localhost:8081/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "qwen3.5:2b",
    "messages": [{"role": "user", "content": "What is 2+2? Answer with just the number."}],
    "max_tokens": 10,
    "temperature": 0,
    "grammar": "root ::= [0-9]+"
  }'
# → {"content": "4"}

Example: Force JSON object with specific keys

curl http://localhost:8081/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "qwen3.5:2b",
    "messages": [{"role": "user", "content": "Extract: Name is John Smith, age 35."}],
    "max_tokens": 100,
    "temperature": 0,
    "grammar": "root ::= \"{\" ws name-kv \",\" ws age-kv \"}\" ws\nname-kv ::= \"\\\"name\\\"\" ws \":\" ws \"\\\"\" [a-zA-Z ]+ \"\\\"\"\nage-kv ::= \"\\\"age\\\"\" ws \":\" ws [0-9]+\nws ::= [ \\t\\n]*"
  }'
# → {"content": "{\"name\": \"John Smith\", \"age\": 35}"}

Example: Force JSON array

curl http://localhost:8081/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "qwen3.5:2b",
    "messages": [{"role": "user", "content": "List 3 fruits as a JSON array"}],
    "max_tokens": 50,
    "temperature": 0.3,
    "grammar": "root ::= \"[\" ws item (\",\" ws item)* ws \"]\"\nitem ::= \"\\\"\" [a-zA-Z]+ \"\\\"\"\nws ::= [ \\t\\n]*"
  }'
# → {"content": "[\"Apple\", \"Banana\", \"Orange\"]"}

Design decisions

1. Grammar applied before top-k, not after. The grammar filter must see the full vocabulary to correctly mask tokens. Applying it after top-k would miss valid tokens that top-k already discarded, potentially leaving zero candidates.

2. Non-owning pointer in Sampler. The FlmGrammar* in Sampler is non-owning — AutoModel owns the grammar via unique_ptr and automatically clears it after generation completes. This prevents grammar state from leaking between requests.

3. Token-piece cache built lazily. The first grammar request triggers a one-time O(vocab_size) scan that decodes every token ID to its text representation. This cache is reused across requests and invalidated on model switch.

4. Zero overhead when unused. When no grammar is active, the only cost is a null-pointer check per sampling step. No allocations, no iterations.

Testing

Tested on AMD Strix Halo hardware (XDNA NPU) with Qwen3.5-2B and Qwen3.5-4B models:

• Digit-only grammar: Forces model to output only digits ✓
• JSON object grammar: Forces well-formed JSON with required keys ✓
• JSON array grammar: Forces valid JSON arrays ✓
• Complex schema grammar: 9+ field JSON extraction from document images ✓
• No grammar (regression): Normal generation unaffected ✓
• Grammar parse errors: Logged to stderr, generation proceeds without constraint ✓

License

The grammar engine is derived from llama.cpp's llama-grammar.h / llama-grammar.cpp, which is licensed under MIT. The ported code retains the same license terms. All llama.cpp-specific dependencies (ggml, llama context, vocabulary interface) have been removed — the port is fully self-contained.

[Cytrus14]

I've been testing your PR with Qwen3.5-9B. Without it, tool calling would often fail due to malformed JSON output. Now it works without issues. Great job!

[Copilot]

Pull request overview

Adds GBNF grammar-constrained decoding to FastFlowLM by porting a llama.cpp-derived grammar engine and integrating it into the sampling pipeline and REST API so outputs can be forced to match a specified grammar (e.g., JSON).

Changes:

• Introduces a self-contained GBNF grammar parser/enforcer module (FlmGrammar) and hooks it into token sampling (logit masking + state advancement).
• Adds per-request grammar lifecycle management in AutoModel (token-piece cache + set/clear API + auto cleanup).
• Extends REST request parsing to accept grammar and optional grammar_root.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/include/grammar/grammar.hpp	Declares grammar data structures, parser API, and FlmGrammar interface.
src/common/grammar/grammar.cpp	Implements GBNF parsing + pushdown automaton + logit filtering + state advancement.
src/include/modules/sampler.hpp	Adds optional grammar pointer + setter to Sampler.
src/common/modules/sampler.cpp	Applies grammar masking before top-k and advances grammar after sampling.
src/include/AutoModel/automodel.hpp	Adds set_grammar / clear_grammar APIs and per-request grammar members.
src/common/AutoModel/automodel.cpp	Builds token-piece cache, creates/clears grammar, and clears grammar after generation.
src/server/rest_handler.cpp	Parses grammar / grammar_root from requests and forwards to the engine.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

clear_grammar() is called only at the end of _shared_generate(), but there are early return paths above (e.g., immediate EOS after prefill, MAX_L reached) that bypass this call and can leak grammar state into the next request. Consider an RAII/scope guard to ensure clear_grammar() runs on all exits (including early returns and exceptions).

[Copilot]

Grammar masking can legally reject every token in the current state. If apply() makes all logits -inf, softmax_inplace() will compute exp(-inf - -inf) and produce NaNs, breaking sampling. After applying the grammar, detect the "no valid tokens" case and fall back deterministically (e.g., allow EOG if grammar permits, otherwise disable/clear grammar for this step or pick the highest-logit unmasked token) so the sampler never feeds an all -inf set into softmax.

[Copilot]

is newly included here but isn’t used in this header (only a forward decl + raw pointer are added). Consider removing it to keep includes minimal and avoid unnecessary rebuilds.

Suggested change

#include <memory>

[Copilot]

AutoModel::set_grammar() assigns grammar_ directly from FlmGrammar::create(). If create() fails, the previous grammar_ is destroyed but Sampler still retains the old non-owning pointer (set_grammar() is only called on success), which can lead to a dangling pointer / UAF during sampling. Create into a temporary unique_ptr and only swap + set sampler pointer on success, otherwise explicitly clear_grammar() (or at least sampler->set_grammar(nullptr)).

Suggested change

	grammar_ = FlmGrammar::create(grammar_str, grammar_root,
	token_pieces_cache_, eos_token_ids);
	if (grammar_ && this->sampler) {
	this->sampler->set_grammar(grammar_.get());
	} else if (!grammar_) {
	fprintf(stderr, "AutoModel::set_grammar: failed to parse grammar\n");
	}
	auto new_grammar = FlmGrammar::create(grammar_str, grammar_root,
	token_pieces_cache_, eos_token_ids);
	if (!new_grammar) {
	fprintf(stderr, "AutoModel::set_grammar: failed to parse grammar\n");
	clear_grammar();
	return;
	}
	grammar_ = std::move(new_grammar);
	if (this->sampler) {
	this->sampler->set_grammar(grammar_.get());
	}

[hvico]

Thanks for the review! Pushed 6a42ba6 addressing all four Copilot observations, rebased on current main. Each fix verified at runtime in a container build against qwen3.5:4b on an AMD Strix Halo NPU.

What changed

1. RAII guard for clear_grammar() (src/common/AutoModel/automodel.cpp)
Replaced the single trailing call at the end of _shared_generate() with a scope-local guard at the top of the function:

struct GrammarGuard {
    AutoModel * self;
    ~GrammarGuard() { self->clear_grammar(); }
} grammar_guard{this};

Now the early-return paths (immediate EOS after prefill, MAX_L reached) and any exception also clear grammar state. Verified with a curl sequence: grammar-constrained request followed by an unconstrained one — the second returns normal text ("Hello!") with no residual constraint.

2. All-masked grammar fallback (src/common/modules/sampler.cpp)
Snapshot logits before grammar->apply(), detect whether any finite logit survives, and on !isfinite(max) restore the pre-grammar snapshot + skip the downstream grammar->accept() call (since the sampled token is outside the grammar, advancing the parser would corrupt it). Warn on stderr.

bool grammar_fell_back = false;
// ... after grammar->apply()
if (!std::isfinite(max_l)) {
    fprintf(stderr, "Sampler: grammar masked all tokens; falling back...");
    std::copy(pre_grammar.begin(), pre_grammar.end(), this->logits.begin());
    grammar_fell_back = true;
}
// ... later
if (this->grammar && !grammar_fell_back) {
    this->grammar->accept(sampled_index);
}

Without the grammar_fell_back skip, the process crashes inside the grammar state machine when accept() receives an unexpected token. Verified with grammar: "root ::= <[999999]>" (token ID that doesn't exist — every step is all-masked): server stays alive, log shows 4× Sampler: grammar masked all tokens; falling back..., output is unconstrained ("Hello! How can"), and the subsequent no-grammar request returns "OK" normally.

3. Removed unused <memory> include (src/include/modules/sampler.hpp)
Header only uses a forward decl + raw pointer, so the include was unnecessary.

4. set_grammar build-and-swap (src/common/AutoModel/automodel.cpp)
Construct into a local unique_ptr first; if FlmGrammar::create() fails, call clear_grammar() (which resets the sampler's non-owning pointer) and bail out. Only std::move into grammar_ and wire the sampler pointer on success.

auto new_grammar = FlmGrammar::create(grammar_str, grammar_root,
                                      token_pieces_cache_, eos_token_ids);
if (!new_grammar) {
    fprintf(stderr, "AutoModel::set_grammar: failed to parse grammar\n");
    clear_grammar();
    return;
}
grammar_ = std::move(new_grammar);
if (this->sampler) this->sampler->set_grammar(grammar_.get());

Verified with grammar: "this is not a valid gbnf grammar": log shows the parse error + our AutoModel::set_grammar: failed to parse grammar, the request completes unconstrained ("Hello! How can I help you today?"), and the next request works normally.

Runtime validation summary

Test	Grammar	Result
Digits only	root ::= [0-9]+	"23" (valid)
Forced JSON	{ "name": <letters>, "age": <digits> }	{"name": "Elara", "age": 34}
Copilot #4	malformed (no ::=)	parse-error logged, falls back to unconstrained, next request unaffected
Copilot #2	root ::= <[999999]> (all-masked)	no NaN, no crash, fallback warning logged 4×, follow-up normal
Copilot #1	any grammar + follow-up unconstrained	follow-up unconstrained — no residual

[ZaneNi]

I've been testing your PR with Qwen3.5-9B. Without it, tool calling would often fail due to malformed JSON output. Now it works without issues. Great job!

Hi, @Cytrus14, do you mind sharing your grammar for Qwen3.5? Thank you!

[Cytrus14]

@ZaneNi
Sure! I've used the json GBNF grammar from llama.cpp (link). In my use case, I simply read the included file as plain text and set the grammar field to its contents.