[digantdesai]

Stack from ghstack (oldest at bottom):

• [executorch][native-vulkan] Lower rms_norm without a weight on the Vulkan engine #25097
• [executorch][vulkan] Make the et_vk.rms_norm weight optional #25096
• -> [executorch][native-vulkan] Run bf16 methods in f32 on the Vulkan engine #25095
• [executorch][native] Record rms_norm op samples #25094

ET-VK has no bf16 storage, so the Native Vulkan engine rejected every bf16 method when it compiled. It now runs them in f32 on the device, the way it already runs int64 as int32. device_dtype maps BFloat16 to Float, and bf16 constants are widened to f32 host copies at load, next to the int64 narrowing. The method still reports bf16 from input_dtype and output_dtype, so callers pass and receive bf16 as before.

A bf16 tensor holds bf16 values at the host boundary whatever the host dtype: set_input rounds BFloat16 or Float host data to bf16, and get_output rounds the f32 result to the nearest bf16 before writing BFloat16 or Float. Other host dtypes are rejected. Between ops, values stay f32, so a method with several ops skips the bf16 rounding eager applies after each one, and its results match eager within bf16 tolerance rather than exactly. For example, (1 + 2^-8) - 1 returns 2^-8, where eager returns 0.

Every widened constant stays on the heap until prepack releases it, so loading a bf16 model needs about twice its bf16 weight size in extra host memory.

Memory planning now sizes each request by its device dtype, which doubles bf16 requests and halves int64 ones. ET-VK sizes each shared object from the tensors bound to it, so this only keeps the planner's ordering consistent with the device sizes. buffer_attention_fits also counts SDPA attention weights at the device element size, so a bf16 attention no longer picks buffer storage up to twice the storage buffer limit.

LowerRMSNorm also accepts bf16 norms, which now run the f32 et_vk.rms_norm kernel. Quantized linears and embeddings with bf16 scales still fail, because FuseQuantizedLinear and FuseQuantizedEmbedding match only f32 and f16 scales.

Authored with Claude Code.

Differential Revision: D131119348

NOTE FOR REVIEWERS: This PR has internal Meta-specific changes or comments, please review them on Phabricator!

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 1e7adcb with merge base 925c1bd ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• pull / unittest / macos / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.