More ROCm support

[glen-amd]

Goals

• Bring the support for AMD ROCm to TorchServe
• Let the TorchServe community be aware of the support for AMD GPUs

Git PRs

• 740 add generic support for different gpu hardware #3371
• ROCm support for Dockerfile and build_image.sh #3381

Already done

• NVIDIA device information via CLI nvidia-smi
• AMD device information via CLI amd-smi/rocm-smi
• More...

TODOs

Must do

• Support for the latest ROCm release
  • Currently supported
    • choices=["rocm6.0", "rocm6.1", "rocm6.2"]
      • In the file "ts_scripts/install_dependencies.py"
  • Latest release by 20250311
    • ROCm 6.3 (rocm6.3)

Nice to do

• NVML instead of CLI nvidia-smi
  • NVML has both C/C++ APIs and Python bindings.
  • TODO: JNI for Java
• AMD SMI lib instead of CLI amd-smi/rocm-smi
  • AMD SMI lib has both C/C++ APIs and Python bindings.
  • TODO: JNI for Java

Exploration in the TorchServe master branch

• Commands
  • find . -type f,l | xargs grep --color=always -nri cuda
  • find . -type f,l | xargs grep --color=always -nriE '\Wnv'
  • find . -type f,l | xargs grep --color=always -nri '_nv'
• File types

Parts

Requirement files

Docker files

Config files

• ts_scripts/spellcheck_conf/wordlist.txt

Build scripts

Frontend

Backend

• cpp/src/backends/handler/base_handler.cc

Documentation

Examples

CI

Regression tests

GitHub workflows

Benchmarks

• benchmarks/install_dependencies.sh
• benchmarks/benchmark.py
  • "nvidia-docker"

Notes

Code name examples

• NVIDIA/CUDA
  • cu92, cu101, cu102, cu111, cu113, cu116, cu117, cu118, cu121
• AMD/ROCm
  • rocm5.9, rocm6.0, rocm6.1, rocm6.2, rocm6.3

---

---

Description

Please read our CONTRIBUTING.md prior to creating your first pull request.

Please include a summary of the feature or issue being fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes #(issue)

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Feature/Issue validation/testing

Please describe the Unit or Integration tests that you ran to verify your changes and relevant result summary. Provide instructions so it can be reproduced.
Please also list any relevant details for your test configuration.

• Test A
  Logs for Test A

• Test B
  Logs for Test B

Checklist:

• Did you have fun?
• Have you added tests that prove your fix is effective or that this feature works?
• Has code been commented, particularly in hard-to-understand areas?
• Have you made corresponding changes to the documentation?

[glen-amd]

@jakki-amd / @smedegaard / @agunapal - can you please make initial review? Thanks.

[jeffdaily]

Why are you using HIP for the device here instead of CUDA? ROCm PyTorch masquerades as the CUDA device. Though a HIP device does exist as a dispatch key, no kernels are registered to it.

[glen-amd]

I will revert this change.

Since kHIP is explicitly defined in PyTorch (e.g., for clearer semantics I guess), I expected that automatic re-mapping would happen internally for kernel registration, lookup, etc. However, I don't really see this in PyTorch.

[jakki-amd]

nit: I think it would be more consistent to follow the same naming convention as with CUDA flag naming, meaning using rocm61 instead of rocm6.1 as CUDA flags are also given like cu111, not cu11.1.

[glen-amd]

I specifically checked the naming convention of both CUDA and ROCm and confirmed internally with some AMDers, then decided to use something like rocm6.3 instead of rocm63.

[jakki-amd]

I don't see how defining cuda label to be rocm6.1 without changing ci-gpu CI script would work. Some work was done regarding CI scripts in this branch here, but the branch is out-of-date.

[jakki-amd]

Left few minor comments, otherwise looks good!