Extend on-device sampling support for dual QPC VLMs

[quic-xiyushi]

Overview

On-device sampling can significantly reduce host overhead and improve inference throughput; however, so far it has only been implemented for QEffForCausalLM models. This PR extends on-device sampling support to the language decoder of dual QPC vision language models, QEffCausalLMForTextImageToTextModel. In addition, it fixes the bug in gumbel noise so that it correctly simulates a multinomial distribution for random sampling.

Implementation details

class _QEffAutoModelForImageTextToTextDualQPC:

def __init__(
        self,
        model: nn.Module,
        continuous_batching: bool = False,
        qaic_config: Optional[dict] = None,
        **kwargs,
    ):
        # Omitting unchanged parts
        self.lang_model = QEffCausalLMForTextImageToTextModel(model, qaic_config=qaic_config, **kwargs)
        # ---Sampling---
        # Note: SamplerTransform should be applied after all other transforms
        # are done. The role of the sampler is to just add nodes at the output of the
        # previous transform function.
        self.lang_model.model, _ = SamplerTransform.apply(self.lang_model.model, qaic_config, **kwargs)

Usage

The usage is the similar to enable on-device sampling for QEffForCausalLM.

from QEfficient import QEFFAutoModelForImageTextToText

model_id = "Qwen/Qwen2.5-VL-3B-Instruct"

qeff_model = QEFFAutoModelForImageTextToText.from_pretrained(
    model_id,
    attn_implementation="eager",
    kv_offload=True,
    continuous_batching=True,
    qaic_config={
        "include_sampler": True,
        "return_pdfs": False,
        "max_top_k_ids": 512,
    },
)

[quic-hemagnih]

Can you please add the CI test cases.

[ochougul]

Does this mean we are enabling sampling only for intern model?
Will other VLMs also be supported?

[quic-xiyushi]

Other VLMs are also supposed to be supported. But currently only InternVL and Qwen VL 2.5 have been tested.

[quic-xiyushi]

Can you please add the CI test cases.

@quic-hemagnih CI added. Please review this PR again. Thank you!

[asmigosw]

Please add these both vision_embeds and image_idx in docs Args list.

[quic-xiyushi]

Since vision_embeds and image_idx come from the original forward method for VLMs and are not specifically added for on-device sampling, and because the docstring is intended only for parameters introduced for on-device sampling support, function docsstring was not added for vision_embeds and image_idx. Instead, I added added a note stating:
"The vision_embeds and image_idx parameters are optional and are used only for VLMs when supported by the original forward function."

In addition, to make this clearer, I reordered the arguments so that vision_embeds and image_idx appear right after num_logits_to_keep, before the on-device sampling arguments.

[quic-mamta]

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

[quic-xiyushi]

I have updated the dtype of these two parameters so that they are consistent with lines 27–28. For the function docstring, since vision_embeds and image_idx come from the original forward method for VLMs and are not specifically added for on-device sampling, and because the docstring is intended only for parameters introduced for on-device sampling support, function docsstring was not added for vision_embeds and image_idx. Instead, I added added a note stating:
"The vision_embeds and image_idx parameters are optional and are used only for VLMs when supported by the original forward function."

In addition, to make this clearer, I reordered the arguments so that vision_embeds and image_idx appear right after num_logits_to_keep, before the on-device sampling arguments.

[quic-mamta]

Please resolve the conflicts.

[ochougul]

merge if CI is passing

[quic-xiyushi]

Please resolve the conflicts.

Done.