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

[seberg]

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]

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.

[seberg]

/ok to test 2c99b8a

[seberg]

/ok to test c640173

[seberg]

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]

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.

[leofang]

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)?

[seberg]

Right, if we use it as a submodule that'll solve quite a few things (especially as we can just rely on the code being available and fixed if we pull in from current main).

Sounds good, that seems like a nice-to-have for future things as well.

If we were to do it, we probably can absorb the fix from upstream quickly, and avoid dynamic linking to cudart

The upstream fix isn't about cudart, that part is still needed. The fix is that the function itself had incorrect checks.

[seberg]

Converting to draft for now. It seems that CCCL should be able to fix the nvrtc dependency and then this all changes enough that this was a lot of effort for nothing (which is good of course).
(The other fixes needed should already have landed in CCCL main and would be usable if vendoring.)

I can open a hot-fix to just bake in the status-quo that we are never using CCCL until CCCL is ready for us here.

[leofang]

One note on the PR description:

TMA was never used and never tested in CI

I don't think it's true? We actually have a fallback path in which we call the driver APIs ourselves.

bot-generated code tracing

Why the CCCL path was dead in CI (pre-#2428)

1. On main, cuda_core/build_hooks.py only puts <cuda_path>/include on the compiler search path, but CCCL headers ship under <cuda_path>/include/cccl/…. So __has_include(<cuda/tma>) and __has_include("dlpack.h") both resolve to 0 (tensor_map.cpp).

2. With both macros 0, cuda_core_cccl_make_tma_descriptor_tiled compiles down to return 1; with the fixed sentinel message "CCCL <cuda/tma> and/or <dlpack/dlpack.h> not available at build time".

3. _tensor_map.pyx treats exactly that string as a soft fallback trigger — otherwise it silently continues to the driver-API path. See cuda_core/cuda/core/_tensor_map.pyx on main:

   if _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes):    # line 642
       ...
       with nogil:
           rc = cuda_core_cccl_make_tma_descriptor_tiled(...)          # line 670
       if rc == 0:
           ...
           return desc                                                 # line 698
       msg = errbuf[:].split(b"\0", 1)[0].decode(...)                  # line 700
       # If CCCL isn't available at build time, fall back to the direct
       # driver API path to preserve functionality on older toolchains.
       if "not available at build time" not in msg:                    # line 703
           raise ValueError(f"Failed to build TMA descriptor via CCCL: {msg}")
   cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]                    # line 706  <-- driver path
   byte_strides = _compute_byte_strides(...)
   ...
   HANDLE_RETURN(cydriver.cuTensorMapEncodeTiled(...))                 # line 740

   Note the indentation: line 706 is at the same level as if _tma_dtype_to_dlpack(...) on line 642, so it's outside the CCCL branch. When CCCL is compiled out, the sentinel check on line 703 is false, no exception is raised, and execution falls into the driver-encode block below.

4. Make sure CCCL TMA is compiled and use it (very narrowly) #2428's build_hooks.py patch is literally "Match nvcc: CCCL lives under include/cccl" — it prepends include/cccl — which is what makes the CCCL path start compiling (and then fall over at link time because CCCL pulls in CUDA-runtime symbols, per bullets 2/3 of the PR body).

What the "driver-API path" is

Direct calls to the driver's cuTensorMapEncode* family in _tensor_map.pyx:

• TensorMapDescriptor.from_tiled(...) → cydriver.cuTensorMapEncodeTiled(...) (_tensor_map.pyx:740)
• TensorMapDescriptor.from_im2col(...) → cydriver.cuTensorMapEncodeIm2col(...) (_tensor_map.pyx:891)
• TensorMapDescriptor.from_im2col_wide(...) → cydriver.cuTensorMapEncodeIm2colWide(...) (_tensor_map.pyx:1033)
• TensorMapDescriptor.replace_address(...) → cydriver.cuTensorMapReplaceAddress(...) (_tensor_map.pyx:1081)

These are driver entrypoints in libcuda (cuda.bindings.cydriver) — no CCCL headers involved. The pyx converts DLPack/CAI shape+strides itself (_compute_byte_strides plus the row-major → col-major reversal at _tensor_map.pyx:709-727) and hands raw CUtensorMap*-shaped args to the driver.

The CCCL path is essentially just a header-only C++ wrapper on top of that same driver call: cuda::make_tma_descriptor(DLTensor, box, …) in <cuda/tma> does DLPack validation + reversal + calls cuTensorMapEncodeTiled itself. Same descriptor, same driver entrypoint — the difference is just which side (C++/CCCL vs the pyx) massages DLPack into the driver's argument shape.

Evidence TMA is exercised via the driver path

The CI matrix (ci/test-matrix.yml) has rtxpro6000 (Blackwell, cc 12.0 → tensor_map_access_supported=True) rows. In run 34293223805, job Test linux-64 / Python 3.13, CUDA 13.3.0 (local), GPU rtxpro6000, every test_tensor_map.py::…::test_from_tiled_*, test_from_im2col_*, test_strided_memory_view_as_tensor_map* PASSED.

Those tests all go through TensorMapDescriptor.from_tiled → _from_tiled with default np.float32, so _tma_dtype_to_dlpack returns True → CCCL branch entered → shim returns rc=1 with the sentinel string → no raise → line 706 kicks in → cuTensorMapEncodeTiled succeeds. If the CCCL branch had actually errored out, those tests would all be red.

[seberg]

I don't think it's true? We actually have a fallback path in which we call the driver APIs ourselves.

Yes, unclear written: The CCCL version was never used/tested, the fallback seems all fine.