* 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>
9bad70b0c3