Skip to content

Make sure CCCL TMA is compiled and use it (very narrowly) - #2428

Open
seberg wants to merge 3 commits into
NVIDIA:mainfrom
seberg:tensor-map-linking
Open

Make sure CCCL TMA is compiled and use it (very narrowly)#2428
seberg wants to merge 3 commits into
NVIDIA:mainfrom
seberg:tensor-map-linking

Conversation

@seberg

@seberg seberg commented Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

The core issues here is that:

  • Previously, these CCCL headers were usually just not found by the host compiler, TMA was never used and never tested in CI. In my local setup (conda) and apparently occasional others they are found. This now inserts the includes so that they are generally founds.
  • CCCL had some issues, mainly checking checking the strong stride, which means we can't rely on it actually working.
  • After linking correctly, we run into the issue that CCCL needs to be soft-linked because it requires the CUDA runtime to be loaded. I.e. if TMA is successfully compiled all that happened previously was that CUDA runtime symbols were not linked correctly.

This should be ready for review, two notes though:

  • The CCCL path is now used extremely narrowly due to issues there, this should be fixed with CCCL 3.5 (fixed on main).
  • I don't like try/except style, so I removed that. Of course the try/except style could be used also just try on current CCCL knowing that it is wrong and things will fail most of the time.

@copy-pr-bot

copy-pr-bot Bot commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@github-actions github-actions Bot added the cuda.core Everything related to the cuda.core module label Jul 27, 2026
@seberg

seberg commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2c99b8a

@seberg
seberg force-pushed the tensor-map-linking branch from 2c99b8a to 679bdf5 Compare July 27, 2026 14:11
@seberg seberg mentioned this pull request Jul 28, 2026
Closed
1 task
@seberg
seberg force-pushed the tensor-map-linking branch from 679bdf5 to c640173 Compare July 29, 2026 08:34
@seberg

seberg commented Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c640173

seberg
seberg commented Jul 29, 2026
View reviewed changes
Comment thread cuda_core/cuda/core/_tensor_map.pyx
make_tma = _get_cccl_make_tma_fn()
if (make_tma != NULL
and device_type == _kDLCUDA
and box_dim[0] == box_dim[rank - 1]

@seberg seberg Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is so restrictive, that I am not sure it still makes sense. But the specific rules also feel too complex to bother with.
(Right now, this path seems useless to me, the only reason to use it and keep it is to exercise it a bit in the hope that it'll be the only path eventually.)

@github-actions

github-actions Bot commented Jul 29, 2026

Copy link
Copy Markdown
Doc Preview CI

🚀 View preview at
https://nvidia.github.io/cuda-python/pr-preview/pr-2428/

https://nvidia.github.io/cuda-python/pr-preview/pr-2428/cuda-core/

https://nvidia.github.io/cuda-python/pr-preview/pr-2428/cuda-bindings/

https://nvidia.github.io/cuda-python/pr-preview/pr-2428/cuda-pathfinder/


Preview will be ready when the GitHub Pages deployment is complete.

seberg added 2 commits August 28, 2026 14:13
@seberg
Make sure <cccl/tma> is found and soft-linked
c38c25e 
In my conda setup, compiling `tensor_map` created linking issues and
it seems like the reason this didn't fail here might be that it is
missing when it shouldn't be.

Adding the `cccl` headers breaks CI the same way, but moving things
to `_tensor_map_cccl.pxd` and soft-linking it should then unbreak
things again.

This now uses CCCL in an *extremely* limited fashion because CCCL
has some bugs/issues around validation.
(Even if fixed, we still need to support all dtypes in DLPack, though.)
@seberg
Fix MSVC CCCL preprocessor issue and add comment about CCCL fix
0d6e970
@seberg
seberg force-pushed the tensor-map-linking branch from c640173 to 0d6e970 Compare August 28, 2026 12:30
@seberg
seberg marked this pull request as ready for review August 28, 2026 12:31
@seberg seberg self-assigned this Aug 28, 2026
@seberg seberg added this to the cuda.core 1.3.0 milestone Aug 28, 2026
@seberg seberg added P1 Medium priority - Should do bug Something isn't working labels Aug 28, 2026
@github-actions github-actions Bot added the CI/CD CI/CD infrastructure label Aug 28, 2026
@seberg seberg changed the title WIP: Make sure TMA is compiled in (for now, see if this fails CI) WIP: Make sure CCCL TMA is compiled and use it (very narrowly) Aug 28, 2026
@seberg seberg changed the title WIP: Make sure CCCL TMA is compiled and use it (very narrowly) Make sure CCCL TMA is compiled and use it (very narrowly) Aug 28, 2026
leofang
leofang requested changes Aug 28, 2026
View reviewed changes

@leofang leofang left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Catching up. I think the problem is that the TMA support was vibe-coded in rush to a large extent, and that we never had a chance to sit down and sort out the build-time dependency issue. The ideal situation should be that CCCL is used as a submodule just like what we do in CuPy. If we were to do it, we probably can absorb the fix from upstream quickly, and avoid dynamic linking to cudart (which has its own class of problems, hence we have been avoiding this pretty much everywhere)?

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

Reviewers

@leofang leofang leofang requested changes

Requested changes must be addressed to merge this pull request.

Assignees

@seberg seberg

Labels

bug Something isn't working CI/CD CI/CD infrastructure cuda.core Everything related to the cuda.core module P1 Medium priority - Should do

Projects

None yet

Milestone

cuda.core 1.3.0

Development

Successfully merging this pull request may close these issues.

2 participants

@seberg @leofang