Skip to content

[Bug] Fix reverse-mode ndarray gradients on Vulkan by binding the grad buffer directly.#801

Closed
duburcqa wants to merge 2 commits into
duburcqa/broaden_ndarray_ad_test_backendsfrom
duburcqa/fix_vulkan_ndarray_grad
Closed

[Bug] Fix reverse-mode ndarray gradients on Vulkan by binding the grad buffer directly.#801
duburcqa wants to merge 2 commits into
duburcqa/broaden_ndarray_ad_test_backendsfrom
duburcqa/fix_vulkan_ndarray_grad

Conversation

@duburcqa

Copy link
Copy Markdown
Contributor

Issue: #

Stacked on #800 -> #799 -> #796 -> #795. Review/merge those first; this PR is based on #800 (which adds Metal to the ndarray-autodiff test matrix) and additionally enables Vulkan once the gradient-buffer bug below is fixed. Base retargets down the stack as each lands.

Brief Summary

copilot:summary

Walkthrough

copilot:walkthrough

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 160f30c446

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread tests/python/test_ad_ndarray.py Outdated
…d buffer directly.

Reverse-mode autodiff over an `ndarray` argument produced all-zero gradients on the
Vulkan/MoltenVK backend: the forward pass and the gradient value were computed correctly,
but the accumulate never reached the host `.grad` buffer. Root cause: on backends with
`spirv_has_physical_storage_buffer` the ndarray data and gradient pointers are passed as
buffer_device_address `u64`s embedded in the args struct, and MoltenVK misreads the SECOND
such pointer member (the gradient slot) - it resolves to a garbage address, so the reverse
accumulate no-ops. The data slot (first pointer) reads fine, and quadrants' own Metal
backend reads the identical SPIR-V correctly; only MoltenVK's SPIRV-Cross mishandles it.
This is independent of the member's offset or alignment (verified: a 16-aligned grad slot
still fails) and of the buffer type (uniform vs storage).

Route the gradient ndarray through the bound-storage-buffer descriptor path instead of the
physical-address path - the same path already used on backends without
physical_storage_buffer, and the reason field gradients already work on Vulkan. The data
pointer keeps the physical-address path. This sidesteps the MoltenVK struct-member defect
with no change to the args layout or the data path, and no regression on Metal/CUDA/CPU.

With this fix every backend supports ndarray reverse-mode autodiff, so the ndarray autodiff
tests no longer need the `archs_support_ndarray_ad` arch restriction: drop the list and let
the tests run on the full expected-arch matrix (the tests fail on Vulkan without this fix).
@duburcqa
duburcqa force-pushed the duburcqa/fix_vulkan_ndarray_grad branch from 160f30c to e9ba11a Compare July 20, 2026 19:21
@github-actions

Copy link
Copy Markdown

@github-actions

Copy link
Copy Markdown

@duburcqa

Copy link
Copy Markdown
Contributor Author

Closing: the bound-buffer route this PR takes is capped by Metal's 31 buffer binding slots and cannot compile a real reverse-mode kernel (Genesis's rigid backward already overflows it). Root cause and both rejected attempts documented in #802. A real fix requires MoltenVK to read the gradient device address correctly on the physical-storage-buffer path.

@duburcqa duburcqa closed this Jul 20, 2026
@github-actions

Copy link
Copy Markdown

@github-actions

Copy link
Copy Markdown

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant