Support LLVM21 + blackwell family

[johnnynunez]

Jax not support llvm21 and spark, thor and gb300

[johnnynunez]

last jax is not supporting GB300, Spark and Thor

#7 604.3 INFO: Build completed successfully, 17262 total actions
#7 604.3 
#7 604.3 
#7 604.3 2025-08-29 11:23:18,076 - INFO - Building jax-cuda-plugin for linux aarch64...
#7 604.3 2025-08-29 11:23:18,076 - INFO - [EXECUTING] ./bazel-7.4.1-linux-arm64 build --repo_env=HERMETIC_PYTHON_VERSION=3.12 --verbose_failures=true --config=cuda_libraries_from_stubs //jaxlib/tools:jax_cuda13_plugin_wheel
#7 604.3 INFO: Options provided by the client:
#7 604.3   Inherited 'common' options: --isatty=0 --terminal_columns=80
#7 604.3 INFO: Reading rc options for 'build' from /opt/jax/.bazelrc:
#7 604.3   Inherited 'common' options: --noenable_bzlmod --announce_rc --experimental_cc_shared_library --experimental_repo_remote_exec
#7 604.3 INFO: Reading rc options for 'build' from /opt/jax/.bazelrc:
#7 604.3   'build' options: --spawn_strategy=local --enable_platform_specific_config --incompatible_enable_cc_toolchain_resolution --repo_env USE_HERMETIC_CC_TOOLCHAIN=1 --define=grpc_no_ares=true --define=tsl_link_protobuf=true -c opt --output_filter=DONT_MATCH_ANYTHING --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir. --copt=-DNB_DOMAIN=jax --legacy_external_runfiles=false --repo_env=USE_PYWRAP_RULES=True --copt=-DGRPC_BAZEL_BUILD --host_copt=-DGRPC_BAZEL_BUILD --action_env=GRPC_BAZEL_RUNTIME=1 --repo_env=PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=upb --action_env=PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=upb --repo_env=RULES_PYTHON_ENABLE_PYSTAR=0
#7 604.3 INFO: Reading rc options for 'build' from /opt/jax/.jax_configure.bazelrc:
#7 604.3   'build' options: --repo_env=HERMETIC_PYTHON_VERSION=3.12 --verbose_failures=true --config=clang_local --action_env=CLANG_COMPILER_PATH=/usr/lib/llvm-21/bin/clang --repo_env=CC=/usr/lib/llvm-21/bin/clang --repo_env=CXX=/usr/lib/llvm-21/bin/clang++ --repo_env=BAZEL_COMPILER=/usr/lib/llvm-21/bin/clang --config=clang --config=mkl_aarch64_threadpool --config=cuda13 --action_env=CLANG_CUDA_COMPILER_PATH=/usr/lib/llvm-21/bin/clang --config=build_cuda_with_nvcc --repo_env=HERMETIC_CUDA_VERSION=13.0.0 --repo_env=HERMETIC_CUDNN_VERSION=9.12.0 --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES=sm_87,sm_89,sm_90,sm_100,sm_103,sm_110,sm_120,sm_121 --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES=sm_87,sm_89,sm_90,sm_100,sm_103,sm_110,sm_120,sm_121 --@local_config_cuda//:cuda_compiler=nvcc --action_env=TF_NVCC_CLANG=0 --action_env=CUDA_HOME=/usr/local/cuda --repo_env=USE_HERMETIC_CC_TOOLCHAIN=1
#7 604.3 INFO: Found applicable config definition build:clang_local in file /opt/jax/.bazelrc: --noincompatible_enable_cc_toolchain_resolution --repo_env USE_HERMETIC_CC_TOOLCHAIN=0
#7 604.3 INFO: Found applicable config definition build:clang in file /opt/jax/.bazelrc: --copt=-Wno-gnu-offsetof-extensions --copt=-Qunused-arguments --copt=-Werror=mismatched-tags --copt=-Wno-error=c23-extensions
#7 604.3 INFO: Found applicable config definition build:mkl_aarch64_threadpool in file /opt/jax/.bazelrc: --define=build_with_mkl_aarch64=true --@compute_library//:openmp=false -c opt
#7 604.3 INFO: Found applicable config definition build:cuda13 in file /opt/jax/.bazelrc: --config=cuda_common --config=cuda_v13
#7 604.3 INFO: Found applicable config definition build:cuda_common in file /opt/jax/.bazelrc: --repo_env TF_NEED_CUDA=1 --repo_env TF_NCCL_USE_STUB=1 --@local_config_cuda//:enable_cuda --@local_config_cuda//cuda:include_cuda_libs=true --linkopt=-Wl,--disable-new-dtags
#7 604.3 INFO: Found applicable config definition build:cuda_v13 in file /opt/jax/.bazelrc: --repo_env=HERMETIC_CUDA_VERSION=13.0.0 --repo_env=HERMETIC_CUDNN_VERSION=9.12.0 --repo_env=HERMETIC_NVSHMEM_VERSION=3.3.20 --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=sm_87,sm_89,sm_90,sm_100,compute_103,compute_110,compute_120,compute_121
#7 604.3 INFO: Found applicable config definition build:build_cuda_with_nvcc in file /opt/jax/.bazelrc: --action_env=TF_NVCC_CLANG=1 --@local_config_cuda//:cuda_compiler=nvcc
#7 604.3 INFO: Found applicable config definition build:cuda_libraries_from_stubs in file /opt/jax/.bazelrc: --@local_config_cuda//cuda:include_cuda_libs=false
#7 604.3 INFO: Found applicable config definition build:linux in file /opt/jax/.bazelrc: --config=posix --copt=-Wno-unknown-warning-option --copt=-Wno-stringop-truncation --copt=-Wno-array-parameter --copt=-Wno-deprecated-register --copt=-Wno-register
#7 604.3 INFO: Found applicable config definition build:posix in file /opt/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++17 --host_cxxopt=-std=c++17
#7 604.4 + true
#7 604.4 + f=/root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/cutlass_archive/include/cutlass/matrix.h
#7 604.4 + '[' -f /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/cutlass_archive/include/cutlass/matrix.h ']'
#7 604.4 + sleep 1
#7 604.5 Computing main repo mapping: 
#7 604.6 Loading: 
#7 604.6 Loading: 0 packages loaded
#7 605.3 Analyzing: target //jaxlib/tools:jax_cuda13_plugin_wheel (0 packages loaded, 0 targets configured)
#7 605.3 Analyzing: target //jaxlib/tools:jax_cuda13_plugin_wheel (0 packages loaded, 0 targets configured)
#7 605.3 
#7 605.4 + true
#7 605.4 + f=/root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/cutlass_archive/include/cutlass/matrix.h
#7 605.4 + '[' -f /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/cutlass_archive/include/cutlass/matrix.h ']'
#7 605.4 + sleep 1
#7 605.9 INFO: Analyzed target //jaxlib/tools:jax_cuda13_plugin_wheel (14 packages loaded, 2414 targets configured).
#7 606.0 ERROR: /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/xla/xla/stream_executor/cuda/BUILD:404:13: Compiling xla/stream_executor/cuda/delay_kernel_cuda.cu.cc failed: (Exit 1): clang failed: error executing CppCompile command (from target @@xla//xla/stream_executor/cuda:delay_kernel_cuda) 
#7 606.0   (cd /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/execroot/__main__ && \
#7 606.0   exec env - \
#7 606.0     CLANG_COMPILER_PATH=/usr/lib/llvm-21/bin/clang \
#7 606.0     CLANG_CUDA_COMPILER_PATH=/usr/lib/llvm-21/bin/clang \
#7 606.0     CUDA_HOME=/usr/local/cuda \
#7 606.0     GRPC_BAZEL_RUNTIME=1 \
#7 606.0     LD_LIBRARY_PATH=/usr/local/cuda/compat:/usr/local/cuda/lib64:/usr/local/cuda/compat:/usr/local/cuda/lib64: \
#7 606.0     PATH=/usr/local/cuda/bin:/opt/venv/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \
#7 606.0     PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=upb \
#7 606.0     PWD=/proc/self/cwd \
#7 606.0     TF_NVCC_CLANG=0 \
#7 606.0   /usr/lib/llvm-21/bin/clang -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections '-std=c++14' -MD -MF bazel-out/aarch64-opt/bin/external/xla/xla/stream_executor/cuda/_objs/delay_kernel_cuda/delay_kernel_cuda.cu.pic.d '-frandom-seed=bazel-out/aarch64-opt/bin/external/xla/xla/stream_executor/cuda/_objs/delay_kernel_cuda/delay_kernel_cuda.cu.pic.o' -fPIC '-DGOOGLE_PROTOBUF_USING_BAZEL=1' '-DEIGEN_MAX_ALIGN_BYTES=64' -DEIGEN_ALLOW_UNALIGNED_SCALARS '-DEIGEN_USE_AVX512_GEMM_KERNELS=0' -DHAVE_SYS_UIO_H -DTF_USE_SNAPPY -DTENSORFLOW_USE_NUMA '-DLLVM_ON_UNIX=1' '-DHAVE_BACKTRACE=1' '-DBACKTRACE_HEADER=<execinfo.h>' '-DLTDL_SHLIB_EXT=".so"' '-DLLVM_PLUGIN_EXT=".so"' '-DLLVM_ENABLE_LLVM_EXPORT_ANNOTATIONS=1' '-DLLVM_ENABLE_PLUGINS=1' '-DLLVM_ENABLE_THREADS=1' '-DHAVE_DEREGISTER_FRAME=1' '-DHAVE_LIBPTHREAD=1' '-DHAVE_PTHREAD_GETNAME_NP=1' '-DHAVE_PTHREAD_H=1' '-DHAVE_PTHREAD_SETNAME_NP=1' '-DHAVE_REGISTER_FRAME=1' '-DHAVE_SETENV_R=1' '-DHAVE_STRERROR_R=1' '-DHAVE_SYSEXITS_H=1' '-DHAVE_SYS_IOCTL_H=1' '-DHAVE_UNISTD_H=1' -D_GNU_SOURCE '-DHAVE_GETAUXVAL=1' '-DHAVE_MALLINFO=1' '-DHAVE_SBRK=1' '-DHAVE_STRUCT_STAT_ST_MTIM_TV_NSEC=1' -DHAVE_BUILTIN_THREAD_POINTER '-DLLVM_NATIVE_ARCH="AArch64"' '-DLLVM_NATIVE_ASMPARSER=LLVMInitializeAArch64AsmParser' '-DLLVM_NATIVE_ASMPRINTER=LLVMInitializeAArch64AsmPrinter' '-DLLVM_NATIVE_DISASSEMBLER=LLVMInitializeAArch64Disassembler' '-DLLVM_NATIVE_TARGET=LLVMInitializeAArch64Target' '-DLLVM_NATIVE_TARGETINFO=LLVMInitializeAArch64TargetInfo' '-DLLVM_NATIVE_TARGETMC=LLVMInitializeAArch64TargetMC' '-DLLVM_NATIVE_TARGETMCA=LLVMInitializeAArch64TargetMCA' '-DLLVM_HOST_TRIPLE="aarch64-unknown-linux-gnu"' '-DLLVM_DEFAULT_TARGET_TRIPLE="aarch64-unknown-linux-gnu"' '-DLLVM_VERSION_MAJOR=22' '-DLLVM_VERSION_MINOR=0' '-DLLVM_VERSION_PATCH=0' '-DLLVM_VERSION_STRING="22.0.0git"' -D__STDC_LIMIT_MACROS -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS '-DLLVM_HAS_AArch64_TARGET=1' '-DLLVM_HAS_AMDGPU_TARGET=1' '-DLLVM_HAS_ARM_TARGET=1' '-DLLVM_HAS_NVPTX_TARGET=1' '-DLLVM_HAS_PowerPC_TARGET=1' '-DLLVM_HAS_RISCV_TARGET=1' '-DLLVM_HAS_SystemZ_TARGET=1' '-DLLVM_HAS_X86_TARGET=1' '-DBLAKE3_USE_NEON=0' -DBLAKE3_NO_AVX2 -DBLAKE3_NO_AVX512 -DBLAKE3_NO_SSE2 -DBLAKE3_NO_SSE41 '-DCLANG_CUDA=1' -iquote external/xla -iquote bazel-out/aarch64-opt/bin/external/xla -iquote external/com_google_protobuf -iquote bazel-out/aarch64-opt/bin/external/com_google_protobuf -iquote external/com_google_absl -iquote bazel-out/aarch64-opt/bin/external/com_google_absl -iquote external/zlib -iquote bazel-out/aarch64-opt/bin/external/zlib -iquote external/tsl -iquote bazel-out/aarch64-opt/bin/external/tsl -iquote external/eigen_archive -iquote bazel-out/aarch64-opt/bin/external/eigen_archive -iquote external/ml_dtypes_py -iquote bazel-out/aarch64-opt/bin/external/ml_dtypes_py -iquote external/snappy -iquote bazel-out/aarch64-opt/bin/external/snappy -iquote external/hwloc -iquote bazel-out/aarch64-opt/bin/external/hwloc -iquote external/com_googlesource_code_re2 -iquote bazel-out/aarch64-opt/bin/external/com_googlesource_code_re2 -iquote external/llvm-project -iquote bazel-out/aarch64-opt/bin/external/llvm-project -iquote external/local_config_cuda -iquote bazel-out/aarch64-opt/bin/external/local_config_cuda -iquote external/cuda_cudart -iquote bazel-out/aarch64-opt/bin/external/cuda_cudart -iquote external/cuda_cublas -iquote bazel-out/aarch64-opt/bin/external/cuda_cublas -iquote external/cuda_cccl -iquote bazel-out/aarch64-opt/bin/external/cuda_cccl -iquote external/cuda_nvtx -iquote bazel-out/aarch64-opt/bin/external/cuda_nvtx -iquote external/cuda_nvcc -iquote bazel-out/aarch64-opt/bin/external/cuda_nvcc -iquote external/cuda_cusolver -iquote bazel-out/aarch64-opt/bin/external/cuda_cusolver -iquote external/cuda_cufft -iquote bazel-out/aarch64-opt/bin/external/cuda_cufft -iquote external/cuda_cusparse -iquote bazel-out/aarch64-opt/bin/external/cuda_cusparse -iquote external/cuda_curand -iquote bazel-out/aarch64-opt/bin/external/cuda_curand -iquote external/cuda_cupti -iquote bazel-out/aarch64-opt/bin/external/cuda_cupti -iquote external/cuda_nvml -iquote bazel-out/aarch64-opt/bin/external/cuda_nvml -iquote external/cuda_nvjitlink -iquote bazel-out/aarch64-opt/bin/external/cuda_nvjitlink -iquote external/cuda_crt -iquote bazel-out/aarch64-opt/bin/external/cuda_crt -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/protobuf -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/internal_visibility -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/micro_string -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/arena -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/arena_align -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/port -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/stubs/_virtual_includes/lite -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/arena_allocation_policy -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/arena_cleanup -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/string_block -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/protobuf_lite -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/endian -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/varint_shuffle -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/io/_virtual_includes/io -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/io/_virtual_includes/io_win32 -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/third_party/utf8_range/_virtual_includes/utf8_validity -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/third_party/utf8_range/_virtual_includes/utf8_range -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/io/_virtual_includes/gzip_stream -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/stubs/_virtual_includes/stubs -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/io/_virtual_includes/printer -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/io/_virtual_includes/zero_copy_sink -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/io/_virtual_includes/tokenizer -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/any_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/api_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/source_context_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/type_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/compiler/_virtual_includes/plugin_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/descriptor_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/duration_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/empty_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/field_mask_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/struct_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/timestamp_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/wrappers_proto -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/protobuf_layering_check_legacy -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/compiler/_virtual_includes/importer -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/util/_virtual_includes/delimited_message_util -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/util/_virtual_includes/differencer -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/util/_virtual_includes/field_mask_util -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/util/_virtual_includes/json_util -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/json -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/parser -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/descriptor_traits -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/lexer -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/message_path -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/zero_copy_buffered_stream -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/untyped_message -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/util/_virtual_includes/type_resolver -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/_virtual_includes/descriptor_legacy -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/unparser -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/json/_virtual_includes/writer -Ibazel-out/aarch64-opt/bin/external/com_google_protobuf/src/google/protobuf/util/_virtual_includes/time_util -Ibazel-out/aarch64-opt/bin/external/llvm-project/third-party/siphash/_virtual_includes/siphash -Ibazel-out/aarch64-opt/bin/external/local_config_cuda/cuda/_virtual_includes/cuda_headers -Ibazel-out/aarch64-opt/bin/external/cuda_cudart/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_cublas/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_cccl/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_nvtx/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_nvcc/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_cusolver/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_cufft/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_cusparse/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_curand/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_cupti/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_nvml/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_nvjitlink/_virtual_includes/headers -Ibazel-out/aarch64-opt/bin/external/cuda_crt/_virtual_includes/headers -isystem external/zlib -isystem bazel-out/aarch64-opt/bin/external/zlib -isystem external/eigen_archive -isystem bazel-out/aarch64-opt/bin/external/eigen_archive -isystem external/eigen_archive/mkl_include -isystem bazel-out/aarch64-opt/bin/external/eigen_archive/mkl_include -isystem external/hwloc/hwloc -isystem bazel-out/aarch64-opt/bin/external/hwloc/hwloc -isystem external/hwloc/include -isystem bazel-out/aarch64-opt/bin/external/hwloc/include -isystem external/llvm-project/llvm/include -isystem bazel-out/aarch64-opt/bin/external/llvm-project/llvm/include -isystem external/local_config_cuda/cuda -isystem bazel-out/aarch64-opt/bin/external/local_config_cuda/cuda -isystem external/cuda_cudart/include -isystem bazel-out/aarch64-opt/bin/external/cuda_cudart/include -isystem external/cuda_cublas/include -isystem bazel-out/aarch64-opt/bin/external/cuda_cublas/include -isystem external/cuda_cccl/include/cccl -isystem bazel-out/aarch64-opt/bin/external/cuda_cccl/include/cccl -isystem external/cuda_nvtx/include -isystem bazel-out/aarch64-opt/bin/external/cuda_nvtx/include -isystem external/cuda_nvcc/include -isystem bazel-out/aarch64-opt/bin/external/cuda_nvcc/include -isystem external/cuda_cusolver/include -isystem bazel-out/aarch64-opt/bin/external/cuda_cusolver/include -isystem external/cuda_cufft/include -isystem bazel-out/aarch64-opt/bin/external/cuda_cufft/include -isystem external/cuda_cusparse/include -isystem bazel-out/aarch64-opt/bin/external/cuda_cusparse/include -isystem external/cuda_curand/include -isystem bazel-out/aarch64-opt/bin/external/cuda_curand/include -isystem external/cuda_cupti/include -isystem bazel-out/aarch64-opt/bin/external/cuda_cupti/include -isystem external/cuda_nvml/include -isystem bazel-out/aarch64-opt/bin/external/cuda_nvml/include -isystem external/cuda_nvjitlink/include -isystem bazel-out/aarch64-opt/bin/external/cuda_nvjitlink/include -isystem external/cuda_crt/include -isystem bazel-out/aarch64-opt/bin/external/cuda_crt/include '-fvisibility=hidden' -Wno-sign-compare -Wno-unknown-warning-option -Wno-stringop-truncation -Wno-array-parameter -Wno-deprecated-register -Wno-register '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' '-DNB_DOMAIN=jax' -DGRPC_BAZEL_BUILD -Wno-gnu-offsetof-extensions -Qunused-arguments '-Werror=mismatched-tags' '-Wno-error=c23-extensions' '-std=c++17' -x cuda '-DGOOGLE_CUDA=1' '--no-cuda-include-ptx=all' '--cuda-feature=+ptx83' '--cuda-gpu-arch=sm_87' '--cuda-gpu-arch=sm_89' '--cuda-gpu-arch=sm_90' '--cuda-gpu-arch=sm_100' '--cuda-gpu-arch=sm_103' '--cuda-gpu-arch=sm_110' '--cuda-gpu-arch=sm_120' '--cuda-gpu-arch=sm_121' '-Xcuda-fatbinary=--compress-all' '-nvcc_options=expt-relaxed-constexpr' -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/xla/xla/stream_executor/cuda/delay_kernel_cuda.cu.cc -o bazel-out/aarch64-opt/bin/external/xla/xla/stream_executor/cuda/_objs/delay_kernel_cuda/delay_kernel_cuda.cu.pic.o)
#7 606.0 # Configuration: e25c07e5ab2ab0cdbb4db9d809c42cbccf51350f08462e121af64de4c09a9c07
#7 606.0 # Execution platform: @@local_execution_config_platform//:platform
#7 606.0 clang: error: unknown argument: '-Xcuda-fatbinary=--compress-all'
#7 606.0 clang: error: unknown argument: '-nvcc_options=expt-relaxed-constexpr'
#7 606.0 clang: warning: CUDA version is newer than the latest partially supported version 12.8 [-Wunknown-cuda-version]
#7 606.0 clang: error: unsupported CUDA gpu architecture: sm_103
#7 606.1 Target //jaxlib/tools:jax_cuda13_plugin_wheel failed to build
#7 606.1 INFO: Elapsed time: 1.801s, Critical Path: 0.10s
#7 606.1 INFO: 88 processes: 67 internal, 21 local.
#7 606.1 ERROR: Build did NOT complete successfully
#7 606.1 Traceback (most recent call last):
#7 606.1   File "/opt/jax/build/build.py", line 760, in <module>
#7 606.1     asyncio.run(main())
#7 606.1   File "/usr/lib/python3.12/asyncio/runners.py", line 194, in run
#7 606.1     return runner.run(main)
#7 606.1            ^^^^^^^^^^^^^^^^
#7 606.1   File "/usr/lib/python3.12/asyncio/runners.py", line 118, in run
#7 606.1     return self._loop.run_until_complete(task)
#7 606.1            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#7 606.1   File "/usr/lib/python3.12/asyncio/base_events.py", line 687, in run_until_complete
#7 606.1     return future.result()
#7 606.1            ^^^^^^^^^^^^^^^
#7 606.1   File "/opt/jax/build/build.py", line 706, in main
#7 606.1     raise RuntimeError(f"Command failed with return code {result.return_code}")
#7 606.1 RuntimeError: Command failed with return code 1
#7 ERROR: process "/bin/sh -c /tmp/JAX/install.sh || /tmp/JAX/build.sh" did not complete successfully: exit code: 1
------
 > [3/3] RUN /tmp/JAX/install.sh || /tmp/JAX/build.sh:
606.1            ^^^^^^^^^^^^^^^^
606.1   File "/usr/lib/python3.12/asyncio/runners.py", line 118, in run
606.1     return self._loop.run_until_complete(task)
606.1            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
606.1   File "/usr/lib/python3.12/asyncio/base_events.py", line 687, in run_until_complete
606.1     return future.result()
606.1            ^^^^^^^^^^^^^^^
606.1   File "/opt/jax/build/build.py", line 706, in main
606.1     raise RuntimeError(f"Command failed with return code {result.return_code}")
606.1 RuntimeError: Command failed with return code 1
------

 �[33m1 warning found (use docker --debug to expand):
�[0m - InvalidDefaultArgInFrom: Default value for ARG ${BASE_IMAGE} results in empty or invalid base image name (line 10)
Dockerfile:27
--------------------
  25 |     
  26 |     # attempt to install JAX from pip, and fall back to building it if the installation fails
  27 | >>> RUN /tmp/JAX/install.sh || /tmp/JAX/build.sh
  28 |         
  29 |     
--------------------
ERROR: failed to solve: process "/bin/sh -c /tmp/JAX/install.sh || /tmp/JAX/build.sh" did not complete successfully: exit code: 1

[hawkinsp]

I don't think this is problem with JAX itself, rather you are building with an incompatible toolchain. I've asked some of my colleagues to take a look.

[johnnynunez]

I don't think this is problem with JAX itself, rather you are building with an incompatible toolchain. I've asked some of my colleagues to take a look.

yeah, i tried to force nvcc but it is still taking clang20...

[yuriivcs]

Hi,
Looks like it's a problem with toolchains configuration. Usually it happens when you add more parameters with clang path than you need, and bazel uses clang instead of NVCC wrapper.
Could you share the original command which you run?

[johnnynunez]

Hi, Looks like it's a problem with toolchains configuration. Usually it happens when you add more parameters with clang path than you need, and bazel uses clang instead of NVCC wrapper. Could you share the original command which you run?

it is in the scripts:

#!/usr/bin/env bash
# JAX builder for Jetson (architecture: ARM64, CUDA support)
set -ex

echo "Building JAX for Jetson"

# Clone JAX repository
git clone --branch "jax-v${JAX_BUILD_VERSION}" --depth=1 --recursive https://github.com/google/jax /opt/jax || \
git clone --depth=1 --recursive https://github.com/google/jax /opt/jax

cd /opt/jax

# Build jaxlib from source with detected versions
BUILD_FLAGS='--disable_nccl '
BUILD_FLAGS+='--cuda_compute_capabilities="sm_87,sm_89,sm_90,sm_100,sm_103,sm_110,sm_120,sm_121" '
BUILD_FLAGS+='--cuda_version=13.0.0 --cudnn_version=9.12.0 '
# BUILD_FLAGS+='--bazel_options=--repo_env=LOCAL_CUDA_PATH="/usr/local/cuda-12.8"'
# BUILD_FLAGS+='--bazel_options=--repo_env=LOCAL_CUDNN_PATH="/opt/nvidia/cudnn/"'
BUILD_FLAGS+='--output_path=$PIP_WHEEL_DIR '
BUILD_FLAGS+='--clang_path=/usr/lib/llvm-20/bin/clang'

python3 build/build.py requirements_update

# Start background process to monitor and patch matrix.h
(
    while true; do
        if [ -f "/root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/cutlass_archive/include/cutlass/matrix.h" ]; then
            echo -e "\e[1;33m[PATCH] Found matrix.h, applying patch...\e[0m"
            sed -i 's/set_slice3x3/set_slice_3x3/g' /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/cutlass_archive/include/cutlass/matrix.h
            echo -e "\e[1;32m[PATCH] Patch applied successfully!\e[0m"
            break
        fi
        sleep 1
    done
) &

# Run the build
python3 build/build.py build $BUILD_FLAGS --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt

# Build the jax pip wheels
pip3 wheel --wheel-dir=$PIP_WHEEL_DIR --no-deps --verbose .

# Upload the wheels to mirror
twine upload --verbose $PIP_WHEEL_DIR/jaxlib-*.whl || echo "failed to upload wheel to ${TWINE_REPOSITORY_URL}"
twine upload --verbose $PIP_WHEEL_DIR/jax_cuda12_pjrt-*.whl || echo "failed to upload wheel to ${TWINE_REPOSITORY_URL}"
twine upload --verbose $PIP_WHEEL_DIR/jax_cuda12_plugin-*.whl || echo "failed to upload wheel to ${TWINE_REPOSITORY_URL}"
twine upload --verbose $PIP_WHEEL_DIR/jax-*.whl || echo "failed to upload wheel to ${TWINE_REPOSITORY_URL}"

# Install them into the container
cd $PIP_WHEEL_DIR/
pip3 install jaxlib*.whl jax_cuda12_plugin*.whl jax_cuda12_pjrt*.whl opt_einsum
pip3 install --no-dependencies jax*.whl
cd /opt/jax

PD: https://github.com/dusty-nv/jetson-containers/blob/master/packages/ml/jax/build.sh

But i tried a lot of things:
removing BUILD_FLAGS+='--clang_path=/usr/lib/llvm-20/bin/clang'
removing llvm20 from docker, force to use nvcc etc, and always when is generating cuda plugin, says that sm_103, or sm_110 is not supported when in cuda 13.0 yes

and it is compatible:
openxla/xla@b16d4bb

[hawkinsp]

FWIW I will note that we are now shipping JAX CUDA 13 aarch64 nightly wheels:

uv pip install -U --pre jax jaxlib "jax-cuda13-plugin[with-cuda]" jax-cuda13-pjrt -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --extra-index-url https://pypi.nvidia.com

They won't work on Thor at the moment (the most recent nightly doesn't have openxla/xla#30514 but the next nightly should).

[johnnynunez]

FWIW I will note that we are now shipping JAX CUDA 13 aarch64 nightly wheels:

uv pip install -U --pre jax jaxlib "jax-cuda13-plugin[with-cuda]" jax-cuda13-pjrt -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --extra-index-url https://pypi.nvidia.com

They won't work on Thor at the moment (the most recent nightly doesn't have openxla/xla#30514 but the next nightly should).

but what am i doing wrong? seems that is from toolchain

[yuriivcs]

@johnnynunez let me check.
What version of JAX_BUILD_VERSION you are you using?

[johnnynunez]

@johnnynunez let me check. What version of JAX_BUILD_VERSION you are you using?

from main (last commit)

[johnnynunez]

llvm21 bazel-contrib/toolchains_llvm@7cfc09d @yuriivcs @hawkinsp

[yuriivcs]

Do you expect it to be built by NVCC or Clang only?

[johnnynunez]

Do you expect it to be built by NVCC or Clang only?

I have seen that it is not compatible with clang, so I was waiting to compile it with nvcc but in the process it always taking clang (not expected result)

[yuriivcs]

Yes, you are right. But the reason for your issue is that bazel takes Clang, so I just wanted to clarify.