Skip to content

Add Support for Guided Decoding to On Device Sampling #624

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
quic-sanising wants to merge 34 commits into
base: main
Choose a base branch
from
Open

Add Support for Guided Decoding to On Device Sampling #624

quic-sanising wants to merge 34 commits into from
+614 −178

Conversation

@quic-sanising
Copy link
Contributor

@quic-sanising quic-sanising commented Nov 19, 2025 •
edited
Loading

✨ Add Support for Guided Decoding to On Device Sampling

📌 Overview

This PR introduces guided decoding capabilities in On Device Sampling for QEffForCausalLM and QEffCausalLMForTextImageToTextModel models.



🚀 Motivation

As outlined in this blog on structured decoding, structured decoding represents a fundamental shift in controlling LLM outputs. Instead of relying on post-processing, constraints are enforced during token generation via logits manipulation. This approach ensures:

  • Format compliance at generation time.
  • Reduced error rates for structured outputs.
  • Performance improvements through optimized backends like XGrammar, which can deliver up to 5× faster token generation under load.

The constraints are provided through token_bitmasks which is a Boolean matrix of shape (batch_size, vocab_size). Here, each element indicates whether a token should be kept (1) or masked (0). During sampling, this mask is applied to the logits before token selection, ensuring that only allowed tokens are considered.

By performing this operation directly on the device, we eliminate host-device transfers, reduce latency, and improve throughput for structured decoding workloads.



🛠️ Implementation Details

The guided decoding logic is injected via include_guided_decoding=True during model loading. No changes to the model architecture are required.

from QEfficient import QEFFAutoModelForCausalLM as AutoModelForCausalLM

# Load model with On Device Sampler enabled
qeff_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    continuous_batching=True,
    qaic_config={
        "include_sampler": True,
        "return_pdfs": False,
        "max_top_k_ids": 512,
        "include_guided_decoding": True,
    },
)

# Compile as usual
qeff_model.compile(
    prefill_seq_length=128,
    ctx_len=256,
    full_batch_size=16,
    num_devices=4,
    num_speculative_tokens=0,
    mxint8_kv_cache=True,
    mxfp6_matmul=True,
)

To disable guided decoding, simply set include_guided_decoding=False.

quic-xiyushi and others added 13 commits November 10, 2025 09:16
@quic-xiyushi
Extend on-device sampling support for dual QPC VLMs
409da24 
Signed-off-by: quic-xiyushi <[email protected]>
@quic-xiyushi
Fix random_numbers shape
e06e175 
Signed-off-by: quic-xiyushi <[email protected]>
@quic-xiyushi
Update example with new random sampling logic
3e242ce 
Signed-off-by: quic-sanising <[email protected]>
Signed-off-by: sanising <[email protected]>
@quic-xiyushi
Update to align with recent VLM CB changes
1a01d57 
Signed-off-by: quic-xiyushi <[email protected]>
Update tests with new random sampling logic
30d6061 
Signed-off-by: sanising <[email protected]>
Add code to perform guided decoding
78ef180 
Signed-off-by: sanising <[email protected]>
Add bitmask to example inputs and dynamic axes
1fafcdb 
Signed-off-by: sanising <[email protected]>
Rename bitmask to token_bitmasks
18ab856 
Signed-off-by: sanising <[email protected]>
Fix typo
b1c049c 
Signed-off-by: sanising <[email protected]>
Merge branch 'main' into guided_decoding_simple
e16e846
Add flag to enable guided decoding
1515497 
Signed-off-by: sanising <[email protected]>
@quic-xiyushi
Merge remote-tracking branch 'origin/main' into HEAD
d02d04d
Add flag to enable guided decoding
97e4baf 
Signed-off-by: sanising <[email protected]>
@quic-sanising
Copy link
Contributor Author

quic-sanising commented Nov 19, 2025

Depends on #597

sanising and others added 6 commits November 19, 2025 14:05
@quic-sanising quic-sanising marked this pull request as ready for review November 20, 2025 19:37
@quic-sanising quic-sanising changed the title Add Guided Decoding Add Support for Guided Decoding to On Device Sampling Nov 20, 2025
@quic-sanising
Copy link
Contributor Author

quic-sanising commented Nov 21, 2025

Ready for review

quic-mamta
quic-mamta reviewed Dec 8, 2025
View reviewed changes
QEfficient/transformers/sampler/sampler.py
top_ps: Optional[torch.Tensor] = None,
min_ps: Optional[torch.Tensor] = None,
random_numbers: Optional[torch.Tensor] = None,
vision_embeds: Optional[torch.Tensor] = None,
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please keep dtype of these 2 consistent as per lines 27-28. also update function docstring for these newly added args.

tests/transformers/sampler/test_sampler.py
1, # spec_length
False, # is_vlm
),
# pytest.param(
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason for this test to be disabled?

tests/transformers/sampler/test_sampler.py
additional_configs["config"] = config
additional_configs["kv_offload"] = True
qeff_class = QEFFAutoModelForImageTextToText
assert isinstance(prompts, tuple)
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add error message in this case for vlms.

tests/transformers/sampler/test_sampler.py
prompts = prompts[1]
else:
qeff_class = QEFFAutoModelForCausalLM
spec_length -= 1
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

num_hidden_layers should be specified here or does this test require all model layers?

tests/transformers/sampler/test_sampler.py
"temperatures": np.array(100.1, dtype=np.float32).repeat(full_batch_size).reshape(-1, 1),
"top_ks": np.array(54720, dtype=np.int32).repeat(full_batch_size).reshape(-1, 1),
"temperatures": np.array(4.0, dtype=np.float32).repeat(full_batch_size).reshape(-1, 1),
"top_ks": np.array(512, dtype=np.int32).repeat(full_batch_size).reshape(-1, 1),
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How these numbers are changed?

tests/transformers/sampler/test_sampler.py
"top_ps": np.array(1.0, dtype=np.float32).repeat(full_batch_size).reshape(-1, 1),
"min_ps": np.array(0.0, dtype=np.float32).repeat(full_batch_size).reshape(-1, 1),
"random_numbers": np.array(0.0, dtype=np.float32).repeat(full_batch_size).reshape(-1, 1),
"random_numbers": np.zeros((full_batch_size, 512), dtype=np.float32),
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For greedy sampling, top_k = 1 and top_p = 0 should be used here?

tests/transformers/sampler/test_sampler.py
else:
additional_configs["num_hidden_layers"] = 2
qeff_class = QEFFAutoModelForCausalLM
spec_length -= 1
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we create a function(with inputs is_vlm, model, prompts, spec_length and return qeff_class, additional_params, additional_configs, spec_length) using lines 565-580 and reuse to avoid duplication?

tests/transformers/sampler/test_sampler.py
qaic_config={
"include_sampler": True,
"return_pdfs": False,
"max_top_k_ids": 1024,
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we define this test parameter in constants and then reuse?

examples/performance/on_device_sampling.py
include_sampler = None
return_pdfs = None
max_top_k_ids = None
include_guided_decoding = None
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These should be used as boolean variables here?

QEfficient/transformers/models/pytorch_transforms.py
QEffMptForCausalLM,
QEffPhi3ForCausalLM,
QEffQwen2ForCausalLM,
QEffQwen_2_5_vl_DecoderWrapper,
Copy link
Contributor

@quic-mamta quic-mamta Dec 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this supported for Intern and Qwen models only?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@quic-mamta quic-mamta quic-mamta left review comments

@quic-rishinr quic-rishinr Awaiting requested review from quic-rishinr quic-rishinr is a code owner

@ochougul ochougul Awaiting requested review from ochougul ochougul is a code owner

@quic-hemagnih quic-hemagnih Awaiting requested review from quic-hemagnih quic-hemagnih is a code owner

@quic-amitraj quic-amitraj Awaiting requested review from quic-amitraj quic-amitraj is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@quic-sanising @quic-mamta @quic-xiyushi