Fix Qwen3 Embedding Float16 DType

[tpendragon]

What does this PR do?

Without this, batches were erroring on CPU because it was trying to compare the F32 attention mask with the F16 output tensors.

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,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

[alvarobartt]

Hey @tpendragon thanks for the PR, but we just ran it from main on both CPU and MPS with --dtype float16 and couldn't reproduce the issue you mention; could you please let us know how to reproduce in case we missed anything? In any case, could you check from main instead in case the relevant fix is already there? Thanks in advance 🤗

[tpendragon]

@alvarobartt Ah, sorry, I'm running on a Mac with CPU only and running the rust as a crate (https://github.com/pulibrary/dpul-collections/blob/334456f531b41ce084dd8dd822d792199df298c8/native/embedanything/src/lib.rs#L44.) Maybe I missed a step in setting that up? I'd expect this branch of code to get hit if it was run with the candle backend on CPU with dtype 16.

[tpendragon]

@alvarobartt oh and importantly it only breaks if two inputs are batched together that have different token lengths, so padding happens. If the batches happen singly it doesn't trigger the bug.

[alvarobartt]

Fair @tpendragon thanks for the details on reproducing, I'll try on both CPU and MPS with the aforementioned settings and come back to you! I'm adding another comment on the PR with what IMO is not required 🤗 Thanks again for the PR!

[alvarobartt]

Thanks again for the PR! Didn't reproduce your issue yet, but the PR looks fair, only left some nits / comments 👍🏻

[alvarobartt]

Also here the linting is failing, you can install it as pip install pre-commit --upgrade, then pre-commit install from the root directory of the repository to run that on every commit, or just run it one to fix the issue with pre-commit run --all-files 🤗

[alvarobartt]

Hey @tpendragon I already tested and it's indeed fixed with your PR, thanks! Could you please run pre-commit run --all-files to fix the linting issues? Otherwise feel free to take the following patch and just run git apply on your branch and then commit that 🤗

diff --git a/backends/candle/src/models/qwen3.rs b/backends/candle/src/models/qwen3.rs
index 064ab6d..1330992 100644
--- a/backends/candle/src/models/qwen3.rs
+++ b/backends/candle/src/models/qwen3.rs
@@ -519,7 +519,8 @@ impl Qwen3Model {

             let attention_bias = if masking {
                 let attention_bias =
-                    Tensor::from_vec(attention_bias, (batch_size, 1, 1, max_length), &self.device)?.to_dtype(self.dtype)?;
+                    Tensor::from_vec(attention_bias, (batch_size, 1, 1, max_length), &self.device)?
+                        .to_dtype(self.dtype)?;
                 // Broadcast once instead of at every layer
                 let attention_bias = attention_bias
                     .broadcast_as((batch_size, self.num_attention_heads, max_length, max_length))?
diff --git a/backends/src/lib.rs b/backends/src/lib.rs
index d333951..be40b09 100644
--- a/backends/src/lib.rs
+++ b/backends/src/lib.rs
@@ -150,11 +150,8 @@ impl Backend {
         }

         max_input_length = std::cmp::min(max_input_length, max_warmup_length);
-        let mut seq_lengths: Vec<usize> = generate_bucket_sizes(
-            seq_bucket_size,
-            max_input_length,
-            seq_len_exp_base,
-        );
+        let mut seq_lengths: Vec<usize> =
+            generate_bucket_sizes(seq_bucket_size, max_input_length, seq_len_exp_base);
         if let Some(&last) = seq_lengths.last() {
             if last < max_input_length {
                 seq_lengths.push(max_input_length);

[tpendragon]

@alvarobartt Ran it - thanks!

[alvarobartt]

@alvarobartt Ran it - thanks!

Hey @tpendragon fmt seems to still be failing, could you please check that? Thanks in advance 🤗

[tpendragon]

@alvarobartt I have no idea what's going on, I can't get fmt to throw an error, even in a nix shell.