Install CuTeDSL into JAX containers and remove old CuTeDSL + JAX project

[mgoldfarb-nvidia]

CuTeDSL + JAX is now formally a part of CuTeDSL in the cutlass.jax sub-module. Users can now install directly from the official CuTeDSL release https://pypi.org/project/nvidia-cutlass-dsl/. Examples and documentation can be found at https://github.com/NVIDIA/cutlass/tree/main

[gpupuck]

There is job failure https://github.com/NVIDIA/JAX-Toolbox/actions/runs/23462382532/job/68271153801 that I don't know how to explain.

[olupton]

LGTM, I think the tests all failed because of a JAX-side bug fixed by jax-ml/jax#36185

[gpupuck]

Ready to merge?

[mgoldfarb-nvidia]

Ready to merge?

Appears the tests are failing for some reason:

=== /jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py ===
Traceback (most recent call last):
  File "/jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py", line 33, in <module>
    import cutlass
  File "/usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/jax/types.py", line 65, in <module>
    @jax.tree_util.register_dataclass
     ^^^^^^^^^^^^^
AttributeError: module 'jax' has no attribute 'tree_util'

This is surprising since jax.tree_util.register_dataclass is a valid module inside jax.

[olupton]

Ready to merge?

Appears the tests are failing for some reason:

=== /jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py ===
Traceback (most recent call last):
  File "/jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py", line 33, in <module>
    import cutlass
  File "/usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/jax/types.py", line 65, in <module>
    @jax.tree_util.register_dataclass
     ^^^^^^^^^^^^^
AttributeError: module 'jax' has no attribute 'tree_util'

This is surprising since jax.tree_util.register_dataclass is a valid module inside jax.

I suspect it's because PYTHONPATH=${CUTLASS_EXAMPLES_ROOT} puts a directory containing a jax folder first in the path.

[mgoldfarb-nvidia]

Ready to merge?

Appears the tests are failing for some reason:

=== /jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py ===
Traceback (most recent call last):
  File "/jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py", line 33, in <module>
    import cutlass
  File "/usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/jax/types.py", line 65, in <module>
    @jax.tree_util.register_dataclass
     ^^^^^^^^^^^^^
AttributeError: module 'jax' has no attribute 'tree_util'

This is surprising since jax.tree_util.register_dataclass is a valid module inside jax.

I suspect it's because PYTHONPATH=${CUTLASS_EXAMPLES_ROOT} puts a directory containing a jax folder first in the path.

Ready to merge?

Appears the tests are failing for some reason:

=== /jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py ===
Traceback (most recent call last):
  File "/jax-cutlass-src/cutlass/examples/python/CuTeDSL/jax/cutlass_call_basic.py", line 33, in <module>
    import cutlass
  File "/usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/jax/types.py", line 65, in <module>
    @jax.tree_util.register_dataclass
     ^^^^^^^^^^^^^
AttributeError: module 'jax' has no attribute 'tree_util'

This is surprising since jax.tree_util.register_dataclass is a valid module inside jax.

I suspect it's because PYTHONPATH=${CUTLASS_EXAMPLES_ROOT} puts a directory containing a jax folder first in the path.

Okay we are down to one more test to running - the sharding test... Log has no output at all from the last run. edit: exit code 139, appears to be a segfault.