Skip to content

fix(candle): always apply attention masks for Qwen3 and Gemma3 batches - #883

Open
malaiwah wants to merge 1 commit into
huggingface:mainfrom
malaiwah:codex/fix-qwen3-gemma3-causal-mask
Open

fix(candle): always apply attention masks for Qwen3 and Gemma3 batches#883
malaiwah wants to merge 1 commit into
huggingface:mainfrom
malaiwah:codex/fix-qwen3-gemma3-causal-mask

Conversation

@malaiwah

@malaiwah malaiwah commented Jun 24, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes #882.

Fixes equal-length batched inference for Qwen3 and Gemma3 on the candle backend.

  • Qwen3 now always applies the causal mask for non-bidirectional models, synthesizing a zero base bias when the batch needs no padding.
  • Gemma3 now passes a zero base bias when the batch needs no padding, so each attention layer still adds its causal or sliding-window mask.
  • Adds no-download synthetic-weight regression tests for both models.

Why

The previous code only created an attention bias for multi-sequence batches when padding was required. Equal-length batches need no padding, so attention_bias stayed None.

For Qwen3, that meant the causal mask was skipped.

For Gemma3, the per-layer causal/sliding-window mask is created only when an attention bias tensor is present, so equal-length batches skipped those masks as well.

Single-sequence inference already creates a zero bias and therefore remains correctly masked. As a result, the same text could embed differently alone vs in an equal-length backend batch.

No new runtime dependencies are required. The only manifest change adds candle as a dev-dependency for synthetic safetensors generation in the new tests.

Validation

Using the repository-pinned Rust toolchain from rust-toolchain.toml:

rustc 1.92.0 (ded5c06cf 2025-12-08)
cargo 1.92.0 (344c4567c 2025-10-21)

Positive synthetic-weight tests on the fixed branch, CPU, no model download:

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

Result:

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

Negative evidence, restoring only the old Qwen3 mask gate:

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, restoring only the old Gemma3 no-bias behavior:

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?

Formatting:

cargo fmt -p text-embeddings-backend-candle -- --check

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline?
  • Was this discussed/approved via a GitHub issue or the forum? See Qwen3/Gemma3 candle skip attention masks for equal-length batches #882.
  • Did you make sure to update the documentation with your changes? Not applicable: this is a backend correctness fix with no user-facing API or documentation change.
  • Did you write any new necessary tests? If applicable, did you include or update the insta snapshots? Added synthetic-weight Qwen3 and Gemma3 regression tests; no insta snapshots are needed.

Notes

The original Qwen3 issue was also reproduced with real Qwen/Qwen3-Embedding-0.6B on CPU in fp32 and fp16, and without MKL. These tests keep the upstream regression coverage small and deterministic enough for CI by using tiny random synthetic weights.

@malaiwah
malaiwah marked this pull request as ready for review June 24, 2026 17:51
@malaiwah
malaiwah marked this pull request as draft June 24, 2026 17:55
@malaiwah
malaiwah marked this pull request as ready for review June 24, 2026 18:01
@malaiwah

malaiwah commented Jun 26, 2026

Copy link
Copy Markdown
Author

CUDA validation added on an RTX 5090 host.

Environment:

Host: Ubuntu 24.04, NVIDIA GeForce RTX 5090
Driver: 595.58.03; nvidia-smi reports CUDA 13.2
Container: podman + docker.io/nvidia/cuda:12.8.0-devel-ubuntu24.04
CUDA_COMPUTE_CAP=120

I used a validation-only RUSTFLAGS='--cfg feature="dynamic-linking"' in the container so cudarc links dynamically against the CUDA runtime available there. No repository files were changed for this.

Synthetic no-download regression tests, which are the CI-friendly coverage for the exact equal-length/no-padding mask bug:

cargo test -p text-embeddings-backend-candle \
  --features cuda \
  --test test_qwen3_equal_length_batch \
  --test test_gemma3_equal_length_batch \
  -- --nocapture

Result:

Gemma3 equal-length batch vs single: cos_a=1.000000, cos_b=1.000000
test test_gemma3_equal_length_batch_matches_single ... ok

equal-length batch vs single: cos_a=0.999999, cos_b=1.000000
test test_qwen3_equal_length_batch_matches_single ... ok

I also ran real-weight CUDA checks through the router, comparing a single-item request against an equal-length batch made from the same text repeated twice.

Qwen/Qwen3-Embedding-0.6B:

dim=1024 batch=2
cos_single_vs_batch0=0.999999819
cos_single_vs_batch1=0.999999819
max_abs_diff_batch0=6.6448e-05
max_abs_diff_batch1=6.6448e-05

google/embeddinggemma-300m using HF gated-model access:

dim=768 batch=2
cos_single_vs_batch0=0.999999855
cos_single_vs_batch1=0.999999855
max_abs_diff_batch0=5.952243e-05
max_abs_diff_batch1=5.952243e-05

The synthetic tests remain the upstream regression tests because they are deterministic and require no model download or gated access. The real Qwen3 and Gemma3 checks confirm the same fixed behavior with actual model weights on CUDA.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant