[ET-VK][sdpa] Expose shader override by SS-JIA · Pull Request #24343 · pytorch/executorch

SS-JIA · 2026-10-30T15:02:45Z

Stack from ghstack (oldest at bottom):

Ulterior Motive

Benchmark SDPA shader routes without duplicating the production operator.

Rationale

Expose sdpa_impl_with_shader_override() to assemble the full SDPA operator while selecting its attention-value shader. The registered operator delegates to this function with no override, preserving automatic routing and the existing QK/softmax/AV dispatch chain.

sdpa_impl_with_shader_override(
    graph, args, graph.add_scalar(kShaderOverrideForceTile2));

The override is a graph scalar using the existing SDPA route constants. Validation is provided by SDPA JSON cases and route sweeps in the later benchmark_v2 commit.

Authored with Codex (AI).

Differential Revision: D128509576

[ghstack-poisoned]

pytorch-bot · 2026-10-30T15:02:49Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit b653bda with merge base aed328d ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-10-30T15:03:29Z

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.

JakeStevens

Review automatically exported from Phabricator review in Meta.

[ghstack-poisoned]

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24344 * #24343 * #24342 * #24341 * #24340 * __->__ #24339 ## Ulterior Motive Validate production quantized operators directly in the JSON benchmark. ## Rationale The dispatcher already passes `alpha`, but the shader ignored it. Add it to the push constants and scale the second input before requantization. | Before | After | | --- | --- | | Computes `a + b` for every `alpha`. | Computes `a + alpha * b`. | Correctness coverage is provided by the quantized-add JSON cases in `benchmark_v2`, introduced later in this stack. Authored with Codex (AI). Differential Revision: D128509571 Pull Request resolved: #24339

@manuelcandales

…features Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24344 * #24343 * #24342 * #24341 * __->__ #24340 * #24339 ## Ulterior Motive Run supported shaders with the required Vulkan features enabled, and skip unsupported shaders before driver compilation. ## Rationale Make compiled SPIR-V the source of shader capability requirements and carry those requirements through code generation, device setup, and pipeline validation. GLSL extension declarations can be hidden in includes or unused. Vulkan features must also be enabled on the logical device before shaders can use them. ## Details | Area | Before | After | | --- | --- | --- | | Requirement detection | Six flags inferred from GLSL extension text. | Require compiled SPIR-V, preserve every declared `OpCapability`, and remove GLSL extension matching. | | Runtime validation | Check the six feature flags. | Map capability IDs to arithmetic, storage, image, subgroup, and cooperative-matrix support; reject unsupported or unrecognized IDs before shader-module/pipeline creation. | | Device setup | Cache selected core features and enable `shaderInt16`, `shaderInt64`, and `shaderFloat64`. | Cache the full core feature query and also enable supported shader image/indexing features required by the expanded checks. | | Error reporting | Describe requirements using the fixed `VulkanExtension` enum. | Accept requirement descriptions that include arbitrary SPIR-V capability IDs. | `gen_vulkan_spv.py` emits the capability list into `ShaderInfo.required_capabilities`; named `SpirvCapability` IDs make the runtime mapping readable. `Device` caches feature queries, `Adapter` configures the device and checks support, and `Context`/`ComputeGraph` enforce requirements. Existing flag checks remain for manually constructed metadata without a list. ```python info = getShaderInfo(glsl_path, spv_path) ``` ```cpp context.check_device_capabilities(shader); ``` Use the authoritative Python host tool in shader-generation builds. Keep descriptor, tile, and required-subgroup-size metadata parsing. Generator tests cover mappings, ignored GLSL extension declarations, unknown-capability preservation, and malformed modules; framework-dependent integration tests accompany the later benchmark commit. Authored with Codex (AI). cc @manuelcandales @digantdesai @cbilgin Differential Revision: D128509572 Pull Request resolved: #24340

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24344 * #24343 * #24342 * __->__ #24341 * #24340 * #24339 ## Ulterior Motive Allow the JSON benchmark to sweep workgroup dimensions such as `[10,10,1]`. ## Rationale Store explicit dimensions directly rather than as power-of-two exponents. Keep automatic shape fitting unchanged and reject invocation-count overflow before checking device limits. | Before | After | | --- | --- | | Explicit dimensions must be powers of two; products can overflow. | Explicit dimensions can be arbitrary positive integers; validation checks the product safely. | ```cpp LocalWorkGroup lwg(10, 10, 1); lwg.validate(adapter->max_compute_workgroup_size(), adapter->max_compute_workgroup_invocations()); ``` Keep the four dispatch-grid regression tests and their Buck target in `test/api_abstractions/`, independent of the benchmark framework. Authored with Codex (AI). Differential Revision: D128509574 Pull Request resolved: #24341

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24344 * #24343 * __->__ #24342 * #24341 * #24340 * #24339 ## Ulterior Motive Support separate batch timing and per-shader reports in the JSON benchmark. ## Rationale Expose unscaled 64-bit query ticks with shader/workgroup metadata, allowing callers to apply the device timestamp period without the legacy integer conversion. Add a context toggle to omit per-shader timestamp commands, and expose dispatch argument groups for input and throughput reporting. ```cpp graph.context()->set_shader_profiling_enabled(false); // Encode batch timestamp queries separately. auto ticks = graph.context()->querypool().get_shader_timestamp_ticks(); const auto& args = graph.execute_nodes().front()->args(); ``` Profiling remains enabled by default. Existing timestamp-report APIs are retained. Keep independent tests for empty pools, enabled/disabled instrumentation, GPU output, and argument/timestamp metadata under `test/api_abstractions/`. Authored with Codex (AI). Differential Revision: D128509575 Pull Request resolved: #24342

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * __->__ #24344 * #24343 * #24342 * #24341 * #24340 * #24339 ## Ulterior Motive Replace C++ benchmark case definitions with JSON specifications. ## Rationale Introduce `etvk_ops_bench` with inline tensor arguments, deterministic initialization, dtype/storage and shader/workgroup sweeps, operator routes, independent CPU references, and registered throughput estimates. Generate host tensor data once per case and share it between the staged validation graph and a separate benchmark graph using automatic calibration. Select validation scope with `--validation_mode=all|first|none`, defaulting to all cases. Default to compact `json` reports; provide condensed tables and detailed `json_full` reports with per-batch GPU/operator samples and shader statistics. Record Android GPU clock start/end snapshots per measurement pass, adjacent to the corresponding batch timing arrays. Tables show observed minimum/maximum GPU clocks across each case and profiling-only ranges for shader rows. Document measurement boundaries, clock units, and shader percentile sampling. Migrate 64 sets into 4,560 definitions using production operators directly, without a dependency on the old benchmark directory. Keep quantized references, route checks, and scripts together. Case JSONs and reports omit `legacy_case`. Review framework parsing/execution/reporting first, then operator references/routes and JSON cases, then usage documentation. This remains one commit above the five runtime commits. Runtime-specific generator tests accompany their implementation; workgroup tests and their Buck target live in `test/api_abstractions/`. Benchmark validation also covers rejection of unrecognized SPIR-V capabilities before pipeline creation. Non-constant tensor inputs display their configured memory layout after the dtype; constant tensors omit this suffix because operator prepacking selects their GPU layout. `json` uses untruncated input strings, min/median/max batch timing summaries, and per-pass clock extrema. Float arguments and list elements in input strings use three significant figures. `json_full` preserves tensor descriptions and individual batch samples. Both retain shader timing statistics. The public README describes the binary, layout, JSON schema, validation, and timing. Internal helper scripts and Buck build instructions are documented under `fb/README.md`; migration documents and the TODO file are omitted. CI workflow integration will follow separately. Authored with Codex (AI). Differential Revision: D128509578 Pull Request resolved: #24344

Pull Request resolved: #24343 ## Ulterior Motive Benchmark SDPA shader routes without duplicating the production operator. ## Rationale Expose `sdpa_impl_with_shader_override()` to assemble the full SDPA operator while selecting its attention-value shader. The registered operator delegates to this function with no override, preserving automatic routing and the existing QK/softmax/AV dispatch chain. ```cpp sdpa_impl_with_shader_override( graph, args, graph.add_scalar(kShaderOverrideForceTile2)); ``` The override is a graph scalar using the existing SDPA route constants. Validation is provided by SDPA JSON cases and route sweeps in the later `benchmark_v2` commit. Authored with Codex (AI). ghstack-source-id: 461636572 @exported-using-ghexport Differential Revision: [D128509576](https://our.internmc.facebook.com/intern/diff/D128509576/)

Update

c6477f0

[ghstack-poisoned]

SS-JIA temporarily deployed to cadence October 30, 2026 15:03 — with GitHub Actions Inactive

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 30, 2026

meta-codesync Bot added the meta-exported label Oct 30, 2026

SS-JIA temporarily deployed to cadence October 30, 2026 15:17 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence October 30, 2026 15:19 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence October 30, 2026 15:21 — with GitHub Actions Inactive

shoumikhin approved these changes Oct 30, 2026

View reviewed changes

JakeStevens approved these changes Oct 30, 2026

View reviewed changes

Update

26b2fc6

[ghstack-poisoned]

SS-JIA had a problem deploying to cadence October 30, 2026 18:05 — with GitHub Actions Failure

SS-JIA temporarily deployed to cadence October 30, 2026 18:05 — with GitHub Actions Inactive

Update

b653bda

[ghstack-poisoned]

SS-JIA temporarily deployed to cadence October 31, 2026 07:31 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence October 31, 2026 07:47 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence October 31, 2026 07:50 — with GitHub Actions Inactive

SS-JIA temporarily deployed to cadence October 31, 2026 07:51 — with GitHub Actions Inactive

meta-codesync Bot merged commit 9b7036d into gh/SS-JIA/677/base Oct 31, 2026
237 of 238 checks passed

meta-codesync Bot deleted the gh/SS-JIA/677/head branch October 31, 2026 16:42

meta-codesync Bot deployed to cherry-pick-bot October 31, 2026 16:42 Active

pytorchbot mentioned this pull request Oct 31, 2026

[ET-VK][sdpa] Expose shader override #24390

Merged

Navigation Menu

Repository navigation

[ET-VK][sdpa] Expose shader override - #24343

[ET-VK][sdpa] Expose shader override#24343
meta-codesync[bot] merged 3 commits into
gh/SS-JIA/677/basefrom
gh/SS-JIA/677/head

SS-JIA commented Oct 30, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Oct 30, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 30, 2026

Uh oh!

JakeStevens left a comment

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

SS-JIA commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Ulterior Motive

Rationale

Uh oh!

pytorch-bot Bot commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

github-actions Bot commented Oct 30, 2026

This PR needs a release notes: label

Uh oh!

JakeStevens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

SS-JIA commented Oct 30, 2026 •

edited

Loading

pytorch-bot Bot commented Oct 30, 2026 •

edited

Loading

This PR needs a `release notes:` label