Qwen3/Gemma3 candle skip attention masks for equal-length batches

[malaiwah]

Bug: Qwen3/Gemma3 candle skip attention masks for equal-length batches

System Info

• TEI version: source checkout at huggingface/text-embeddings-inference@main when filed, package version 1.9.3
• Local validation OS: macOS arm64 (Darwin 25.5.0)
• Repository Rust toolchain: rustc 1.92.0, cargo 1.92.0, matching rust-toolchain.toml
• Backend: candle CPU/Metal path
• Affected model families: Qwen3 embedding models and Gemma3 embedding models
• Original real-model reproduction: Qwen/Qwen3-Embedding-0.6B, CPU, fp32 and fp16; also reproduced without MKL
• Synthetic reproductions: tiny random-weight Qwen3 and Gemma3 models, no model download required

Information

The problem arises when using:

• The CLI directly
• Docker, depending on backend/device path

Tasks

The thing I am working on is:

• An officially supported command
• My own modifications

Reproduction

On the candle backend, Qwen3 and Gemma3 can skip their attention masks for multi-sequence batches where every sequence has the same length.

For Qwen3, Qwen3Model::forward only applied the causal mask when a padding bias existed. Equal-length batches need no padding, so attention_bias stayed None and the batch ran without the causal mask.

For Gemma3, Gemma3Model::forward similarly passed None when no padding was required. Each Gemma3 attention layer only builds and adds its causal or sliding-window mask when it receives an attention bias tensor, so equal-length batches skipped those masks too.

I added no-download synthetic-weight tests for both models:

cargo test -p text-embeddings-backend-candle --test test_qwen3_equal_length_batch --test test_gemma3_equal_length_batch

With the fixes applied, both pass:

test test_gemma3_equal_length_batch_matches_single ... ok
test test_qwen3_equal_length_batch_matches_single ... ok

Negative evidence with only the old Qwen3 gate restored:

equal-length batch vs single: cos_a=0.613198, cos_b=0.617658
equal-length batch[0] diverged from single inference (cos=0.61319774); causal mask skipped for no-padding batch?

Negative evidence with only the old Gemma3 no-bias behavior restored:

Gemma3 equal-length batch vs single: cos_a=0.924376, cos_b=0.900559
equal-length Gemma3 batch[0] diverged from single inference (cos=0.92437583); attention mask skipped for no-padding batch?

Expected behavior

The same input should produce the same embedding whether it is embedded alone or inside an equal-length backend batch.

Qwen3 should always apply its causal mask for non-bidirectional models, including no-padding batches.

Gemma3 should always create or receive a base attention bias so every attention layer adds its causal or sliding-window mask, including no-padding batches.

Impact

The router coalesces concurrent requests into backend batches. Any sub-batch that happens to contain equal-length inputs can silently produce wrong embeddings. This is especially plausible for retrieval/RAG workloads with short, similarly shaped queries.

This is independent of dtype or BLAS backend. The synthetic tests keep the regression coverage small and deterministic enough for CI.

Root cause

Qwen3 currently gates causal masking on attention_bias.is_some() in the batched path. Equal-length batches do not build padding bias, so causal masking is skipped.

Gemma3 currently lets attention_bias remain None for equal-length batches. Gemma3Attention::forward only calls create_attention_mask inside the Some(attention_bias) branch, so both full causal and sliding-window masks are skipped.

Expected fix

Always provide a zero base attention bias for equal-length batches that need masking but no padding:

• Qwen3: synthesize a zero bias before applying get_causal_attention_bias.
• Gemma3: synthesize a zero bias in Gemma3Model::forward so every attention layer adds its causal or sliding-window mask.

PR: #883

[HexCore8]

I discovered Hugging Face while exploring artificial intelligence and machine learning resources on GitHub and through various developer communities. As I learned more about its open-source models, tools, and educational resources, I became interested in its mission of making AI accessible and collaborative for everyone.

What I like most about the Hugging Face community is its strong focus on open-source innovation, collaboration, and knowledge sharing. The community brings together developers, researchers, and AI enthusiasts from around the world, creating an environment where people can learn from one another, contribute to meaningful projects, and build impactful AI solutions.

I want to join the Hugging Face community because I am passionate about artificial intelligence, machine learning, and continuous learning. I am eager to connect with like-minded individuals, explore cutting-edge AI technologies, contribute to open-source projects, and strengthen my technical skills. I believe being part of Hugging Face will help me grow both personally and professionally while enabling me to contribute to a community that is shaping the future of AI.

[malaiwah]

Built and pushed a ready-to-use CPU image from the latest TEI release (v1.9.3) with the Qwen3/Gemma3 attention-mask fix applied on top:

docker.io/malaiwah/text-embeddings-inference:cpu-1.9.3-qwen3-gemma3-causal-mask

Image details:

• Base source: v1.9.3
• Patch commit on top: Qwen3 + Gemma3 equal-length batch mask fix from fix(candle): always apply attention masks for Qwen3 and Gemma3 batches #883
• Platform built/tested: linux/amd64
• Pushed image config digest observed after registry pull: 8ebeb96d19c5731de5ebb74dcb296b89d6dfa18cbdcc70800c895be400d0d39b

Example run:

podman run --rm -p 8080:80 -v /path/to/hf-cache:/data \
  docker.io/malaiwah/text-embeddings-inference:cpu-1.9.3-qwen3-gemma3-causal-mask \
  --model-id Qwen/Qwen3-Embedding-0.6B \
  --max-batch-tokens 16384 \
  --hostname 0.0.0.0

Empirical CPU check on Qwen/Qwen3-Embedding-0.6B:

Official ghcr.io/huggingface/text-embeddings-inference:cpu-1.9, TEI 1.9.3
token_counts [11, 11, 3]
[A,A] equal-length batch vs single: cos0=0.142477, cos1=0.142477
[A,C] padded unequal-length batch vs single: cos0=1.000000

Patched docker.io/malaiwah/text-embeddings-inference:cpu-1.9.3-qwen3-gemma3-causal-mask, TEI 1.9.3
token_counts [11, 11, 3]
[A,A] equal-length batch vs single: cos0=1.000000, cos1=1.000000
[A,C] padded unequal-length batch vs single: cos0=1.000000

So the official release image reproduces the equal-length causal-mask failure, while the patched release image keeps equal-length batched embeddings identical to single-input embeddings.

[malaiwah]

Additional local validation: this is also proven on macOS/arm64 using the same synthetic-weight setup as the regression tests in #883.

Environment:

OS: macOS/arm64
rustc 1.92.0 (ded5c06cf 2025-12-08)
cargo 1.92.0 (344c4567c 2025-10-21)

Positive validation on the fixed branch:

cargo test -p text-embeddings-backend-candle --test test_qwen3_equal_length_batch --test test_gemma3_equal_length_batch

test test_gemma3_equal_length_batch_matches_single ... ok
test test_qwen3_equal_length_batch_matches_single ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out
test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out

Using temporary negative worktrees to restore the old behavior:

Qwen3 old mask gate restored:
equal-length batch vs single: cos_a=0.613198, cos_b=0.617658
equal-length batch[0] diverged from single inference (cos=0.61319774)

Gemma3 old no-bias behavior restored:
Gemma3 equal-length batch vs single: cos_a=0.924376, cos_b=0.900559
equal-length Gemma3 batch[0] diverged from single inference (cos=0.92437583)

So the issue is not Linux- or x86-specific; the synthetic tests reproduce the masking bug locally on macOS/arm64, and the fixed branch passes the same tests there.

[malaiwah]

Archaeology: why the earlier Qwen3 fixes missed this edge case

I traced the relevant backends/candle/src/models/qwen3.rs history to check whether this has prior art. The short version is: the previous Qwen3 fixes correctly covered padded batches and single-input inference, but the all-equal-length multi-sequence batch path still had no bias tensor, so the causal mask had nothing to attach to.

On current main before #883, the multi-sequence path only builds attention_bias when padding is required:

let attention_bias = if masking {
    // build padding bias tensor -> Some(...)
    Some(attention_bias)
} else {
    None
};

Later, causal masking is only applied if that bias tensor exists:

let attention_bias = if self.use_bidirectional_attention {
    attention_bias
} else if let Some(attn_bias) = attention_bias {
    Some(self.get_causal_attention_bias(attn_bias)?)
} else {
    None
};

So an equal-length batch (masking == false) falls through as None and runs without the causal mask.

Date	Commit/PR	What changed	Effect on this edge case
2025-06-10	11ffc60 / #627	Added the initial Qwen3Model wiring.	Not the full CPU batch path yet.
2025-06-12	bedb2e5 / #632	Added CPU support for Qwen3 embeddings, including the multi-sequence batch builder.	The attention_bias = if masking { Some(...) } else { None } structure appears here.
2025-06-18	1193add / #646	Added get_causal_attention_bias(...) and applied it only when attention_bias was already Some. Also fixed padding/pooling behavior.	Padded batches get the causal mask; equal-length batches still have attention_bias = None.
2025-06-24	f7aa35b / #648	Fixed single-vs-batch inconsistency by giving the single-sequence path a zero bias, using left padding, and adjusting pooling.	Single-input inference is fixed; multi-sequence equal-length batches still keep None.
2025-07-03	de70b7d / #663	Fixed Qwen3 FP16 dtype handling in batched inference.	Dtype-only; the no-padding None branch remains.
2026-02-12	b59f754 / #808	Added use_bidirectional_attention and projection support.	The causal gate became conditional on use_bidirectional_attention, but the else { None } path was carried forward.
2026-06-24	#883	Always provides a zero base attention bias for causal Qwen3 batches, including no-padding batches.	Equal-length batches now receive the causal mask.

This also explains why the existing Qwen3 test did not catch it. tests/test_qwen3.rs checks a mixed batch containing duplicate text and verifies the duplicate batch entries match single-input inference, which covers the padded/mixed batch behavior fixed by #646/#648. It does not construct the specific problematic case: two different sequences with the same length in the same backend batch.

Related history: #774 was a real Qwen3 CPU mismatch report against cpu-1.7.2, and the maintainer response pointed to the Qwen3 fixes in #646, #648, #650, and #663 that landed by v1.8.0. This issue is narrower: those fixes covered the common padded/single paths, while the equal-length no-padding sub-batch path was still outside the test surface.

The synthetic regression test in #883 closes that gap by using two different same-length sequences and checking batch-vs-single equality without downloading model weights.