[ServeurpersoCom]

Overview

Embedding and encoder models run attention without a KV cache, and that path did not pad the keys/values to a multiple of 256 like the KV cache does. Some CUDA flash attention kernels (head size 512, e.g. EmbeddingGemma 2) need that padding, so attention silently fell back to the CPU for most input lengths, and -fa auto turned flash attention off entirely.

This pads the mask and K/V to 256 in the no-cache path. With EmbeddingGemma 2 on CUDA, 3841 tokens goes from 363 ms to 58 ms, -fa auto keeps flash attention on, and the embeddings are unchanged.

Additional information

Fixes #30398

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[deadprogram]

Hello, I ran into a failure in the yzma test suite, and tracked it back here. The following is an edited version of my automated analysis.

The T5 encoder now aborts in llama_encode since #30416

ggml.c:2072: GGML_ASSERT(ggml_can_repeat(b, a)) failed
  llm_graph_context::build_attn_mha
  llm_graph_context::build_attn(llm_graph_input_attn_no_cache *, ...)
  llama_model_t5::graph<true>
  llama_model_t5encoder::build_arch_graph
  llama_context::encode

T5 passes a relative position bias as kq_b, built from pos_bucket_enc with shape [n_tokens, n_tokens]. The mask and K/V are now padded to 256, so kq is [n_kv, n_tokens] and the add with kq_b fails. With the pad now unconditional it fails for every batch size.

b11574 works and b11576 and b11583 fail. You can repro with https://huggingface.co/callgg/t5-base-encoder-f32/resolve/main/t5base-encoder-q4_0.gguf by calling llama_encode on any short prompt.

Thank you.

[deadprogram]

Please see above @ggerganov

[ServeurpersoCom]

Hello, I ran into a failure in the yzma test suite, and tracked it back here. The following is an edited version of my automated analysis.

The T5 encoder now aborts in llama_encode since #30416

ggml.c:2072: GGML_ASSERT(ggml_can_repeat(b, a)) failed
  llm_graph_context::build_attn_mha
  llm_graph_context::build_attn(llm_graph_input_attn_no_cache *, ...)
  llama_model_t5::graph<true>
  llama_model_t5encoder::build_arch_graph
  llama_context::encode

T5 passes a relative position bias as kq_b, built from pos_bucket_enc with shape [n_tokens, n_tokens]. The mask and K/V are now padded to 256, so kq is [n_kv, n_tokens] and the add with kq_b fails. With the pad now unconditional it fails for every batch size.

b11574 works and b11576 and b11583 fail. You can repro with https://huggingface.co/callgg/t5-base-encoder-f32/resolve/main/t5base-encoder-q4_0.gguf by calling llama_encode on any short prompt.

Thank you.

Thanks for the report and the precise analysis! The T5 relative position bias was not padded along with K and V, it's fixed in #30446. I tested it with the t5-base-encoder q4_0 you linked: lengths from 2 to 511 tokens now work on CUDA and CPU, and the lengths that already worked give bit exact results. Could you also give it a try with the yzma test suite and confirm it works on your side?