From bbcba7fc211c076206961a38e55c0bc380c85ca7 Mon Sep 17 00:00:00 2001
From: DoubleMathew <mmathew23@gmail.com>
Date: Sat, 12 Jul 2025 17:52:24 -0500
Subject: [PATCH] patch falcon h1 inference (#2932)

---
 unsloth/models/falcon_h1.py | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py
index 2cbb78f8ad..0db9c1ca4e 100644
--- a/unsloth/models/falcon_h1.py
+++ b/unsloth/models/falcon_h1.py
@@ -43,6 +43,9 @@ except:
 from transformers.modeling_attn_mask_utils import (
     _prepare_4d_causal_attention_mask_for_sdpa,
 )
+from transformers.utils import (
+    is_torchdynamo_compiling,
+)
 # For Pytorch 2.1.1
 try:
     from transformers.models.falcon_h1.modeling_falcon_h1 import (
@@ -519,7 +522,7 @@ def _FalconH1_fast_forward_inference(attention_fast_forward_inference=FalconH1At
                 attention_mask = attention_mask,
                 do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
             )
-            attention_hidden_states = attention_hidden_states * decoder_layer.attention_out_multiplier
+            attention_hidden_states = attention_hidden_states * decoder_layer.attn_out_multiplier
             mamba_hidden_states = decoder_layer.mamba(
                 hidden_states=X,
                 cache_params=present_key_value,
@@ -595,15 +598,17 @@ def _fast_prepare_inputs_for_generation(
             input_ids = input_ids[:, -cache_position.shape[0] :]
         elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
             input_ids = input_ids[:, cache_position]
-    else:
-        past_key_values = FalconHybridMambaAttentionDynamicCache(
-            self.config,
-            input_ids.shape[0],
-            self.dtype,
-            devices=[
-                self.model.layers[i].mamba.conv1d.weight.device for i in range(self.config.num_hidden_layers)
-            ],
-        )
+    pass
+    # TODO: Wire up Cache to work for inference.
+    # else:
+    #     past_key_values = FalconHybridMambaAttentionDynamicCache(
+    #         self.config,
+    #         input_ids.shape[0],
+    #         self.dtype,
+    #         devices=[
+    #             self.model.layers[i].mamba.conv1d.weight.device for i in range(self.config.num_hidden_layers)
+    #         ],
+    #     )
 
     if attention_mask is not None and position_ids is None:
         # create position_ids on the fly for batch generation