Add sm_86 (Ampere) and sm_89 (Ada) SASS in aarch64 builds

[agirault]

🚀 The feature, motivation and pitch

Request

Add SASS support for sm_86 and sm_89 to the aarch64 (sbsa) wheels

Motivation

Support NVIDIA IGX (aarch64) with discrete GPU in the same build as the SBSA wheels. The IGX supports the NVIDIA A6000 (Ampere) and NVIDIA RTX 6000 Ada.

Alternatives

No response

Additional context

• IGX is based on Arm Cortex-A78AE with a base arch of Armv8.2-a, higher than the Armv8.1-a requirement for NVPL which is needed in current arm64 builds.

• IGX is "Orin" like Jetson Orin (same SOC) but the dGPU stack is the same as SBSA, not as Jetpack (iGPU)

• I assume this would also enable running latest pytorch builds on NVIDIA L40S and NVIDIA L4 in arm datacenters.

cc @malfet @seemethere @ptrblck @msaroufim @eqy @jerryzh168 @snadampal @milpuz01 @aditew01 @nikhil-arm @fadara01 @puririshi98

[agirault]

cc @ptrblck

[ptrblck]

sm_89 to the aarch64 (sbsa) wheels

We won't add sm_89 as it's SASS binary compatible with sm_86 and only needed for e.g. FP8 kernels (and will be added in our cmake scripts if sm_86 is detected).

The IGX supports the NVIDIA A6000 (Ampere) and NVIDIA RTX 6000 Ada.

I assume no other GPU architectures are supported, e.g. A100?

[agirault]

We won't add sm_89 as it's SASS binary compatible with sm_86 and only needed for e.g. FP8 kernels

ack ty. was basing my request on https://developer.nvidia.com/cuda-gpus

Is there any docs that can help understand compatibility and extra features for minor versions? Is compatibility by major version (so 8 here_?

and will be added in our cmake scripts if sm_86 is detected

Not catching what you're referring to here.

I assume no other GPU architectures are supported, e.g. A100?

We do not, no need for 8.0 capability for our use cases on IGX+dGPU

[ptrblck]

Is there any docs that can help understand compatibility and extra features for minor versions? Is compatibility by major version (so 8 here_?

From https://docs.nvidia.com/cuda/ada-compatibility-guide/#compatibility-between-ampere-and-ada:

The NVIDIA Ada architecture is based on Ampere’s Instruction Set Architecture ISA 8.0, extending it with new instructions. As a consequence, any binary that runs on Ampere will be able to run on Ada (forward compatibility), but an Ada binary will not be able to run on Ampere.

CUDA applications built using CUDA Toolkit 11.0 through 11.7 are compatible with the NVIDIA Ada GPU architecture as long as they are built to include kernels in Ampere-native cubin (see Compatibility between Ampere and Ada) or PTX format (see Applications Built Using CUDA Toolkit 10.2 or Earlier), or both.

Unless you need to build Ada-specific instructions (e.g. for FP8) you don't need to build for sm_89.
In any case, this RFE targets the addition of sm_86 to our SBSA+CUDA binaries.
CC @nWEIdia @tinglvv

[agirault]

We won't add sm_89 as it's SASS binary compatible with sm_86 and only needed for e.g. FP8 kernels

Just pointing out that sm_89 is in the x86_64 python wheel (just twice). Might not change anything but just in case.

$ cuobjdump --list-elf /torch-venv/lib/python3.12/site-packages/torch/lib/libtorch_cuda.so 
| grep -o "sm_[0-9]\+" | sort -V | uniq -c
    378 sm_75
    378 sm_80
    378 sm_86
      2 sm_89
    381 sm_90
    379 sm_100
    378 sm_120

In any case, this RFE targets the addition of sm_86 to our SBSA+CUDA binaries.

Thank you.

[eqy]

We won't add sm_89 as it's SASS binary compatible with sm_86 and only needed for e.g. FP8 kernels

Just pointing out that sm_89 is in the x86_64 python wheel (just twice). Might not change anything but just in case.

$ cuobjdump --list-elf /torch-venv/lib/python3.12/site-packages/torch/lib/libtorch_cuda.so
| grep -o "sm_[0-9]+" | sort -V | uniq -c
378 sm_75
378 sm_80
378 sm_86
2 sm_89
381 sm_90
379 sm_100
378 sm_120

In any case, this RFE targets the addition of sm_86 to our SBSA+CUDA binaries.

Thank you.

This is expected, see also:

pytorch/aten/src/ATen/native/cuda/RowwiseScaledMM.cu

Line 494 in 1d77280

void f8f8bf16_rowwise_impl_sm89(

pytorch/cmake/Codegen.cmake

Line 119 in 1d77280

"${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/cuda/RowwiseScaledMM.cu"

[agirault]

Right. Excuse the naive question, but would these 2 kernels not be compatible for aarch64 as well? Just wondering why they wouldn't be added if we were to support sm_87 sm_86 globally. Or maybe I misunderstood the initial response, and sm_89 would be added, just conditionally like here?

[eqy]

I'm not sure where 8.7 came from as I thought this discussion was about 8.6 and 8.9. It would not make sense to add 8.9, as @ptrblck stated there are only exceptional cases where 8.9-specific features are used and in those cases that is defined already in the cmake file linked above. 8.9 is conditionally added, and the condition is that it is a kernel that uses 8.9-specific features such as fp8.

There would be no difference between x86 and aarch64 support for 8.6 and 8.9 if 8.6 is added to the builds via e.g., TORCH_CUDA_ARCH_LIST.

[agirault]

That last sentence clarifies it, I now understand what was meant by "not adding". Thank you.