Decoding cannot be constrained: no per-step hook in the sampling path, and `logit_bias` cannot substitute for one

Component: cpp/sampler, cpp/runtime/decoding, experimental/pybind
Version: TensorRT Edge-LLM v0.10.1 (e8b2952), built from source with -DBUILD_PYTHON_BINDINGS=ON
Checkpoint: W4A16_AWQ quantise of nvidia/Cosmos3-Edge (model_type: cosmos3_edge), engine built from the published ONNX via llm_build
Working branch: filipemartinsubrobotics:feat/before-sampling-hook — one commit on e8b2952

Detailed description of the requested feature

There is no way to constrain what the runtime samples. LLMGenerationRequest
carries onTokenGenerated, which the decode loop fires after a token has been
accepted, so it can observe but not shape. There is no counterpart that runs
before sampling, and no equivalent of HuggingFace's LogitsProcessor, vLLM's
logits_processors, or llama.cpp's grammar.

The consequence is that structured output can only be checked after the fact
and discarded when it is wrong, never guaranteed.

Proposed shape

The mirror of the callback you already accept, in the same place on the request,
with the same std::optional<std::function> shape and lifetime:

using BeforeSamplingHook = std::function<bool(BeforeSamplingInfo const&)>;

Invoked once per active slot immediately before sampling. It receives the tokens
accepted so far and a cleared bitmask over the output vocabulary, sets the bits
it will allow, and returns whether to enforce. Cleared bits become -inf on the
device. Absent — which is every existing caller — it costs nothing.

A bitmask rather than a sparse list because its cost does not depend on how much
it forbids, and because it is what xgrammar and outlines already emit.

This does not ask Edge-LLM to own a grammar engine. It makes constrained
decoding a library choice.

We have implemented this

Branch linked above.

• cpp/runtime/decoding/tokenMask.{h,cpp} — staging and enforcement, split so
  the decision logic is testable without a GPU
• cpp/sampler/sampling.cu — applyTokenMaskRepeatedRowsKernel
• enforced at the end of prefill and at every vanilla decode step
• refused, with a message naming disable_spec_decode, when set against a
  speculative engine — draft tokens are proposed and verified outside this path,
  so the constraint would not hold, and silently not holding is the failure mode
  this whole request exists to remove
• an enforced mask that allows no token is refused too, rather than left as an
  all--inf row and undefined sampling
• exposed as LLMGenerationRequest.on_before_sampling, a callable
  (slot_index, generation_step, token_ids) -> Sequence[int] | None. Python
  never receives the mask buffer, which is valid only inside the call

Testing: 13 C++ cases (unittests/cpp/runtime/decoding/tokenMaskTests.cpp),
five of which run the kernel on device, plus 6 Python cases for the binding. The
existing unitTestRuntime suite is unchanged at 309 passed, 1 skipped. Each
guard was checked by mutation rather than assumed — neutering the kernel fails
four tests, removing the empty-mask refusal fails one.

End to end against Cosmos3-Edge INT4-AWQ on an RTX 4060 Ti, eight tokens a run:

hook	sampled ids
none	1073, 3219, 1261, 73916, 4604, 5970, 1513, 97558
allow only 1000	1000 × 8, hook invoked once per token
allow only 2000	2000 × 8
returns None	identical to the unconstrained row
allow {1000, 2000, 3000}	3000, 2000, 1000, 2000, 2000, 2000, 2000, 2000
allow {}	refused

The last row but one is the one we would point at: the model still chooses
within the permitted set rather than being dictated to, which is what a grammar
needs and what a bias cannot give you.

Recorded because it argues for the seam existing upstream rather than being
approximated downstream: the unit tests were green while the binding was
corrupting CPython's refcounts, because handle_request releases the GIL and
the runtime copies the request's hook into the decoding context. Only the
end-to-end run found it. A constraint API is easy to get subtly wrong from
outside the runtime.

We are not opening a PR, since CONTRIBUTING.md asks for an approved issue
first. If a callback is the wrong seam for this runtime, or you have a
preferred shape, we would much rather build to that than hand you a design you
have to unpick. The branch is offered as evidence the gap is real and
closeable, not as a fait accompli.

Timeline

No fixed date, and the branch above is usable today. Impact is should-have
rather than blocker: without a constraint we validate structured answers after
generation and discard the ones that fail, which works but cannot guarantee the
format.

Describe alternatives you've considered

The knobs that exist

This is the part worth checking before anything else, because it is decided by
constants in this repository rather than by our workload.

From cpp/common/inputLimits.h:

constexpr size_t kMaxLogitBiasTokens = 1024;
constexpr float  kMinLogitBias = -100.0F;
constexpr float  kMaxLogitBias =  100.0F;

A grammar is a state machine: which token is legal depends on what has already
been emitted. Against that, logit_bias:

cannot vary by position	it is one map for the whole generation, applied identically at every step
cannot forbid enough	1024 entries against a ~150k vocabulary; a constraint usually has to forbid nearly all of it
cannot forbid absolutely	values clamp to ±100, so a forbidden token is merely unlikely, not impossible

stop_strings bounds where an answer ends, not what shape it takes, so it
does not help either.

The third row is the one that matters most. A ±100 nudge is a preference. Any
API built on it would read like a guarantee and behave like a suggestion, which
in our experience is worse than having nothing — an off-format answer that
usually does not appear is one you stop checking for.

A conventional LogitsProcessor

The usual shape — a processor handed a vocab-sized tensor each step — means a
device-to-host copy of ~150k floats per token. On a 16 GB edge part driving a
robot that is not a reasonable ask.

But a grammar does not need to read the distribution. It needs to say which
tokens are legal. So the constraint can travel the other way, as a mask, and the
logits never have to leave the GPU.

What we fixed on our own side first

So that this is not a missing-homework request. Most of the bad output we saw
was ours, and none of it needed anything from Edge-LLM:

• Answers were truncating because reasoning consumed the token budget before
  the answer was emitted. Passing
  chat_template_kwargs: {"enable_thinking": false} took readable answers from
  20/24 to 24/24 on a fixed 24-frame set.
• Recorded separately because it is counter-intuitive and cost us a day:
  raising max_tokens does not fix that. The reasoning expands to fill the
  budget — 796 characters of it at 250 tokens, 1056 at 1024. Only declining the
  reasoning bounded it.
• Our own OpenAI-compatible shim was accepting stop and logit_bias and
  dropping both before they reached create_generation_request. Also ours, also
  fixed.

After all of that, the truncation problem is gone. What remains is the
guarantee: with a constraint, an off-format answer cannot be generated;
without one it can only be detected afterwards and thrown away.

Target hardware/use case

Cosmos3-Edge as the perception and reasoning model for a search-and-rescue
ground robot, asking one structured question per look.

• RTX 4060 Ti 16 GB, SM89, CUDA 12.8
• Edge-LLM v0.10.1 (e8b2952), Python bindings
• Cosmos3-Edge INT4-AWQ, engine from the published ONNX via llm_build
• Served through our own OpenAI-compatible shim, because
  tensorrt-edgellm-serve cannot build this multimodal checkpoint (#208)

---

A smaller note, not a request: grammar is a llama.cpp extension, and an
OpenAI-compatible server that does not implement it has no reason to accept it.
Ours did — it took the field and ignored it, which is indistinguishable from
honouring it until you measure the output. That was our bug and it is fixed. But
if a constraint API does land, a way to ask "is this enforced here" would save
the next caller the same measurement.

---

Found during the NVIDIA / OpenHackathons / Oracle Open Models Codefest 2026, while
running an INT4-AWQ Cosmos3-Edge reasoner as the perception model for an offline-first
search-and-rescue robotics entry (Team UBR Stack). The production target is a Jetson Orin
Nano; the RTX 4060 Ti is a bench machine used for evaluation.

Start by reading cpp/runtime/decoding/tokenMask.{h,cpp}, cpp/sampler/sampling.cu, and the existing LLMGenerationRequest handling in experimental/pybind. Run tokenMaskTests.cpp and the Python binding tests, then compare the proposed branch behavior across vanilla and speculative decoding. Done means an agreed constraint API is documented, tested on device and through Python, and preserves existing callers without silently bypassing enforcement.