Skip to content

Allow rank 3 cp.async - #40222

Open
samanklesaria wants to merge 1 commit into
jax-ml:mainfrom
samanklesaria:push-vmyxvqnmzvnx
Open

Allow rank 3 cp.async#40222
samanklesaria wants to merge 1 commit into
jax-ml:mainfrom
samanklesaria:push-vmyxvqnmzvnx

Conversation

@samanklesaria

Copy link
Copy Markdown
Contributor

Closes #40201

Currently, cp.async refuses any GMEM ref that is not exactly 2-D. This PR relaxes that restriction. To test it,tests/mosaic/gpu_test.py::AsyncCopyTest::test_cp_async_tiled already covers this path for rank 2, and _test_cp_async is rank-agnostic, so we just add a shape parameterization of (2, 64, 128) to the test.

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.

Support GMEM ranks above two in the tiled cp.async GMEM→SMEM copy (pre-Hopper)

1 participant