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

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 commented Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c640173

make_tma = _get_cccl_make_tma_fn()
if (make_tma != NULL
and device_type == _kDLCUDA
and box_dim[0] == box_dim[rank - 1]

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

Copy link
Copy Markdown

seberg added 2 commits August 28, 2026 14:13
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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

Development

Successfully merging this pull request may close these issues.

1 participant