[ROCm] Add ROCm support and norm op.

[WhatGhost]

Add AMD GPU (ROCm) support: build path + first norm operator

Background

HPC-Ops is currently a CUDA-only operator library targeting NVIDIA SM90. This PR starts bringing it to AMD GPUs (ROCm).

The goal is incremental: this first PR lands the compilation path for the ROCm backend plus one operator (fused_rmsnorm_with_scale) as a working end-to-end example. More operators will be ported in follow-up PRs, and the performance of the ported operators will continue to be optimized over time.

Target hardware for this PR: AMD MI350X (gfx950 / CDNA4).

1. HIP kernel with simple adaptation and optimization.
The original CUDA norm operator was converted to HIP and placed under a new src/amd/ tree. Made some changes to adapt it to HIP。It also includes a simple optimization

2. ROCm build path for AMD GPUs.
CMakeLists.txt gains an opt-in USE_ROCM branch that compiles only the ported src/amd/** sources with hipcc, leaving the CUDA path byte-for-byte unchanged when USE_ROCM is off，and hpc/__init__.py tolerates operators not yet built on ROCm (only the norm op is compiled for now) so importing hpc works.

Testing

The norm operator was validated for correctness (60/60 parametrized pytest cases pass on gfx950) and benchmarked against four references:

• HPC — this operator (fused RMSNorm + scale + fp8 quant)
• Triton — a Triton RMSNorm + scale + fp8 quant kernel
• aiter — AMD's aiter rms_norm (plain RMSNorm, bf16)
• torch — a native PyTorch RMSNorm (bf16)

Results (gfx950 / MI350X, latency μs, lower is better)

Measured with CUDA events, median-of-100 × 3 rounds (min), warmup 20, pinned to an idle GPU.

Columns marked +fp8 (HPC and Triton) perform RMSNorm + scale + fp8 quantization — i.e. they emit an fp8 output and therefore do strictly more work. The aiter and torch columns do a plain RMSNorm with bf16 output. The two groups are not the same workload; the +fp8 group is doing extra quantization on top of the normalization.

hidden = 5120

batch	HPC (+fp8)	Triton (+fp8)	aiter	torch
16	6.54	13.62	10.64	39.42
256	7.80	14.92	13.10	42.52
1024	8.14	13.84	12.14	56.74
4096	14.40	13.52	17.12	129.46
8192	23.56	21.22	28.90	280.82

hidden = 4096

batch	HPC (+fp8)	Triton (+fp8)	aiter	torch
16	7.00	13.46	12.04	41.44
256	7.16	13.64	12.08	41.08
1024	7.72	13.14	11.48	50.84
4096	11.60	13.92	14.44	109.02
8192	19.04	16.84	23.56	211.56

hidden = 320

batch	HPC (+fp8)	Triton (+fp8)	aiter	torch
16	7.28	13.48	11.84	40.72
256	7.16	13.52	11.88	41.42
1024	7.20	13.44	12.08	41.62
4096	7.44	12.92	12.18	41.12
8192	7.76	13.60	11.44	47.68

Summary: HPC is fastest across small/medium batches and all hidden=320 sizes (while also doing the extra fp8 quant), and beats aiter/torch everywhere. On large batches with large hidden the picture is mixed: at batch 4096 it is competitive (faster than Triton at hidden=4096, within ~7% at hidden=5120), while at batch 8192 it trails Triton by ~11-13%; closing that batch-8192 gap is future optimization work.

Benchmark script

"""hpc vs triton / aiter / torch。"""
import sys
from statistics import median

import torch
import triton
import triton.language as tl
import hpc  


@triton.jit
def _rmsnorm_fp8(x_ptr, w_ptr, o_ptr, inv_scale, eps, H: tl.constexpr, BLK: tl.constexpr):
    row = tl.program_id(0)
    off = tl.arange(0, BLK)
    mask = off < H
    x = tl.load(x_ptr + row * H + off, mask=mask, other=0.0).to(tl.float32)
    w = tl.load(w_ptr + off, mask=mask, other=0.0).to(tl.float32)
    rms = tl.rsqrt(tl.sum(x * x) / H + eps)
    y = x * rms * w * inv_scale
    tl.store(o_ptr + row * H + off, y.to(tl.float8e4nv), mask=mask)


def triton_fp8(x, w, inv_scale, eps, H, BLK):
    o = torch.empty_like(x, dtype=torch.float8_e4m3fn)
    _rmsnorm_fp8[(x.shape[0],)](x, w, o, inv_scale, eps, H, BLK)
    return o


def torch_bf16(x, w, eps):
    rms = torch.rsqrt(torch.mean(x.float() ** 2, -1, keepdim=True) + eps)
    return (x * rms * w.float()).to(torch.bfloat16)


def bench_us(fn, warmup=20, iters=100, rounds=3):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    meds = []
    for _ in range(rounds):
        evs = [(torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True))
               for _ in range(iters)]
        for s, e in evs:
            s.record(); fn(); e.record()
        torch.cuda.synchronize()
        meds.append(median([s.elapsed_time(e) * 1000.0 for s, e in evs]))
    return min(meds)


def main():
    from aiter.ops.rmsnorm import rms_norm as aiter_rms  # noqa
    dev = "cuda"
    eps = 1e-6
    scale = 2.5
    st = torch.tensor([scale], dtype=torch.float32, device=dev)
    inv = 1.0 / scale
    print(f"GPU {torch.cuda.get_device_name(0)}  |  hpc from {hpc.__file__}")
    print(f"{'hidden':>6} {'batch':>6} | {'HPC(+fp8)':>10} {'Tri(+fp8)':>10} | "
          f"{'aiter':>8} {'torch':>8}")
    print("-" * 62)
    for hidden in [5120, 4096, 320]:
        BLK = triton.next_power_of_2(hidden)
        w2 = torch.rand((1, hidden), dtype=torch.bfloat16, device=dev).contiguous()
        w1 = w2.reshape(-1).contiguous()
        for bs in [16, 256, 1024, 4096, 8192, 16384]:
            x = torch.randn(bs, hidden, dtype=torch.bfloat16, device=dev)
            t_hpc = bench_us(lambda: torch.ops.hpc.fused_rmsnorm_with_scale(x, w2, st, eps, False))
            t_tri = bench_us(lambda: triton_fp8(x, w1, inv, eps, hidden, BLK))
            t_ai = bench_us(lambda: aiter_rms(x, w1, eps))
            t_to = bench_us(lambda: torch_bf16(x, w2, eps))
            print(f"{hidden:>6} {bs:>6} | {t_hpc:>10.2f} {t_tri:>10.2f} | {t_ai:>8.2f} {t_to:>8.2f}")


if __name__ == "__main__":
    main()

[WhatGhost]

@reed-lau Hi, could you please help to review this PR. Thanks!