From 9bad70b0c39efa3da4c9a8ceaf74439ebdd397e1 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Wed, 25 Feb 2026 09:21:04 -0600 Subject: [PATCH] Fix/pr 3699 leftpad prefill main (#4100) * Fix left-padding masks and positions in batched decode/prefill * Fix batched generation with left padding * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix attention mask handling, padding_idx zeroing, and Mistral batched generation 1. attention_dispatch.py: Fall back from flash/xformers to SDPA when an attention_mask is present, since flash attention only supports causal masking via flag and cannot consume arbitrary padding masks. 2. gemma2.py: Apply attention_mask during decode inference for bsz > 1. Guard against boolean SWA/GA flags with isinstance check. Slice mask to match K/V length when sliding window is active. Remove dead commented-out SDPA branch (SDPA does not support softcapping). 3. granite.py: Apply attention_mask during decode inference for bsz > 1. Remove dead commented-out SDPA branch and misleading comment. 4. mistral.py: Fix 2D-to-4D padding mask conversion -- convert 0/1 mask to additive format (0 for keep, -inf for mask) before combining with the causal mask. Force SDPA backend when attention_mask is present. 5. llama.py: Skip zeroing embed_tokens.weight[padding_idx] when the embedding is weight-tied to lm_head, since zeroing the shared weight forces logit(pad) = 0 which is higher than real token logits in models like Gemma, causing the decoder to emit pad tokens as gibberish. Also add eos != pad guard, clean up unused _seq_length variable, and fix get_max_cache_shape handling. 6. vision.py: Same padding_idx fix as llama.py for the vision model loading path. Tested on gemma-2b-it, gemma-2-2b-it, Llama-3.2-1B, Mistral-7B-v0.3, Qwen2.5-0.5B, Qwen3-0.6B with flash-attn 2.8.3 active. All outputs coherent, zero crashes, zero resize warnings. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Inference path optimizations: eliminate per-layer GPU-CPU sync, cache inspect.signature, add Granite SDPA split * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * More inference path optimizations across model files - gemma: hoist rotary_seq_len computation to model level (eliminates N per-layer GPU-CPU syncs from position_ids.max().item()), pre-convert attention mask to bool once for all layers, use scalar float multiply instead of torch.tensor allocation for embedding scaling - gemma2: use in-place tanh_() for softcap attention, use scalar float multiply for embedding scaling - granite: pre-convert attention mask to bool once for all layers - cohere: use in-place neg_() for rotary embedding (consistent with all other model files) - falcon_h1: use in-place mul_() for key_multiplier scaling - llama: use in-place tanh_() for logit softcapping * Revert scalar multiply for Gemma/Gemma2 embedding scaling The original torch.tensor(..., dtype=hidden_states.dtype) is intentional: sqrt(3072) rounds to 55.5 in bfloat16 vs 55.4256 in float32. A plain scalar multiply may compute at higher precision internally, producing different results. Restore the explicit dtype-cast tensor to match the training path in LlamaModel_fast_forward. * Fix hardcoded cuda:0 device strings and add Cohere .eq(0) bool mask Replace 15 hardcoded "cuda:0" with f"{DEVICE_TYPE_TORCH}:0" across gemma.py, gemma2.py, cohere.py, and falcon_h1.py to support multi-GPU and non-CUDA devices (XPU, etc.). Add .eq(0) bool mask pre-conversion in CohereModel_fast_forward_inference for batched inference consistency with llama.py, granite.py, and gemma.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Disable flex_attention for Mllama (Llama 3.2 Vision) Mllama's _update_causal_mask uses the deprecated make_flex_block_causal_mask which creates a BlockMask with Q_LEN=KV_LEN=total_seq_len. During decode with KV cache, q_len=1 but the block_mask still has Q_LEN=total_seq_len, causing a ValueError. This is an upstream transformers issue -- newer models use flex_attention_mask from masking_utils which handles decode correctly via cache_position, but mllama has not been updated yet. Add mllama to the exclusion list in prefer_flex_attn_if_supported alongside gpt_oss so it falls back to sdpa, which works correctly for both training and inference. * Fix off-by-one in sliding window K/V slicing for gemma2, qwen3, falcon_h1, cohere The old formula `slicing_tokens = 1 - sliding_window` uses negative indexing that keeps `sliding_window - 1` tokens instead of `sliding_window`. For example with sliding_window=32 and kv_seq_len=100, `1-32 = -31` keeps indices 69..99 (31 tokens) instead of the correct 68..99 (32 tokens). Replace with `start = kv_seq_len - sliding_window` to match the fix already applied in llama.py and the canonical definition in transformers masking_utils (sliding_window_overlay: kv_idx > q_idx - W, which keeps exactly W tokens). Also add attention_mask slicing after K/V trim in qwen3, falcon_h1, and cohere to prevent mask/K dimension mismatch during batched SDPA inference, matching the pattern already used in llama.py. Currently only gemma2 (sliding_window=4096) is actively affected. The other three models have sliding_window=None in their configs so the code path is not triggered, but this keeps it correct for any future models that set it. * Fix Gemma2 softcapping order: apply mask after softcap, not before The attention mask must be applied AFTER logit softcapping, not before. Both the Google DeepMind reference implementation (google-deepmind/gemma, gm/nn/_modules.py lines 254-277) and transformers' eager_attention_forward (gemma2/modeling_gemma2.py lines 187-193) use this order: 1. logits = Q @ K^T * scale 2. logits = tanh(logits / softcap) * softcap # softcap first 3. logits = logits + mask # mask after 4. probs = softmax(logits) The PR had the mask addition before softcapping, which causes tanh to clamp the -inf mask values to -softcap instead of preserving them as -inf for softmax. While the practical impact is small (masked positions get ~1e-23 probability instead of exact zero), this should match upstream. * Clarify GQA condition precedence and remove stale comments Add explicit parentheses to grouped query attention conditions in llama.py, qwen3.py, granite.py to make operator precedence clear. The expression `bsz == 1 or not X and Y` relies on Python binding `not` > `and` > `or` which is correct but easy to misread. Remove dead commented-out code (`# else: # Knn, Vnn = Knn, Vnn`) and stale mask comments (`# if attention_mask ...`) from the bsz==1 fast path in llama, qwen3, cohere, falcon_h1, gemma2 inference functions. These were leftover from the pre-batched-inference structure and no longer apply. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- unsloth/models/_utils.py | 11 +- unsloth/models/cohere.py | 47 +++--- unsloth/models/falcon_h1.py | 34 +++-- unsloth/models/gemma.py | 12 +- unsloth/models/gemma2.py | 39 ++--- unsloth/models/granite.py | 57 ++++--- unsloth/models/llama.py | 225 +++++++++++++++++++++------- unsloth/models/mistral.py | 22 ++- unsloth/models/qwen3.py | 42 ++++-- unsloth/models/vision.py | 15 +- unsloth/utils/attention_dispatch.py | 83 +++++++++- 11 files changed, 436 insertions(+), 151 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 5f4927b92f..dafdec4e1e 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -235,11 +235,14 @@ def prefer_flex_attn_if_supported(model_class, config): model_class, "_supports_flex_attn", False ): return None - # GPT-OSS uses eager attention during inference since flex attention - # returns incorrect results (likely due to left padding issues). - # Skip setting flex_attention to avoid BlockMask type errors. + # GPT-OSS and Mllama use eager/sdpa attention during inference since + # flex attention returns incorrect results or errors out. + # GPT-OSS: left padding issues cause incorrect outputs. + # Mllama: _update_causal_mask uses make_flex_block_causal_mask which + # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during + # decode q_len=1, causing ValueError. Needs transformers update. model_type = getattr(config, "model_type", "") if config else "" - if model_type == "gpt_oss": + if model_type in ("gpt_oss", "mllama"): return None if config is not None: setattr(config, "_attn_implementation", "flex_attention") diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index c33317ee02..4251f3acd9 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -188,7 +188,9 @@ def CohereDecoderLayer_fast_forward( self, "_flag_for_generation" ): # past_key_value is not None: out_weight = torch.empty( - self.input_layernorm.weight.shape, dtype = torch.float32, device = "cuda:0" + self.input_layernorm.weight.shape, + dtype = torch.float32, + device = f"{DEVICE_TYPE_TORCH}:0", ) # Self Attention @@ -254,6 +256,7 @@ def CohereAttention_fast_forward_inference( position_ids, do_prefill = False, attention_mask = None, + **kwargs, ): Xn = hidden_states bsz, _, hd = hidden_states.size() @@ -277,26 +280,28 @@ def CohereAttention_fast_forward_inference( self.paged_attention = torch.empty( (KV_CACHE_INCREMENT + seq_len + 1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, - device = "cuda:0", + device = f"{DEVICE_TYPE_TORCH}:0", ) self.paged_attention_K = self.paged_attention[:, 0] self.paged_attention_V = self.paged_attention[:, 1] self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) self.temp_QA = torch.empty( - (2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0" + (2, bsz, 1, attention_size), dtype = dtype, device = f"{DEVICE_TYPE_TORCH}:0" ) self.temp_KV = torch.empty( - (2, bsz, 1, n_kv_heads * head_dim), dtype = dtype, device = "cuda:0" + (2, bsz, 1, n_kv_heads * head_dim), + dtype = dtype, + device = f"{DEVICE_TYPE_TORCH}:0", ) self.RH_Q = torch.empty( - (bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0" + (bsz, n_heads, 1, head_dim), dtype = dtype, device = f"{DEVICE_TYPE_TORCH}:0" ) # 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 = f"{DEVICE_TYPE_TORCH}:0" ) else: self.temp_O = self.temp_QA[1][:, :, :hidden_size] @@ -304,17 +309,21 @@ def CohereAttention_fast_forward_inference( self.attention = torch.empty( (bsz, n_heads, 1, KV_CACHE_INCREMENT + seq_len), dtype = dtype, - device = "cuda:0", + device = f"{DEVICE_TYPE_TORCH}:0", ) self.scalar = 1.0 / math_sqrt(self.head_dim) self.half_head_dim = head_dim // 2 # Cohere has QK layernorms if self.use_qk_norm: self.q_norm_out_weight = torch.empty( - self.q_norm.weight.shape, dtype = torch.float32, device = "cuda:0" + self.q_norm.weight.shape, + dtype = torch.float32, + device = f"{DEVICE_TYPE_TORCH}:0", ) self.k_norm_out_weight = torch.empty( - self.k_norm.weight.shape, dtype = torch.float32, device = "cuda:0" + self.k_norm.weight.shape, + dtype = torch.float32, + device = f"{DEVICE_TYPE_TORCH}:0", ) else: self.q_norm_out_weight = None @@ -355,7 +364,7 @@ def CohereAttention_fast_forward_inference( RH_Q = self.RH_Q RH_Q[:, :, :, :h] = Qn[:, :, :, h:] RH_Q[:, :, :, h:] = Qn[:, :, :, :h] - torch.neg(RH_Q[:, :, :, :h], out = RH_Q[:, :, :, :h]) + RH_Q[:, :, :, :h].neg_() Qn *= cos Qn.addcmul_(RH_Q, sin) @@ -364,7 +373,7 @@ def CohereAttention_fast_forward_inference( ] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0") RH_K[:, :, :, :h] = Kn[:, :, :, h:] RH_K[:, :, :, h:] = Kn[:, :, :, :h] - torch.neg(RH_K[:, :, :, :h], out = RH_K[:, :, :, :h]) + RH_K[:, :, :, :h].neg_() Kn *= cos Kn.addcmul_(RH_K, sin) @@ -379,10 +388,11 @@ def CohereAttention_fast_forward_inference( # Handle sliding windows sliding_window = getattr(self.config, "sliding_window", None) if sliding_window is not None and kv_seq_len > sliding_window: - # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193 - slicing_tokens = 1 - sliding_window - Knn = Kn[:, :, slicing_tokens:, :] # .contiguous() - Vnn = Vn[:, :, slicing_tokens:, :] # .contiguous() + start = kv_seq_len - sliding_window + Knn = Kn[:, :, start:, :] # .contiguous() + Vnn = Vn[:, :, start:, :] # .contiguous() + if attention_mask is not None: + attention_mask = attention_mask[..., start:] else: Knn, Vnn = Kn, Vn @@ -397,9 +407,6 @@ def CohereAttention_fast_forward_inference( ) Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) - # else: - # Knn, Vnn = Knn, Vnn - # pass # Attention if bsz == 1: @@ -408,7 +415,6 @@ def CohereAttention_fast_forward_inference( A = torch_matmul( Qn, Knn.transpose(2, 3), out = self.attention[:, :, :, :cached_len] ) - # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched A[:] = torch_nn_functional_softmax( A, dim = -1, dtype = torch.float32 ) # .to(A.dtype) @@ -453,6 +459,9 @@ def CohereModel_fast_forward_inference( seq_len, sliding_window = getattr(self.config, "sliding_window", None), ) + # Pre-convert to bool once for all layers (avoids per-layer .eq(0)) + if attention_mask is not None and attention_mask.dtype != torch.bool: + attention_mask = attention_mask.eq(0) else: attention_mask = None diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index 428f49d727..6e3b16b21b 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -197,6 +197,7 @@ def FalconH1Attention_fast_forward_inference( position_ids, do_prefill = False, attention_mask = None, + **kwargs, ): """ https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406 @@ -265,7 +266,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] @@ -292,7 +293,7 @@ def FalconH1Attention_fast_forward_inference( Qn = fast_linear_forward(self.q_proj, Xn, out = self.temp_QA[0]) Kn = fast_linear_forward(self.k_proj, Xn, out = self.temp_KV[0]) - Kn = Kn * self.config.key_multiplier + Kn.mul_(self.config.key_multiplier) Vn = fast_linear_forward(self.v_proj, Xn, out = self.temp_KV[1]) Qn = Qn.view( bsz, 1, n_heads, head_dim @@ -343,10 +344,11 @@ def FalconH1Attention_fast_forward_inference( # Handle sliding windows sliding_window = getattr(self.config, "sliding_window", None) if sliding_window is not None and kv_seq_len > sliding_window: - # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193 - slicing_tokens = 1 - sliding_window - Knn = Kn[:, :, slicing_tokens:, :] # .contiguous() - Vnn = Vn[:, :, slicing_tokens:, :] # .contiguous() + start = kv_seq_len - sliding_window + Knn = Kn[:, :, start:, :] # .contiguous() + Vnn = Vn[:, :, start:, :] # .contiguous() + if attention_mask is not None: + attention_mask = attention_mask[..., start:] else: Knn, Vnn = Kn, Vn @@ -361,9 +363,6 @@ def FalconH1Attention_fast_forward_inference( ) Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) - # else: - # Knn, Vnn = Knn, Vnn - # pass # Attention if bsz == 1: @@ -372,7 +371,6 @@ def FalconH1Attention_fast_forward_inference( A = torch_matmul( Qn, Knn.transpose(2, 3), out = self.attention[:, :, :, :cached_len] ) - # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched A[:] = torch_nn_functional_softmax( A, dim = -1, dtype = torch.float32 ) # .to(A.dtype) @@ -533,11 +531,19 @@ def _FalconH1_fast_forward_inference( bsz, q_len, hd = X.shape assert q_len == 1 # Get saved buffers to reduce memory movement - residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") - _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") + residual = torch.empty( + (bsz, q_len, hd), dtype = torch.float32, device = f"{DEVICE_TYPE_TORCH}:0" + ) + _XX = torch.empty( + (2, bsz, q_len, hd), dtype = torch.float32, device = f"{DEVICE_TYPE_TORCH}:0" + ) XX, XX2 = _XX[0], _XX[1] - variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0") - temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0") + variance = torch.empty( + (bsz, q_len, 1), dtype = torch.float32, device = f"{DEVICE_TYPE_TORCH}:0" + ) + temp_mlp = torch.empty( + (2, bsz, 1, mlp_size), dtype = X.dtype, device = f"{DEVICE_TYPE_TORCH}:0" + ) temp_gate, temp_up = temp_mlp[0], temp_mlp[1] seq_len = past_key_values[0][0].shape[-2] if bsz != 1: diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 55a8c8697f..cf543ae094 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -97,7 +97,9 @@ def GemmaDecoderLayer_fast_forward( self, "_flag_for_generation" ): # past_key_value is not None: out_weight = torch.empty( - self.input_layernorm.weight.shape, dtype = torch.float32, device = "cuda:0" + self.input_layernorm.weight.shape, + dtype = torch.float32, + device = f"{DEVICE_TYPE_TORCH}:0", ) # Self Attention @@ -191,6 +193,7 @@ def GemmaModel_fast_forward_inference( bsz, q_len, hd = hidden_states.shape seq_len = past_key_values[0][0].shape[-2] + kv_seq_len = seq_len + 1 if bsz != 1: attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( attention_mask, @@ -198,6 +201,12 @@ def GemmaModel_fast_forward_inference( hidden_states, seq_len, ) + # Pre-convert to bool once for all layers (avoids per-layer .eq(0)) + if attention_mask is not None and attention_mask.dtype != torch.bool: + attention_mask = attention_mask.eq(0) + + # Compute rotary_seq_len once to avoid per-layer GPU-CPU sync from .item() + rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1) next_decoder_cache = [] for idx, decoder_layer in enumerate(self.model.layers): @@ -217,6 +226,7 @@ def GemmaModel_fast_forward_inference( position_ids = position_ids, attention_mask = attention_mask, do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"), + rotary_seq_len = rotary_seq_len, ) hidden_states += residual diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 03e77f6504..e59b8d5ebd 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -222,7 +222,9 @@ def Gemma2DecoderLayer_fast_forward( self, "_flag_for_generation" ): # past_key_value is not None: out_weight = torch.empty( - self.input_layernorm.weight.shape, dtype = torch.float32, device = "cuda:0" + self.input_layernorm.weight.shape, + dtype = torch.float32, + device = f"{DEVICE_TYPE_TORCH}:0", ) # Self Attention @@ -352,7 +354,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 ) @@ -399,7 +401,7 @@ def Gemma2Attention_fast_forward_inference( RH_Q = self.RH_Q RH_Q[:, :, :, :h] = Qn[:, :, :, h:] RH_Q[:, :, :, h:] = Qn[:, :, :, :h] - torch.neg(RH_Q[:, :, :, :h], out = RH_Q[:, :, :, :h]) + RH_Q[:, :, :, :h].neg_() Qn *= cos Qn.addcmul_(RH_Q, sin) @@ -408,7 +410,7 @@ def Gemma2Attention_fast_forward_inference( ] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0") RH_K[:, :, :, :h] = Kn[:, :, :, h:] RH_K[:, :, :, h:] = Kn[:, :, :, :h] - torch.neg(RH_K[:, :, :, :h], out = RH_K[:, :, :, :h]) + RH_K[:, :, :, :h].neg_() Kn *= cos Kn.addcmul_(RH_K, sin) @@ -423,10 +425,9 @@ def Gemma2Attention_fast_forward_inference( # Handle sliding windows sliding_window = self.config.sliding_window if use_sliding_window and kv_seq_len > sliding_window: - # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193 - slicing_tokens = 1 - sliding_window - Knn = Kn[:, :, slicing_tokens:, :] # .contiguous() - Vnn = Vn[:, :, slicing_tokens:, :] # .contiguous() + start = kv_seq_len - sliding_window + Knn = Kn[:, :, start:, :] # .contiguous() + Vnn = Vn[:, :, start:, :] # .contiguous() else: Knn, Vnn = Kn, Vn @@ -441,28 +442,32 @@ def Gemma2Attention_fast_forward_inference( ) Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) - # else: - # Knn, Vnn = Knn, Vnn - # pass # Attention - # if bsz == 1: + # [TODO] Gemma2 uses manual matmul for all batch sizes because SDPA does + # not support softcapping (tanh logit scaling). If a future PyTorch adds + # a softcap param to scaled_dot_product_attention, consider using SDPA + # for bsz > 1 to match the llama/qwen3 pattern. Qn *= ( self.scalar ) # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963 # It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows A = torch_matmul(Qn, Knn.transpose(2, 3), out = self.attention[:, :, :, :cached_len]) - # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched + # Softcapping must happen BEFORE the mask is applied. + # Reference: google-deepmind/gemma _modules.py and transformers gemma2 eager_attention_forward A *= self.reciprocal_t - torch_tanh(A, out = A) + A.tanh_() A *= self.t # Logit softcapping + if attention_mask is not None and isinstance(attention_mask, torch.Tensor): + # Slice mask to match K/V when sliding window is active + if attention_mask.shape[-1] != A.shape[-1]: + attention_mask = attention_mask[:, :, :, -A.shape[-1] :] + A += attention_mask + A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32) # .to(A.dtype) A = torch_matmul(A, Vnn, out = Qn) - # else: - # A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False) - # pass A = A.transpose(1, 2) A = A.reshape(bsz, 1, attention_size) A = fast_linear_forward(self.o_proj, A, out = self.temp_O) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 168df90f4c..79ac41c43f 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -323,8 +323,7 @@ def GraniteAttention_fast_forward_inference( (2, bsz, 1, n_kv_heads * head_dim), dtype = dtype, device = device ) 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 ) @@ -362,7 +361,7 @@ def GraniteAttention_fast_forward_inference( RH_Q = self.RH_Q RH_Q[:, :, :, :h] = Qn[:, :, :, h:] RH_Q[:, :, :, h:] = Qn[:, :, :, :h] - torch.neg(RH_Q[:, :, :, :h], out = RH_Q[:, :, :, :h]) + RH_Q[:, :, :, :h].neg_() Qn *= cos Qn.addcmul_(RH_Q, sin) @@ -371,7 +370,7 @@ def GraniteAttention_fast_forward_inference( ] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0") RH_K[:, :, :, :h] = Kn[:, :, :, h:] RH_K[:, :, :, h:] = Kn[:, :, :, :h] - torch.neg(RH_K[:, :, :, :h], out = RH_K[:, :, :, :h]) + RH_K[:, :, :, :h].neg_() Kn *= cos Kn.addcmul_(RH_K, sin) @@ -385,7 +384,7 @@ def GraniteAttention_fast_forward_inference( # Grouped query attention _, _, cached_len, _ = Kn.shape - if n_groups != 1: + if bsz == 1 or ((not SDPA_HAS_GQA) and n_groups != 1): Kn = Kn[:, :, None, :, :].expand( bsz, n_kv_heads, n_groups, cached_len, head_dim ) @@ -394,20 +393,39 @@ def GraniteAttention_fast_forward_inference( ) Kn = Kn.reshape(bsz, n_heads, cached_len, head_dim) Vn = Vn.reshape(bsz, n_heads, cached_len, head_dim) - # else: - # Kn, Vn = Kn, Vn - # pass - Qn *= self.scaling - A = torch_matmul(Qn, Kn.transpose(2, 3), out = self.attention[:, :, :, :cached_len]) - - # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched - - A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32) # .to(A.dtype) - A = torch_matmul(A, Vn, out = Qn) - # else: - # A = scaled_dot_product_attention(Qn, Kn, Vn, attn_mask = attention_mask, is_causal = False) - # pass + # Attention + if bsz == 1: + Qn *= self.scaling + A = torch_matmul( + Qn, Kn.transpose(2, 3), out = self.attention[:, :, :, :cached_len] + ) + A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32) + A = torch_matmul(A, Vn, out = Qn) + else: + if ( + attention_mask is not None + and attention_mask.dim() == 4 + and attention_mask.dtype != torch.bool + ): + attention_mask = attention_mask.eq(0) + if SDPA_HAS_GQA: + A = scaled_dot_product_attention( + Qn, + Kn, + Vn, + attn_mask = attention_mask, + scale = self.scaling, + enable_gqa = True, + ) + else: + A = scaled_dot_product_attention( + Qn, + Kn, + Vn, + attn_mask = attention_mask, + scale = self.scaling, + ) A = A.transpose(1, 2) A = A.reshape(bsz, 1, attention_size) A = fast_linear_forward(self.o_proj, A, out = self.temp_O) @@ -442,6 +460,9 @@ def GraniteModel_fast_forward_inference( hidden_states, seq_len, ) + # Pre-convert to bool once for all layers (avoids per-layer .eq(0)) + if attention_mask is not None and attention_mask.dtype != torch.bool: + attention_mask = attention_mask.eq(0) else: attention_mask = None diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 342366f02e..f80a55fdd2 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -36,6 +36,7 @@ from ..utils.attention_dispatch import ( AttentionConfig, AttentionContext, run_attention, + SDPA, select_attention_backend, ) from torch.nn.functional import scaled_dot_product_attention @@ -213,11 +214,22 @@ def _fast_prepare_inputs_for_generation( **kwargs, ): past_key_values = kwargs.get("past_key_values", None) + original_attention_mask = attention_mask # Handle inputs_embeds - only use on FIRST generation step (no cache) # This fixes GitHub issue #3798: inputs_embeds was ignored use_inputs_embeds = inputs_embeds is not None and past_key_values is None + if input_ids is not None and input_ids.numel() > 0: + bs, seq_length = input_ids.shape + device = input_ids.device + elif inputs_embeds is not None: + bs, seq_length, _ = inputs_embeds.shape + device = inputs_embeds.device + else: + bs, seq_length = 1, 0 + device = "cuda" if torch.cuda.is_available() else "cpu" + if past_key_values is not None: # Check for uninitialized DynamicCache if len(past_key_values) == 0: @@ -234,16 +246,47 @@ def _fast_prepare_inputs_for_generation( use_inputs_embeds = inputs_embeds is not None else: if input_ids is not None and input_ids.numel() > 0: - bs, cache_length = input_ids.shape + bs = input_ids.shape[0] input_ids = input_ids[:, [-1]] device = input_ids.device + seq_length = 1 elif inputs_embeds is not None: - bs, cache_length, _ = inputs_embeds.shape + bs, seq_length, _ = inputs_embeds.shape device = inputs_embeds.device else: - bs, cache_length = 1, 0 + bs, seq_length = 1, 0 device = "cuda" if torch.cuda.is_available() else "cpu" + if hasattr(past_key_values, "get_seq_length"): + past_len = int(past_key_values.get_seq_length()) + else: + # legacy tuple cache: (layer, (K,V)) + past_len = int(past_key_values[0][0].shape[-2]) + + max_cache_len = None + if hasattr(past_key_values, "get_max_cache_shape"): + m = past_key_values.get_max_cache_shape() + max_cache_len = int(m) if m is not None and m > 0 else None + elif hasattr(past_key_values, "get_max_length"): + m = past_key_values.get_max_length() + max_cache_len = int(m) if m is not None else None + + # ensure cache_position + cache_position = kwargs.get("cache_position", None) + if cache_position is None: + kwargs["cache_position"] = torch.arange( + past_len, + past_len + seq_length, + device = device, + dtype = torch.long, + ) + else: + if ( + hasattr(cache_position, "device") + and cache_position.device != device + ): + kwargs["cache_position"] = cache_position.to(device) + # Get to the base model base_model = self if hasattr(base_model, "base_model_prefix"): @@ -252,45 +295,49 @@ def _fast_prepare_inputs_for_generation( if hasattr( base_model, "_prepare_4d_causal_attention_mask_with_cache_position" ): + if not hasattr(base_model, "_unsloth_mask_needs_device"): - def needs_device_kw(fn) -> bool: - try: - sig = inspect.signature(inspect.unwrap(fn)) - return "device" in sig.parameters - except: - # transformers <= 4.51.3 includes device arg but > 4.51.3 does not - return transformers_version < Version("4.52.0") + def _check_needs_device(fn) -> bool: + try: + sig = inspect.signature(inspect.unwrap(fn)) + return "device" in sig.parameters + except: + # transformers <= 4.51.3 includes device arg but > 4.51.3 does not + return transformers_version < Version("4.52.0") - kwargs = { - "sequence_length": 1, - "target_length": cache_length, + base_model._unsloth_mask_needs_device = _check_needs_device( + base_model._prepare_4d_causal_attention_mask_with_cache_position + ) + + if max_cache_len is not None: + target_length = max_cache_len + elif ( + original_attention_mask is not None + and original_attention_mask.dim() == 2 + ): + target_length = original_attention_mask.shape[-1] + else: + target_length = past_len + seq_length + + mask_kwargs = { + "sequence_length": seq_length, + "target_length": target_length, "dtype": self.dtype, - "cache_position": torch.arange( - cache_length, cache_length + 1, device = device - ), + "cache_position": kwargs["cache_position"], "batch_size": bs, "config": self.config, "past_key_values": past_key_values, } - try: - if needs_device_kw( - base_model._prepare_4d_causal_attention_mask_with_cache_position - ): - kwargs["device"] = device - except: - print( - f"Unsloth: Could not inspect signature of {base_model._prepare_4d_causal_attention_mask_with_cache_position}" - ) + if base_model._unsloth_mask_needs_device: + mask_kwargs["device"] = device attention_mask = ( base_model._prepare_4d_causal_attention_mask_with_cache_position( attention_mask, - **kwargs, + **mask_kwargs, ) ) else: - if attention_mask is not None: - attention_mask = attention_mask[:, [-1]] if transformers_version <= Version("4.52.4"): logger.warning_once( f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method " @@ -299,8 +346,17 @@ def _fast_prepare_inputs_for_generation( "issue on GitHub." ) - if "cache_position" in kwargs: - kwargs["position_ids"] = kwargs["cache_position"] + if kwargs.get("position_ids", None) is None: + if original_attention_mask is not None and original_attention_mask.dim() == 2: + position_ids = original_attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(original_attention_mask == 0, 1) + position_ids = position_ids[:, -seq_length:] + kwargs["position_ids"] = position_ids + elif kwargs.get("cache_position", None) is not None: + cp = kwargs["cache_position"] + if cp.dim() == 1: + cp = cp.unsqueeze(0).expand(bs, -1) + kwargs["position_ids"] = cp result = { "attention_mask": attention_mask, @@ -330,6 +386,7 @@ def LlamaAttention_fast_forward_inference( position_ids, do_prefill = False, attention_mask = None, + rotary_seq_len = None, ): """ https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406 @@ -398,7 +455,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] @@ -435,10 +492,19 @@ def LlamaAttention_fast_forward_inference( # Need to do it prior 2 steps before hitting full on short KV cache # or else error - self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2) - cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) - cos = cos[position_ids].unsqueeze(1) - sin = sin[position_ids].unsqueeze(1) + # ensure correct shape + if position_ids.dim() == 1: + position_ids = position_ids[:, None] + position_ids = position_ids.to(Qn.device) + + if rotary_seq_len is None: + rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1) + self.rotary_emb.extend_rope_embedding(Vn, rotary_seq_len + 1) # +1 slack + cos, sin = self.rotary_emb.get_cached(rotary_seq_len, Qn.device.index or 0) + + cos = cos[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype) + sin = sin[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype) + h = self.half_head_dim RH_Q = self.RH_Q @@ -469,15 +535,17 @@ def LlamaAttention_fast_forward_inference( sliding_window = getattr(self.config, "sliding_window", None) if sliding_window is not None and kv_seq_len > sliding_window: # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193 - slicing_tokens = 1 - sliding_window - Knn = Kn[:, :, slicing_tokens:, :] # .contiguous() - Vnn = Vn[:, :, slicing_tokens:, :] # .contiguous() + start = kv_seq_len - sliding_window + Knn = Kn[:, :, start:, :] # .contiguous() + Vnn = Vn[:, :, start:, :] # .contiguous() + if attention_mask is not None: + attention_mask = attention_mask[..., start:] else: Knn, Vnn = Kn, Vn # Grouped query attention _, _, cached_len, _ = Knn.shape - if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1: + if bsz == 1 or ((not SDPA_HAS_GQA) and n_groups != 1): Knn = Knn[:, :, None, :, :].expand( bsz, n_kv_heads, n_groups, cached_len, head_dim ) @@ -486,9 +554,6 @@ def LlamaAttention_fast_forward_inference( ) Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) - # else: - # Knn, Vnn = Knn, Vnn - # pass # when qlen==vlen and attn_mask is None, we should use causal attention Q_len = Qn.shape[-2] @@ -504,12 +569,23 @@ def LlamaAttention_fast_forward_inference( A = torch_matmul( Qn, Knn.transpose(2, 3), out = self.attention[:, :, :, :cached_len] ) - # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched A[:] = torch_nn_functional_softmax( A, dim = -1, dtype = torch.float32 ) # .to(A.dtype) A = torch_matmul(A, Vnn, out = Qn) + # --- attention_mask fixup for SDPA if user passes 2D padding mask else: + if attention_mask is not None and attention_mask.dim() == 2: + attention_mask = attention_mask[:, None, None, :].to(torch.bool) + # is it more appropriate to use _prepare_4d_causal_attention_mask_for_sdpa? + elif ( + attention_mask is not None + and attention_mask.dim() == 4 + and attention_mask.dtype != torch.bool + ): + # Decode is more stable with boolean keep masks than additive bf16 masks. + attention_mask = attention_mask.eq(0) + if SDPA_HAS_GQA: A = scaled_dot_product_attention( Qn, @@ -663,6 +739,8 @@ def LlamaAttention_fast_forward( rotary_emb = self.rotary_emb rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len) cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index) + cos = cos.to(device = Q.device, dtype = Q.dtype) + sin = sin.to(device = Q.device, dtype = Q.dtype) rope_position_ids = position_ids if rope_position_ids is None and seq_info is not None: @@ -682,7 +760,11 @@ def LlamaAttention_fast_forward( # Attention module use_varlen = seq_info is not None and past_key_value is None - backend = select_attention_backend(use_varlen) + backend = ( + SDPA if attention_mask is not None else select_attention_backend(use_varlen) + ) + + # should dropout be hardcoded to 0.0? config = AttentionConfig( backend = backend, n_kv_heads = n_kv_heads, @@ -1257,7 +1339,8 @@ def _LlamaModel_fast_forward_inference( ) seq_len = past_key_values[0][0].shape[-2] - if bsz != 1: + kv_seq_len = seq_len + 1 + if attention_mask is not None: attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( attention_mask, (bsz, q_len), @@ -1265,9 +1348,15 @@ def _LlamaModel_fast_forward_inference( seq_len, sliding_window = getattr(self.config, "sliding_window", None), ) + # Pre-convert to bool once for all layers (avoids per-layer .eq(0)) + if attention_mask is not None and attention_mask.dtype != torch.bool: + attention_mask = attention_mask.eq(0) else: attention_mask = None + # Compute rotary_seq_len once to avoid per-layer GPU-CPU sync from .item() + rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1) + next_decoder_cache = [] for idx, decoder_layer in enumerate(self.model.layers): @@ -1290,6 +1379,7 @@ def _LlamaModel_fast_forward_inference( position_ids = position_ids, attention_mask = attention_mask, do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"), + rotary_seq_len = rotary_seq_len, ) X += residual @@ -1528,7 +1618,7 @@ def CausalLM_fast_forward(fast_forward_inference): logits = logit_softcapping * logits else: logits *= 1.0 / logit_softcapping - torch.tanh(logits, out = logits) + logits.tanh_() logits *= logit_softcapping if not return_dict: @@ -2586,16 +2676,39 @@ class FastLlamaModel: model._old_generate = model.generate unsloth_fast_generate.__doc__ = model._old_generate.__doc__ model.generate = types.MethodType(unsloth_fast_generate, model) - # Set weight[padding_idx] = 0 - with torch.no_grad(): - for name, module in model.named_modules(): - if type(module) is torch.nn.Embedding: - if ( - getattr(module, "weight", None) is not None - and getattr(module, "padding_idx", None) is not None - ): - if module.padding_idx < module.weight.shape[0]: - module.weight[module.padding_idx] = 0 + # Set weight[padding_idx] = 0 for embeddings that are NOT tied with the + # lm_head. When weights are tied, zeroing the padding row also zeros + # the corresponding lm_head row, forcing logit = 0 for the pad token. + # This is higher than the (negative) logits for real tokens in models + # like Gemma, causing the decoder to emit and produce gibberish. + # Skip entirely if eos_token == pad_token to avoid zeroing EOS embedding. + eos_token_id = ( + getattr(tokenizer, "eos_token_id", None) if tokenizer is not None else None + ) + pad_token_id = ( + getattr(tokenizer, "pad_token_id", None) if tokenizer is not None else None + ) + if tokenizer is not None and eos_token_id != pad_token_id: + lm_head = getattr(model, "lm_head", None) + lm_head_weight = ( + getattr(lm_head, "weight", None) if lm_head is not None else None + ) + with torch.no_grad(): + for name, module in model.named_modules(): + if type(module) is torch.nn.Embedding: + if ( + getattr(module, "weight", None) is not None + and getattr(module, "padding_idx", None) is not None + ): + if module.padding_idx < module.weight.shape[0]: + # Skip if tied to lm_head + if ( + lm_head_weight is not None + and module.weight.data_ptr() + == lm_head_weight.data_ptr() + ): + continue + module.weight[module.padding_idx] = 0 return model, tokenizer @staticmethod diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 5e893d2b6f..83e9ab9486 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -25,6 +25,7 @@ from ..utils.attention_dispatch import ( AttentionConfig, AttentionContext, run_attention, + SDPA, select_attention_backend, ) from .llama import ( @@ -115,7 +116,9 @@ def MistralAttention_fast_forward( use_varlen = ( seq_info is not None and past_key_value is None and window_size == (-1, -1) ) - backend = select_attention_backend(use_varlen) + backend = ( + SDPA if attention_mask is not None else select_attention_backend(use_varlen) + ) attention_config = AttentionConfig( backend = backend, n_kv_heads = n_kv_heads, @@ -216,13 +219,18 @@ def MistralForCausalLM_fast_forward( bsz, 1, q_len, q_len ) else: - # attention_mask should be [bsz, 1, q_len, q_len] or broadcastable - # Add causal mask to existing attention mask if attention_mask.dim() == 2: - # [bsz, seq_len] -> [bsz, 1, 1, seq_len] - attention_mask = attention_mask[:, None, None, :] - attention_mask = attention_mask.expand(bsz, 1, q_len, q_len) - attention_mask = attention_mask + causal_mask_values[None, None, :, :] + # Convert 0/1 padding mask to additive format: 1->0 (keep), 0->-inf (mask) + padding_mask = torch.where( + attention_mask[:, None, None, :].bool(), + 0.0, + -torch.inf, + ) + attention_mask = causal_mask_values[None, None, :, :] + padding_mask + else: + attention_mask = ( + attention_mask + causal_mask_values[None, None, :, :] + ) attention_mask = attention_mask.to( dtype = _get_dtype(dtype_from_config(self.config)) diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index ea06016d72..b93dddb186 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -21,6 +21,7 @@ from ..utils.attention_dispatch import ( AttentionConfig, AttentionContext, run_attention, + SDPA, select_attention_backend, ) from .llama import ( @@ -139,7 +140,9 @@ def Qwen3Attention_fast_forward( # Attention module use_varlen = seq_info is not None and past_key_value is None - backend = select_attention_backend(use_varlen) + backend = ( + SDPA if attention_mask is not None else select_attention_backend(use_varlen) + ) attention_config = AttentionConfig( backend = backend, n_kv_heads = n_kv_heads, @@ -181,6 +184,7 @@ def Qwen3Attention_fast_forward_inference( position_ids, do_prefill = False, attention_mask = None, + **kwargs, ): """ https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406 @@ -249,7 +253,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] @@ -329,24 +333,42 @@ def Qwen3Attention_fast_forward_inference( # Handle sliding windows sliding_window = getattr(self.config, "sliding_window", None) if sliding_window is not None and kv_seq_len > sliding_window: - # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193 - slicing_tokens = 1 - sliding_window - Knn = Kn[:, :, slicing_tokens:, :] # .contiguous() - Vnn = Vn[:, :, slicing_tokens:, :] # .contiguous() + start = kv_seq_len - sliding_window + Knn = Kn[:, :, start:, :] # .contiguous() + Vnn = Vn[:, :, start:, :] # .contiguous() + if attention_mask is not None: + attention_mask = attention_mask[..., start:] else: Knn, Vnn = Kn, Vn # when qlen==vlen and attn_mask is None, we should use causal attention Q_len = Qn.shape[-2] K_len = Knn.shape[-2] + if attention_mask is not None and attention_mask.dim() == 2: + attention_mask = attention_mask[:, None, None, :].to(torch.bool) + elif ( + attention_mask is not None + and attention_mask.dim() == 4 + and attention_mask.dtype != torch.bool + ): + attention_mask = attention_mask.eq(0) if attention_mask is None and Q_len == K_len: is_causal = True else: is_causal = False + use_sdpa_gqa = SDPA_HAS_GQA + if ( + use_sdpa_gqa + and isinstance(attention_mask, torch.Tensor) + and attention_mask.dim() >= 3 + and attention_mask.shape[0] > 1 + ): + # Avoid SDPA GQA drift for batched masked decode. + use_sdpa_gqa = False # Grouped query attention _, _, cached_len, _ = Knn.shape - if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1: + if bsz == 1 or ((not use_sdpa_gqa) and n_groups != 1): Knn = Knn[:, :, None, :, :].expand( bsz, n_kv_heads, n_groups, cached_len, head_dim ) @@ -355,9 +377,6 @@ def Qwen3Attention_fast_forward_inference( ) Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) - # else: - # Knn, Vnn = Knn, Vnn - # pass # Attention if bsz == 1: @@ -366,13 +385,12 @@ def Qwen3Attention_fast_forward_inference( A = torch_matmul( Qn, Knn.transpose(2, 3), out = self.attention[:, :, :, :cached_len] ) - # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched A[:] = torch_nn_functional_softmax( A, dim = -1, dtype = torch.float32 ) # .to(A.dtype) A = torch_matmul(A, Vnn, out = Qn) else: - if SDPA_HAS_GQA: + if use_sdpa_gqa: A = scaled_dot_product_attention( Qn, Knn, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 24e72e1535..e6c859d44e 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1410,9 +1410,15 @@ class FastBaseModel: m.for_training = functools.partial(FastBaseModel.for_training, m) m.for_inference = functools.partial(FastBaseModel.for_inference, m) m = m.model - # Set weight[padding_idx] = 0 + # Set weight[padding_idx] = 0 for embeddings that are NOT tied with the + # lm_head. When weights are tied, zeroing the padding row also zeros + # the corresponding lm_head row, forcing logit = 0 for the pad token. # Only do this if tokenizer is defined since eos_token == pad_token sometimes! pad_token_id = getattr(tokenizer, "pad_token_id", None) + lm_head = getattr(model, "lm_head", None) + lm_head_weight = ( + getattr(lm_head, "weight", None) if lm_head is not None else None + ) if ( tokenizer is not None and getattr(tokenizer, "eos_token_id", None) != pad_token_id @@ -1428,6 +1434,13 @@ class FastBaseModel: module.padding_idx == pad_token_id and module.padding_idx < module.weight.shape[0] ): + # Skip if tied to lm_head + if ( + lm_head_weight is not None + and module.weight.data_ptr() + == lm_head_weight.data_ptr() + ): + continue module.weight[module.padding_idx] = 0 return model diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index a7620549be..72d52ab376 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -20,6 +20,7 @@ from __future__ import annotations from dataclasses import dataclass from typing import Any, Optional, Tuple +import torch from torch import Tensor from torch.nn.functional import scaled_dot_product_attention @@ -119,6 +120,19 @@ def run_attention( backend = config.backend if backend == FLASH_VARLEN and context.seq_info is None: backend = FLASH_DENSE if HAS_FLASH_ATTENTION else SDPA + + # [TODO] Flash attention does not support arbitrary attention masks (only + # causal via flag). When a padding mask is present (e.g. left-padded + # batched generation), fall back to SDPA which consumes attn_mask. + # xFormers also does not thread context.attention_mask through, so the + # same fallback applies. + if context.attention_mask is not None and backend in ( + FLASH_DENSE, + FLASH_VARLEN, + XFORMERS, + ): + backend = SDPA + flash_dense_kwargs = config.flash_dense_kwargs or {} flash_varlen_kwargs = config.flash_varlen_kwargs or {} sdpa_kwargs = config.sdpa_kwargs or {} @@ -234,14 +248,79 @@ def run_attention( else: q_len_local = Q.shape[-2] k_len_local = K.shape[-2] + # ---- SDPA mask normalization for left padding / 2D masks ---- + if local_mask is not None and isinstance(local_mask, torch.Tensor): + local_mask = local_mask.to(device = Q.device) + + if local_mask.dim() == 2: + # key padding keep mask: (bsz, k_len), 1/True = real token + if local_mask.dtype == torch.bool: + key_keep = local_mask + else: + # tokenizer attention_mask is typically int 0/1 + key_keep = local_mask != 0 + + past_len = ( + k_len_local - q_len_local + ) # works for prefill (0) and decode + q_pos = torch.arange( + past_len, past_len + q_len_local, device = Q.device + ) + k_pos = torch.arange(k_len_local, device = Q.device) + + causal_keep = ( + k_pos[None, :] <= q_pos[:, None] + ) # True = allowed (SDPA) + if sliding_window is not None: + causal_keep &= k_pos[None, :] >= ( + q_pos[:, None] - (sliding_window - 1) + ) + + # (bsz, 1, q_len, k_len) boolean keep mask + local_mask = ( + causal_keep[None, None, :, :] & key_keep[:, None, None, :] + ) + + elif local_mask.dim() == 3: + # (bsz, q_len, k_len) -> (bsz, 1, q_len, k_len) + local_mask = local_mask[:, None, :, :] + + elif local_mask.dim() == 4: + if local_mask.dtype != torch.bool: + # Use boolean keep masks for better SDPA stability. + local_mask = local_mask.eq(0) + else: + raise ValueError( + f"Unsupported SDPA attention_mask rank: {local_mask.dim()}" + ) + + # Avoid NaNs from fully-masked rows (common with left padding). + if local_mask.dtype == torch.bool: + no_allowed = ~local_mask.any( + dim = -1, keepdim = True + ) # (bsz,1,q_len,1) + local_mask = local_mask | no_allowed + is_causal_local = local_mask is None and q_len_local == k_len_local kwargs = dict(sdpa_kwargs) kwargs.setdefault("attn_mask", local_mask) kwargs.setdefault("is_causal", is_causal_local) - if SDPA_HAS_GQA: - kwargs.setdefault("enable_gqa", config.n_groups != 1) + use_sdpa_gqa = SDPA_HAS_GQA and config.n_groups != 1 + if ( + use_sdpa_gqa + and (not requires_grad) + and isinstance(local_mask, torch.Tensor) + and local_mask.dim() >= 3 + and local_mask.shape[0] > 1 + ): + # Batched masked inference has shown row-coupled drift with SDPA GQA. + # Fall back to explicit KV expansion for deterministic row-wise behavior. + use_sdpa_gqa = False + + if use_sdpa_gqa: + kwargs.setdefault("enable_gqa", True) out = scaled_dot_product_attention(Q, K, V, **kwargs) return out.transpose(1, 2)