Repository navigation
[ET-VK][runtime] Align buffer packed dimensions to four - #24774
Merged
Merged
Conversation
Pull Request resolved: #24753 ## Ulterior Motive Make buffer storage behave more like texture storage, so one shader can service both storage types. Textures already pad the packed dimension to four elements; buffers now do the same. This also allows a single vectorized shader for buffer tensors instead of separate scalar and vector variants. ## Rationale **What**: Align the packed dimension of ordinary buffer tensors to four scalar elements. Update dispatch, indexing and tensor contracts for the physical padding. **Why**: Vector tail accesses on odd-sized rows could cross row or allocation boundaries. Aligned physical rows make those accesses safe without changing logical tensor shapes. ## Details - Buffer representation - **Before**: `packed_dim_align` was `1`, so logical and physical row lengths matched. - **After**: `packed_dim_align` is `4`; padded sizes, strides, metadata and allocations include row padding. `kInt8x4` is unchanged. - Implicit packed dim - **Before**: unpadded buffers had no storage for a packed dim beyond the tensor's rank. - **After**: as in textures, that dim is padded to four and innermost, so each element is followed by three padding lanes. `BufferMetadata` and the hashed layout list it first in the dim order with stride 1, so index decoding treats its padding as out of bounds. Such tensors are not contiguous. - Dispatch and indexing - **Before**: buffer kernels dispatched logical numel, and several shaders used logical widths as physical strides. - **After**: ordinary buffers dispatch physical numel. Layout-sensitive shaders index through `BufferMetadata`. - GEMM selection - **Before**: odd dimensions selected scalar linear and matmul shaders. - **After**: every ordinary buffer shape selects the vector shaders. Padding lanes may hold non-finite values, so the last partial K vector is masked after the main K loop; full vectors load unmasked. - SDPA head_dim tail - **Before**: LLM attention weights read the padding lanes of the last partial head_dim vector unmasked. Fused SDPA zeroed them on every loop iteration. - **After**: both paths mask the Q and K padding lanes of the last partial head_dim vector. The tiled kernel masks once after the main loop. The coop kernel masks inside the loop, since each worker covers about one vector. - Native backend views - **Before**: static view, squeeze, unsqueeze, permute, transpose and select of a buffer aliased their source. - **After**: they become copies when the source or view can carry padding, because a view pads its own sizes along its source's packed dim. Transpose runs as `permute_copy`, since ET-VK has no `transpose_copy`. `alias` keeps aliasing. Views in mutated alias groups must stay aliases, so such a buffer view that can carry padding is rejected at build time. - Tensor contracts - **Before**: Q4 embedding indices, fused CE logits and per-row reduction inputs had no layout checks. AdamW checked only that numels match. - **After**: each requires a contiguous width-packed buffer, matching the padded row stride its shader assumes. AdamW requires `m`, `v` and `grad` to match `param`'s sizes and hashed layout, since its shader indexes all four with one padded offset. Smaller changes: - `has_same_physical_element_order` keeps direct buffer copies when padding and element order match. - Unary shaders bound-check against output metadata instead of push constants, so buffers cover padded numel. - Unused scalar linear/matmul templates are removed. - Interleaved RoPE uses vector frequency loads. - Legacy 4W4C QDQ kernels use physical strides. - Direct-copy prepack sends padded buffers through standard prepack. - Q4 embedding reads indices by padded row, and reads 1-D per-row scales with stride 1. - Layer norm and RMS norm prepack buffer weights width-packed, since their shaders index them densely. - Buffer arange writes by logical index and zeroes padding. Zero-padding enforcement is deferred to follow-up commits. Review `Tensor.cpp` and `ComputeGraph.cpp` first, staging and layout-conversion shaders second, operator kernels third, then GEMM, the native backend and tests. Authored with Codex and Claude Code. ghstack-source-id: 467556304 @exported-using-ghexport Differential Revision: [D124585167](https://our.internmc.facebook.com/intern/diff/D124585167/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24774
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 2 Unrelated FailuresAs of commit 41cffc8 with merge base 9517b6b ( NEW FAILURE - The following job has failed:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
SS-JIA
approved these changes
Nov 7, 2026
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #24753 by @SS-JIA
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/SS-JIA/693/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/SS-JIA/693/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/SS-JIA/692/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/SS-JIA/693/orig
Differential Revision: D124585167
@diff-train-skip-merge