Fix self.temp_O buffer shape for batched inference

The temp_O output buffer was allocated as (1, bsz, hidden_size) but
torch.matmul expects (bsz, 1, hidden_size) when bsz > 1. With the
wrong shape, PyTorch silently resizes the output tensor, producing a
UserWarning about deprecated resize behavior that will become an error
in a future release.

This only affects the Pattern A allocation path (when attention_size !=
hidden_size, e.g. Mistral Nemo 12b). The Pattern B alias path already
gets the correct shape from temp_QA[1]. When bsz == 1, the shapes are
identical so there is no behavioral change for single-sequence
inference.

Files changed:
- unsloth/models/llama.py
- unsloth/models/gemma2.py
- unsloth/models/granite.py
- unsloth/models/cohere.py
- unsloth/models/qwen3.py
- unsloth/models/falcon_h1.py

unsloth/models/cohere.py

@@ -296,7 +296,7 @@ def CohereAttention_fast_forward_inference(
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
             self.temp_O = torch.empty(
-                (1, bsz, hidden_size), dtype = dtype, device = "cuda:0"
+                (bsz, 1, hidden_size), dtype = dtype, device = "cuda:0"
             )
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]

unsloth/models/falcon_h1.py

@@ -265,7 +265,7 @@ def FalconH1Attention_fast_forward_inference(
 
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
-            self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+            self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]
 

unsloth/models/gemma2.py

@@ -352,7 +352,7 @@ def Gemma2Attention_fast_forward_inference(
         )
         self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
         # Only for Gemma2
-        self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+        self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         self.attention = torch.empty(
             (bsz, n_heads, 1, KV_CACHE_INCREMENT + seq_len), dtype = dtype, device = device
         )

unsloth/models/granite.py

@@ -324,7 +324,7 @@ def GraniteAttention_fast_forward_inference(
         )
         self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
         # Only for Gemma2
-        self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+        self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         self.attention = torch.empty(
             (bsz, n_heads, 1, KV_CACHE_INCREMENT + seq_len), dtype = dtype, device = device
         )

unsloth/models/llama.py

@@ -397,7 +397,7 @@ def LlamaAttention_fast_forward_inference(
 
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
-            self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+            self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]
 

unsloth/models/qwen3.py

@@ -249,7 +249,7 @@ def Qwen3Attention_fast_forward_inference(
 
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
-            self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+            self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]