[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.

⚠️ 17 Awaiting Approval

As of commit e297841 with merge base 8ecf0fb ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test Native Engines (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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?

[JakeStevens]

The fake implementations for the non-dtype embedding_2bit and embedding_4bit return torch.nn.Embedding(...)(indices). That is always fp32, whatever the scales' dtype. instead, return weight_scales.dtype

[Shakhtar-Sankur]

Thanks @JakeStevens, good catch. The fake kernels for embedding_2bit and embedding_4bit now return weight_scales.dtype, matching the reference ops, and test_embedding_xbit_fake_dtype checks each fake against its reference op for fp16, bf16 and fp32 scales (the fp16 and bf16 cases fail without the fix).

@nil-is-all I've rebased on main. Ready for review, thanks!