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 <danielhanchen@gmail.com>
This commit is contained in:
DoubleMathew
GitHub
parent
73e2ac372a
commit
9bad70b0c3
11 changed files with 437 additions and 152 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <pad> 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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue