[executorch][vulkan] Add payload-bounded constant sharding

[JCNTH]

Stack from ghstack (oldest at bottom):

• [executorch][webgpu] Test linear embedding dynamic reuse #21406
• [executorch][webgpu] Preserve embedding packing on resize #21405
• -> [executorch][vulkan] Add payload-bounded constant sharding #21404

Enable deterministic payload-bounded Vulkan constant sharding

Problem
Large Vulkan exports can make inline PTE named data exceed downstream artifact limits, but NamedDataStore cannot externalize that data into bounded groups.

Solution

• Before: Vulkan constants remain inline with no payload-bounded externalization option.
• After: an opt-in uint64 raw-data cap externalizes constants after graph construction into deterministic shards.

Implementation

• NamedDataStore.externalize_pte_data groups aliases by backing buffer, preserves existing tags, and rejects oversize buffers or collisions atomically.
• parse_compile_options and parse_compile_spec transport external_constants_max_data_bytes as a strict positive uint64.
• VulkanBackend.preprocess applies sharding only when the option is present.

Constraints
The default path is unchanged; buffers are not copied or split. The cap bounds unique raw buffer bytes per external tag, not final FlatTensor file size including alignment and header overhead.

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D113608555

Differential Revision: D113608555

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 6 Unrelated Failures

As of commit d0e8ae3 with merge base 2439920 ():

NEW FAILURE - The following job has failed:

• pull / unittest / windows / windows-job (gh)
  examples/models/test/test_export.py::ExportTest::test_ic3_export_to_executorch

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / test-llama-runner-linux (fp32, xnnpack+custom+qe, linux.2xlarge, executorch-ubuntu-22.04-clang12) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-models-linux (emformer_transcribe, portable, linux.2xlarge) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / test-qnn-passes-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / test-qnn-python-imports-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, operators) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-voxtral-realtime-xnnpack-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' 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.

[JCNTH]

@claude Thoroughly review for redundancy or areas of modularization or any high value changes

[claude]

Claude finished @JCNTH's task in 1m 40s —— View job

---

Review: payload-bounded Vulkan constant sharding

Thorough pass focused on redundancy, modularization, and high-value changes. Overall this is clean, well-tested (atomicity, idempotency, aliased-buffer-counted-once, order-independence are all covered), and the default path is genuinely untouched. Notes below, roughly in priority order.

1. externalize_pte_data externalizes all pte_data, but Vulkan labels it vulkan_constants

backends/vulkan/vulkan_preprocess.py:264-268 calls the method right after build_graph() and tags everything vulkan_constants. The method itself (_named_data_store.py:227-267) moves the entire pte_data map, not just constants. This is correct only if pte_data contains nothing but constants at that point. If any non-constant named data can land in pte_data before this call, it would be silently mislabeled as a constant shard. Worth a one-line assertion or comment documenting that invariant, since the method name (pte_data) and the Vulkan tag (constants) disagree.

2. The positive-uint64 check is validated in three places

• parse_compile_options — vulkan_partitioner.py:291-299 (rejects bool / non-int / <=0 / >= 1<<64)
• _parse_external_constants_max_data_bytes — vulkan_preprocess.py:93-99 (decode: length-8 + <=0)
• externalize_pte_data — _named_data_store.py:218-223 (bool / non-int / <=0)

Each with a slightly different message ("positive uint64" / "encoded as uint64" / "positive integer"), and each asserted by a separate test. Defense-in-depth across the encode → CompileSpec-bytes → decode → store boundaries is defensible, but the encode and decode validators live in the two Vulkan files and could share one helper. At minimum, the duplication is worth a comment explaining it's intentional boundary validation rather than an oversight. Fix this →

3. Greedy packing is key-sorted, not size-sorted

_named_data_store.py:234-251 does first-fit in sorted-key order. This is deterministic (good, and the reason for the sort) but not size-optimal — it can produce more shards than a first-fit-decreasing pass would. Given the PR explicitly documents "buffers are not copied or split" and prioritizes determinism, this is a reasonable trade-off, but a brief comment on why key-order over size-order (reproducibility) would help the next reader who is tempted to "optimize" it.

4. Test coupling to internal structure

test_vulkan_compile_options.py:_preprocess_named_data patches four internal symbols by name (unsafe_remove_auto_functionalized_pass, apply_passes, VkGraphBuilder, serialize_vulkan_graph). This verifies the wiring well, but it's tightly coupled to the current internals of preprocess and will break on unrelated refactors. Acceptable given a real graph would be heavyweight, but flagging as a maintenance cost.

Minor / positive

• Good: placing externalize_pte_data on the generic NamedDataStore rather than in Vulkan makes it reusable by other backends — nice modularization.
• Good: the external_data shallow-copy at _named_data_store.py:253-255 before mutation gives the atomic-on-collision guarantee the tests assert. Note the oversize check runs earlier (line 239) so that path is atomic too.
• _parse_external_constants_max_data_bytes is a single-use helper, which CLAUDE.md discourages, but here it cleanly separates two distinct error messages (encoding vs. value) — I'd keep it.
• Content-addressed tag naming (sha256 of the joined keys) is a solid choice for reproducibility and the collision guard is correct.

No blocking correctness bugs found. The main thing I'd want confirmed before merge is #1 — the pte_data-is-only-constants invariant at the call site.
· branch gh/JCNTH/194/head

[SS-JIA]

Review automatically exported from Phabricator review in Meta.