[Shakhtar-Sankur]

Summary

The reference ops in exir/passes/_quant_patterns_and_replacements.py require weight_zero_points to share weight_scales' dtype, and the variants without a dtype argument return the scales' dtype. The kernels checked the zero points against the output dtype instead (while reading them as the scales' type), and the .out variants read the scales as the output type without checking it. So:

• embedding_byte/2bit/4bit.dtype_out with fp16 scales and zero points and a fp32 output (valid per the reference) aborts;
• fp32 zero points with fp16 scales are accepted and reinterpreted as fp16;
• embedding_byte/2bit/4bit.out with fp16 scales and a fp32 output reads past the end of the scales (AddressSanitizer: heap-buffer-overflow).

This checks the zero points against the scales, and the scales against the output in the .out variants.

Related to #23677 (index bounds check in the same kernels); the two changes are independent and merge cleanly in either order.

Test plan

Built quantized_kernels_test with CMake (Debug, -fsanitize=address).

• With the change: 86/86 quantized kernel tests pass (88/88 together with Add index bounds check to 2-bit and 4-bit quantized embedding kernels #23677).
• Without it, all 6 new tests fail: the two fp16-qparams/fp32-output tests abort, the two zero-point mismatch tests fail to die, and the two .out mismatch tests die in AddressSanitizer (heap-buffer-overflow) instead of the dtype check.
• Python reference cross-check: embedding_byte.dtype with fp16 scales and zero points and dtype=float32 returns [[3.5, 2], [5.5, 13]], the value the new test expects; fp32 zero points with fp16 scales raise an assertion.

Authored with AI assistance. I reviewed the change and ran the tests myself.

cc @larryliu0820 @manuelcandales @JakeStevens @nil-is-all

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23912

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 460b8b9 with merge base d1f7f1f ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Test WebGPU Backend / test-webgpu / test-backend-linux (webgpu, operators) / linux-job (gh) (trunk failure)
  curl: (22) The requested URL returned error:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Shakhtar-Sankur]

@pytorchbot label "release notes: ops & kernels"

[nil-is-all]

@Shakhtar-Sankur could you rebase and ping for review? Thanks

[Shakhtar-Sankur]

@nil-is-all rebased onto main (it now sits on top of #23677), thanks. @JakeStevens @metascroy could you take a look when you have a moment?