Extend On-Device Sampling Support to more Causal Language Models (#553)

### 📢 Expanded On-Device Sampling Support in QEfficient

Excited to share that **On-Device Sampling**—previously available only
for `LlamaForCausalLM`—is now supported across a broader set of
architectures! This enhancement brings faster, more efficient inference
directly to the QAIC device.

#### ✅ Newly Supported Architectures:
1. `FalconForCausalLM`
2. `GemmaForCausalLM`
3. `GPT2LMHeadModel`
4. `GPTJForCausalLM`
5. `GraniteForCausalLM`
6. `GraniteMoeForCausalLM`
7. `LlamaForCausalLM` (existing)
8. `MptForCausalLM`
9. `Phi3ForCausalLM`
10. `Qwen2ForCausalLM`

#### ⚠️ Architectures Still Pending Support:
1. `GPTBigCodeForCausalLM`
2. `InternVLChatModel`
3. `MistralForCausalLM`
4. `MixtralForCausalLM`
5. `LlamaSwiftKVForCausalLM`
6. `Grok1ModelForCausalLM`

We’re actively working to extend support to these models. Contributions,
feedback, and testing from the community are always welcome to help
accelerate this effort!

---------

Signed-off-by: quic-sanising <[email protected]>
Signed-off-by: sanising <[email protected]>
Signed-off-by: Dhiraj Kumar Sah <[email protected]>
Co-authored-by: sanising <[email protected]>
Co-authored-by: Dhiraj Kumar Sah <[email protected]>
Co-authored-by: Hem Agnihotri <[email protected]>

Author: 4 people

QEfficient/transformers/models/pytorch_transforms.py

@@ -676,8 +676,16 @@ class SamplerTransform:
 
     # supported architectures
     _module_mapping = {
-        # Llama
+        QEffFalconForCausalLM,
+        QEffGemmaForCausalLM,
+        QEffGPT2LMHeadModel,
+        QEffGPTJForCausalLM,
+        QEffGraniteForCausalLM,
+        QEffGraniteMoeForCausalLM,
         QEffLlamaForCausalLM,
+        QEffMptForCausalLM,
+        QEffPhi3ForCausalLM,
+        QEffQwen2ForCausalLM,
     }
 
     @classmethod