From 668ded4a1afd6dbed8b87733b32009e667ffe142 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 24 Feb 2026 01:26:38 +0000 Subject: [PATCH] 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 | 2 +- unsloth/models/falcon_h1.py | 2 +- unsloth/models/gemma2.py | 2 +- unsloth/models/granite.py | 2 +- unsloth/models/llama.py | 2 +- unsloth/models/qwen3.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index c33317ee02..25aa3c9960 100644 --- a/unsloth/models/cohere.py +++ b/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] diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index 428f49d727..ac970d0b2f 100644 --- a/unsloth/models/falcon_h1.py +++ b/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] diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 03e77f6504..b6e125d688 100644 --- a/unsloth/models/gemma2.py +++ b/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 ) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 168df90f4c..be753dad38 100644 --- a/unsloth/models/granite.py +++ b/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 ) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a3fc9ab49a..23d29499a1 100644 --- a/unsloth/models/llama.py +++ b/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] diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index ea06016d72..0773a68309 100644 --- a/unsloth/models/qwen3.py +++ b/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]