AMD GPU & ROCm Deep Dive: Can It Challenge CUDA for LLM Inference?

RDNA Family (Consumer GPUs):
- RX 7900 XTX (RDNA 3): 24GB GDDR6, 960 GB/s bandwidth
- RX 7900 XT (RDNA 3): 20GB GDDR6, 800 GB/s bandwidth
- Optimized for graphics rendering (rasterization, ray tracing)
- Cache hierarchy tuned for maximum gaming performance
- ML workload support: possible, but official ROCm support is limited

CDNA Family (Datacenter GPUs):
┌─────────────────────────────────────────────────────────────────┐
│ AMD MI300X (CDNA 3, released 2023)                              │
│                                                                 │
│  • 192GB HBM3 memory (industry record!)                        │
│    → 2.4x more than H100 SXM's 80GB                            │
│  • 5.3 TB/s memory bandwidth                                   │
│    → 1.58x more than H100 SXM's 3.35 TB/s                     │
│  • 304 Compute Units                                           │
│  • 1,307 TFLOPS FP16                                           │
│  • 655 TFLOPS FP32                                             │
│  • MCM (Multi-Chip Module): GPU + CPU HBM integrated           │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

                    AMD MI300X          NVIDIA H100 SXM
Memory:             192GB HBM3          80GB HBM3
Memory Bandwidth:   5.3 TB/s            3.35 TB/s
FP16 Performance:   1,307 TFLOPS        1,979 TFLOPS
FP8 Performance:    2,614 TFLOPS        3,958 TFLOPS
AI Accelerator:     MFMA               4th Gen Tensor Core
TDP:                750W                700W
Estimated Price:    ~$15,000-20,000     ~$30,000-40,000
Memory advantage:   ✅ 2.4x larger      -
Compute advantage:  -                   ✅ ~1.5x faster
Bandwidth advantage: ✅ 1.58x higher    -

NVIDIA Stack:                   AMD Stack:
┌──────────────────────┐        ┌──────────────────────┐
│   PyTorch / JAX      │        │   PyTorch / JAX      │
│   TensorFlow         │        │   TensorFlow         │
└──────────┬───────────┘        └──────────┬───────────┘
           ↓                               ↓
┌──────────────────────┐        ┌──────────────────────┐
│    CUDA Runtime      │        │   ROCm Runtime (HIP) │
└──────────┬───────────┘        └──────────┬───────────┘
           ↓                               ↓
┌──────────────────────┐        ┌──────────────────────┐
│   cuDNN / cuBLAS     │        │  MIOpen / rocBLAS    │
│   cuSPARSE           │        │  rocSPARSE           │
│   cuFFT              │        │  rocFFT              │
└──────────┬───────────┘        └──────────┬───────────┘
           ↓                               ↓
┌──────────────────────┐        ┌──────────────────────┐
│   NVCC compiler      │        │   hipcc compiler     │
│   PTX (IR)           │        │   GCN ISA / AMDGCN   │
└──────────┬───────────┘        └──────────┬───────────┘
           ↓                               ↓
┌──────────────────────┐        ┌──────────────────────┐
│    NVIDIA GPU        │        │     AMD GPU          │
└──────────────────────┘        └──────────────────────┘

// CUDA code (NVIDIA):
#include <cuda_runtime.h>

__global__ void vector_add(float* a, float* b, float* c, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

int main() {
    float *d_a, *d_b, *d_c;
    int n = 1024 * 1024;
    size_t size = n * sizeof(float);

    cudaMalloc(&d_a, size);
    cudaMalloc(&d_b, size);
    cudaMalloc(&d_c, size);

    cudaMemcpy(d_a, h_a, size, cudaMemcpyHostToDevice);
    vector_add<<<n/256, 256>>>(d_a, d_b, d_c, n);
    cudaMemcpy(h_c, d_c, size, cudaMemcpyDeviceToHost);

    cudaFree(d_a);
    return 0;
}

// HIP code (AMD): nearly identical to CUDA!
#include <hip/hip_runtime.h>

__global__ void vector_add(float* a, float* b, float* c, int n) {
    int idx = hipBlockIdx_x * hipBlockDim_x + hipThreadIdx_x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

int main() {
    float *d_a, *d_b, *d_c;
    int n = 1024 * 1024;
    size_t size = n * sizeof(float);

    hipMalloc(&d_a, size);           // cudaMalloc → hipMalloc
    hipMalloc(&d_b, size);
    hipMalloc(&d_c, size);

    hipMemcpy(d_a, h_a, size, hipMemcpyHostToDevice);  // prefix change only
    vector_add<<<n/256, 256>>>(d_a, d_b, d_c, n);      // same <<<>>> syntax
    hipMemcpy(h_c, d_c, size, hipMemcpyDeviceToHost);

    hipFree(d_a);
    return 0;
}

# Convert CUDA to HIP using HIPIFY
hipify-perl cuda_kernel.cu > hip_kernel.hip

# Or using clang-based HIPIFY
hipify-clang cuda_kernel.cu -- -I/usr/local/cuda/include

# Conversion rate: simple CUDA code converts 90%+ automatically
# Custom CUDA intrinsics require manual conversion

# Compile HIP code
hipcc hip_kernel.hip -o hip_kernel

// Same HIP code works on both platforms:
// AMD:    hipcc -arch=gfx942 kernel.hip   (MI300X)
// NVIDIA: hipcc --platform=nvidia kernel.hip  (compiles via CUDA)

// Platform detection
#ifdef __HIP_PLATFORM_AMD__
    // AMD-specific optimization
    __builtin_amdgcn_s_sleep(1);
#elif defined(__HIP_PLATFORM_NVIDIA__)
    // NVIDIA-specific code path
    __nanosleep(1000);
#endif

NVIDIA H100 SM (Streaming Multiprocessor):
┌───────────────────────────────────────────────┐
│  128 CUDA Cores (FP32 compute units)          │
│  64 FP64 Cores                                │
│  4 Tensor Cores (4th gen, FP8/FP16/BF16/INT8) │
│  8 LD/ST units (Load/Store)                   │
│  228KB L1 cache / Shared Memory (configurable)│
│  65,536 x 32-bit registers                    │
│                                               │
│  Single SM FP16 peak: ~60 TFLOPS             │
└───────────────────────────────────────────────┘

AMD MI300X CU (Compute Unit):
┌───────────────────────────────────────────────┐
│  64 Stream Processors (SIMD vector units)     │
│  64 FP64 units                                │
│  4 Matrix Cores (MFMA: Matrix Fused Multiply-Add) │
│  16 LD/ST units                               │
│  64KB L1 cache                                │
│  32KB LDS (Local Data Share = Shared Memory)  │
│  65,536 x 32-bit registers                    │
│                                               │
│  MFMA is AMD's equivalent of Tensor Cores:   │
│  v_mfma_f32_16x16x16f16 (FP16 matrix multiply)│
└───────────────────────────────────────────────┘

NVIDIA Warp:
- 32 threads execute simultaneously (SIMT: Single Instruction, Multiple Threads)
- All threads execute the same instruction

AMD Wavefront (Wave64):
- 64 threads execute simultaneously
- Wave32 mode also supported on RDNA3 and MI300X
- Larger wavefront = higher data parallelism, but more inefficiency on branch divergence

// Check wavefront size in a HIP kernel
__global__ void check_wavefront() {
    // On AMD: warpSize = 64 (or 32 in Wave32 mode)
    // On NVIDIA: always warpSize = 32
    int lane = threadIdx.x % warpSize;
    printf("warpSize: %d, my lane: %d\n", warpSize, lane);
}

# Install PyTorch with ROCm support
# ROCm 6.0 (stable as of 2024)
pip install torch torchvision torchaudio \
  --index-url https://download.pytorch.org/whl/rocm6.0

# Verify ROCm environment
python -c "
import torch
print('PyTorch version:', torch.__version__)
print('ROCm available:', torch.cuda.is_available())  # True on AMD!
print('ROCm version:', torch.version.hip)
print('GPU count:', torch.cuda.device_count())
print('GPU name:', torch.cuda.get_device_name(0))
"
# Example output:
# PyTorch version: 2.3.0+rocm6.0
# ROCm available: True  ← True even on AMD!
# ROCm version: 6.0.0
# GPU count: 1
# GPU name: AMD Instinct MI300X

# Tensor operations on AMD GPU
import torch
device = torch.device("cuda")  # "cuda" works on AMD GPU via ROCm!

x = torch.randn(1000, 1000, device=device)
y = torch.randn(1000, 1000, device=device)
z = torch.matmul(x, y)
print(z.shape)  # torch.Size([1000, 1000])

import torch
import torch.nn.functional as F

# BF16 (BFloat16) support check
device = torch.device("cuda")
a = torch.randn(512, 512, dtype=torch.bfloat16, device=device)
b = torch.randn(512, 512, dtype=torch.bfloat16, device=device)
c = torch.matmul(a, b)  # Uses BF16 MFMA on MI300X

# Flash Attention on AMD (ROCm)
# Install flash-attn with ROCm build
# pip install flash-attn --no-build-isolation

from flash_attn import flash_attn_func

q = torch.randn(2, 512, 8, 64, dtype=torch.float16, device=device)
k = torch.randn(2, 512, 8, 64, dtype=torch.float16, device=device)
v = torch.randn(2, 512, 8, 64, dtype=torch.float16, device=device)

# Flash Attention supported on ROCm
out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)

# Install vLLM for ROCm
pip install vllm  # Automatically picks ROCm version in ROCm environment

# Or build from source
git clone https://github.com/vllm-project/vllm
cd vllm
pip install -e . --no-build-isolation  # Auto-detects ROCm environment

# Serve Llama 3.1 70B on MI300X
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 1 \
  --dtype float16 \
  --device cuda  # "cuda" works on AMD via ROCm!

# Larger model: 405B in INT8 on 4x MI300X (768GB total)
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.1-405B-Instruct \
  --tensor-parallel-size 4 \
  --quantization fp8 \
  --device cuda

# vLLM Python API on AMD
from vllm import LLM, SamplingParams

# Load 70B in FP16 on single MI300X (possible because of 192GB VRAM!)
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    dtype="float16",
    tensor_parallel_size=1,  # Single MI300X is enough
    gpu_memory_utilization=0.85,
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=512,
)

prompts = ["Explain the attention mechanism in transformers."]
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(output.outputs[0].text)

# Build llama.cpp with ROCm support
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
cmake -B build -DLLAMA_HIPBLAS=ON \
  -DCMAKE_HIP_ARCHITECTURES="gfx942"  # MI300X architecture
cmake --build build --config Release -j

# Run Llama 3.1 70B Q4_K_M on MI300X
./build/bin/llama-cli \
  -m models/llama-3.1-70b-q4_k_m.gguf \
  -p "Explain ROCm architecture" \
  -n 200 \
  --n-gpu-layers 999
# Expected: ~15-20 tok/s (Q4_K_M, single MI300X)
# With FP16 load: ~8-10 tok/s (better quality, larger model)

# On consumer AMD RX 7900 XTX:
./build/bin/llama-cli -m models/llama-3.1-8b-q4_k_m.gguf ...
# Expected: ~50-60 tok/s (RX 7900 XTX 24GB)

What a single MI300X 192GB enables:
┌─────────────────────────────────────────────────────────────────┐
│ Llama 3.1 70B FP16 load: ~140GB → fits! (52GB headroom)        │
│   → Single H100 SXM cannot hold this (80GB limit)              │
│                                                                 │
│ Llama 3.1 70B + long context KV cache:                         │
│   Model ~140GB + KV cache 32K context ~20GB = 160GB → fits!    │
│                                                                 │
│ Mixtral 8x7B MoE FP16: ~93GB → single card possible!          │
│                                                                 │
│ Experimental 100B+ research models: test without quantization  │
└─────────────────────────────────────────────────────────────────┘

# Theoretical maximum token generation speed (memory bandwidth limited)
# Batch size 1, single token generation

# 70B FP16 model = 140GB
# Each token generation must read all weights once

def estimate_max_toks_per_sec(memory_bw_gbps, model_size_gb):
    """Theoretical maximum under memory bandwidth limit"""
    return memory_bw_gbps / model_size_gb

# MI300X: 5,300 GB/s / 140 GB = ~38 tok/s (theoretical ceiling)
mi300x_est = estimate_max_toks_per_sec(5300, 140)
print(f"MI300X theoretical ceiling: {mi300x_est:.1f} tok/s")  # ~37.9

# H100 SXM: 3,350 GB/s / 140 GB = ~24 tok/s (theoretical ceiling)
# In practice H100 is faster due to higher compute efficiency
h100_est = estimate_max_toks_per_sec(3350, 140)
print(f"H100 theoretical ceiling: {h100_est:.1f} tok/s")   # ~23.9

CUDA Ecosystem (2024):                  ROCm Ecosystem (2024):
- PyTorch: full support ✅              - PyTorch: supported ✅ (stability improving)
- JAX: full support ✅                  - JAX: experimental support ⚠️
- TensorFlow: full support ✅           - TensorFlow: official support ✅
- FlashAttention: highly optimized ✅  - FlashAttention: supported but slower ⚠️
- cuDNN kernels: 15 years optimized ✅ - MIOpen: optimization ongoing ⚠️
- Triton: full support ✅               - Triton on ROCm: supported (perf gap) ⚠️
- BitsAndBytes: full support ✅        - BitsAndBytes on ROCm: supported ✅
- DeepSpeed: full support ✅           - DeepSpeed on ROCm: supported ✅
- vLLM: full support ✅                 - vLLM on ROCm: official support ✅ (2024~)

# ROCm 6.0 officially supported GPUs (as of 2024):
# ✅ AMD Instinct MI300X
# ✅ AMD Instinct MI250X
# ✅ AMD Instinct MI210
# ✅ AMD Instinct MI100
# ⚠️ RX 7900 XTX: unofficial support (many libraries work but no guarantee)
# ❌ RX 7800 XT and below: unstable support

# Force ROCm on consumer GPU
export HSA_OVERRIDE_GFX_VERSION=11.0.0  # For RX 7900 XTX
export ROCR_VISIBLE_DEVICES=0
rocminfo  # Print detected GPU info

# Use ROCm Docker image (most stable approach)
docker pull rocm/pytorch:rocm6.0_ubuntu22.04_py3.10_pytorch_2.1.1

# Run container with GPU access
docker run -it \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --security-opt seccomp=unconfined \
  --cap-add=SYS_PTRACE \
  -v $(pwd):/workspace \
  rocm/pytorch:rocm6.0_ubuntu22.04_py3.10_pytorch_2.1.1 \
  /bin/bash

# Inside container, verify
python -c "import torch; print(torch.cuda.is_available(), torch.version.hip)"

# Install ROCm 6.0 on Ubuntu 22.04 LTS
# Step 1: Add AMD ROCm package repository
wget https://repo.radeon.com/amdgpu-install/6.0/ubuntu/jammy/amdgpu-install_6.0.60000-1_all.deb
sudo dpkg -i amdgpu-install_6.0.60000-1_all.deb

# Step 2: Install ROCm
sudo amdgpu-install --usecase=hiplibsdk,rocm,ml

# Step 3: Add user to render and video groups
sudo usermod -aG render,video $USER

# Step 4: Re-login and verify
rocminfo | grep "Name:"
# Agent 2: gfx942  (MI300X)

# Step 5: Install PyTorch ROCm
pip install torch torchvision torchaudio \
  --index-url https://download.pytorch.org/whl/rocm6.0