On Device Sampling

[quic-sanising]

✨ On Device Sampling

📌 Overview

This PR introduces On Device Sampling for QEffForCausalLM models, enabling sampling operations to be executed directly on the QAIC device rather than the host CPU. This enhancement significantly reduces host-device communication overhead and improves inference throughput and scalability.

🚀 Motivation

Traditionally, sampling (e.g., greedy, top-k, top-p) is performed on the host CPU after logits are returned from the device. This approach incurs:

• High PCIe traffic due to large logits tensors [batch_size, vocab_size]
• Latency bottlenecks from CPU-bound sampling logic
• Limited scalability due to CPU thread constraints

On Device Sampling addresses these issues by:

• Performing sampling directly on the QAIC device
• Returning only the selected next tokens [batch_size, 1]
• Leveraging the device’s parallelism and optimized compute paths

⚙️ Supported Sampling Strategies

The following sampling techniques are now supported natively on the QAIC device:

1. Repetition Penalty​: Penalize tokens that have appeared in the prompt and the generated text so far. Values > 1 encourage the model to use new tokens, while values < 1 encourage the model to repeat tokens.​

2. Presence Penalty​: Penalize tokens that are present in the generated text so far. Values > 0 encourage the model to use new tokens, while values < 0 encourage the model to repeat tokens.

3. Temperature Scaling​: Adjust the sharpness of the logits distribution. Lower values make the model more deterministic, while higher values make the model more random. Zero means greedy sampling.​

4. Top K​: Sample from the k largest tokens by value.

5. Top P​: Sample from the smallest set of tokens whose cumulative probability is greater than or equal to p.​ Must be in (0, 1]. Set to 1 to consider all tokens.​

6. Min P​: Represents the minimum probability for a token to be considered, relative to the probability of the most likely token. Must be in [0, 1]. Set to 0 to disable this.​

7. Greedy Sampling​: Choose the token with highest value.

8. Random Sampling: Choose a token randomly with its probability of being chosen given by its value.

🛠️ Implementation Details

• Sampler Integration: Sampling logic is injected via include_sampler=True during model loading. No changes to the model architecture are required.

• Memory Optimization: Two scratch buffers of shape [batch_size, vocab_size] are used to track token occurrences for applying repetition and presence penalties efficiently on-device.

• Performance Gains:

  • Reduced PCIe traffic (logits → tokens)
  • Higher throughput via device-level parallelism
  • Scalable to 64+ concurrent inference streams

🧪 Usage

from QEfficient import QEFFAutoModelForCausalLM as AutoModelForCausalLM

# Load model with On Device Sampler enabled
qeff_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    include_sampler=True,
    return_pdfs=False,
)

# 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 On Device Sampling and revert to host-side sampling, simply set include_sampler=False.

[quic-hemagnih]

Can you please rebase the code and see if there are any conflicts. Mostly code looks okay to me. We can plan its merge once all the pending comments are addressed

[ochougul]

I don't see a test case.
everything else is Okay. approving.

[quic-amitraj]

@quic-sanising Could you please add the description for this PR?

[quic-amitraj]

Please add docstring to all the functions and classes.

[quic-amitraj]

Could you please add a docstring to enhance clarity for other developers? Additionally, could we relocate this to the sampler folder? This will help keep modeling_auto.py streamlined and avoid unnecessary complexity.

[quic-sanising]

Added the doc-string. Cannot to relocate this function to sampler as it is a member function of QEFFAutoModelForCausalLM.

[quic-sanising]

@quic-sanising Could you please add the description for this PR?

Done

[quic-hemagnih]

As discussed in the meeting, we will approve and merge this PR. Please have another PR with unit test added for this feature to ensure that its not broken by any subsequent checkins

[quic-sanising]

Moved to #440