Skip to content

Skip Metal-only gated delta kernel tests on other backends(fix CI) - #4522

Merged
nastya236 merged 2 commits into
ml-explore:mainfrom
aleroot:fix_cuda_build
Sep 16, 2026
Merged

nastya236 merged 2 commits into
ml-explore:mainfrom
aleroot:fix_cuda_build

Conversation

@aleroot

@aleroot aleroot commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

The test_gated_delta_{sequential,simdgroup,nax} tests set GATED_DELTA_CHUNK to pick a Metal kernel, but were guarded with mx.is_available(mx.gpu), so they ran on CUDA too. On CUDA GatedDeltaUpdate::use_fallback always returns true, so both sides of the comparison were the same unrolled fallback graph over the large gpu_dims cases. That leaves ~12 GiB cached in the allocator on the T4 runner; test_fft then hits a cuFFT allocation failure, which surfaces as cudaPeekAtLastError() failed: out of memory and aborts the test process.

Pre-fix device footprint of the three tests with the fallback forced on: 12.1 GiB. Post-fix (remaining tests): 0.64 GiB.

I have discovered even another weird thing during this investigation but I will probably raise another PR for it...

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure:

@nastya236 nastya236 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

thank you!

Comment thread python/tests/test_fast_gated_delta.py Outdated
self.assertTrue(hf.dtype == mx.float32, msg="State " + msg)

@unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available")
# The chunk size only selects between Metal kernels. Other backends always

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I think the comment should be that we have fast implementation only on metal for now, it is independent from GATED_DELTA_CHUNK

@nastya236
nastya236 merged commit c948334 into ml-explore:main Sep 16, 2026
29 checks passed
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.

2 participants