hip-quant 0.4.6

HIP/ROCm 7.1 Quantization library for AMD GPUs

🚀 hip-quant

Blazing Fast On-Device Tensor Quantization for AMD GPUs

hip-quant is a standalone Python library and highly optimized HIP C++ backend that quantizes tensors directly on AMD GPUs with no CPU round-trips. The offline GGUF path consumes float32; the PyTorch FP8 training extension accepts float32, float16, and bfloat16 tensors.

It ships two independent APIs that can be used together or separately:

API	Purpose	Requires
NumPy / ctypes (offline)	Offline GGUF-format quantization via hip_quantize.dll	ROCm 7.1, numpy
PyTorch extension (training)	GPU-resident FP8 training ops with full autograd	PyTorch 2.x + ROCm, built _C extension

Hardware Status

Runtime validation is currently on RDNA4. The PyTorch FP8 WMMA kernels target gfx1200 and gfx1201; gfx1200 is treated as the cut-down gfx1201 die with the same relevant FP8 WMMA capabilities.

CDNA support is included for the offline NumPy/DLL quantization path and compatibility tooling. .\build.ps1 -CDNA cross-compiles the portable quantization kernels for gfx90a and gfx942 on the local ROCm 7.1 Windows toolchain. Device reports still classify gfx940 and gfx941 correctly when encountered, but this Windows hipcc does not accept them as offload targets. The gfx12 WMMA FP8 GEMM test is intentionally disabled on CDNA; CDNA GEMM needs an MFMA path.

⚠️ ROCm 7.1 WMMA Driver Crash (Windows gfx12)

The PyTorch FP8 training API (Fp8Linear, Fp8ScaledLinear, Fp8ShadowLinear) uses __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12 for the fused GEMM forward/backward pass. On ROCm 7.1 with Windows gfx1201, these WMMA intrinsics can trigger a GPU TDR (driver timeout) that corrupts GPU memory and may restart the PC.

Root cause: The ROCm 7.1 HIP runtime has a stability issue with gfx12 WMMA instructions on RDNA4. The kernel launch succeeds but the GPU can hang asynchronously, causing subsequent tensor.item() calls to read back corrupted memory — typically manifesting as a ZeroDivisionError at torch_api.py:495 (1.0 / weight_inv_scale with a zeroed GPU value).

Fix: Use the PyTorch venv's ROCm 7.2.1+ runtime, which ships with torch 2.9.1+rocm7.2.1. The system-installed ROCm 7.1 (C:\Program Files\AMD\ROCm\7.1\bin) is shadowed at runtime by the venv's ROCm 7.2 DLLs, and WMMA is stable on 7.2.

Validated local test system:

• GPU: AMD Radeon RX 9070 XT, gfx1201, 16 GB VRAM
• CPU: AMD Ryzen 7 7800X3D, 8 cores / 16 threads
• RAM: 32 GB system memory
• OS/toolchain: Windows, Visual Studio 2022 Build Tools, ROCm installed at C:\Program Files\AMD\ROCm\7.1
• PyTorch venv: C:\venvs\medusa_rocm\Scripts\python.exe
• PyTorch: 2.9.1+rocm7.2.1, HIP runtime: 7.2.53211-158bd99533
• FP8 training verified stable: 10-step Fp8ScaledLinear loop, WMMA stress test (20x 256×256×256), full forward+backward on all module types

Note: The offline NumPy/DLL quantization path (hip_quantize.dll) does not use WMMA and is unaffected. It works with both ROCm 7.1 and 7.2.

---

⚡ Supported Quantization Formats

🔢 Standard & K-Quants (offline API)

• Legacy: Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q8_1
• K-Quants: Q2_K, Q3_K, Q4_K, Q5_K, Q6_K

🧠 I-Quants (Importance Matrix)

Non-linear quants that preserve quality at extreme low bits:

• IQ1_S, IQ2_XXS, IQ2_XS, IQ3_XXS, IQ3_S, IQ4_NL, IQ4_XS

⚖️ Ternary Quants

For models trained to be ternary (BitNet, TriLM):

• TQ1_0 (1.69 bpw), TQ2_0 (2.06 bpw)

🧪 FP8 Formats (both APIs)

Format	Layout	Use case
F8_E4M3	1s·4e·3m, bias=7, max=±448, NaN only	Forward activations & weights
F8_E5M2	1s·5e·2m, bias=15, max=±57344, ±Inf+NaN	Backward gradients

Both use OCP standard semantics with round-to-nearest-even. Math is validated against ml_dtypes reference — 90/90 test cases pass.

---

🛠️ Build

Offline DLL (NumPy API)

Requires hipcc at C:\Program Files\AMD\ROCm\7.1\bin\hipcc.exe:

.\build.ps1

# Include CDNA portable quantization targets as well as RDNA4
.\build.ps1 -CDNA

# Custom target set
.\build.ps1 -Arch "gfx942,gfx1200,gfx1201"

PyTorch Extension (_C)

Requires PyTorch with ROCm support (torch 2.x+rocm):

& "C:\venvs\medusa_rocm\Scripts\python.exe" setup_torch.py build_ext --inplace

To build a PyPI wheel that includes the compiled _C.pyd extension, build with the extension flag from the ROCm/PyTorch environment:

$env:HIP_QUANT_BUILD_TORCH_EXT = "1"
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m build --wheel --no-isolation

Without HIP_QUANT_BUILD_TORCH_EXT=1, python -m build intentionally creates a source-only py3-none-any wheel that still requires the local extension build.

---

📦 Installation

# Binary wheel with fused PyTorch kernels
$env:HIP_QUANT_BUILD_TORCH_EXT = "1"
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m build --wheel --no-isolation
pip install dist/hip_quant-0.4.5-cp312-cp312-win_amd64.whl

# With PyTorch optional dependency declared
pip install "hip-quant[torch]"

---

🐍 Usage

Offline NumPy API

import numpy as np
from hip_quant import quantize

weights = np.random.randn(4096, 4096).astype(np.float32)

# Quantize directly to Q4_K on the GPU — byte-exact match to llama.cpp
q4k_bytes = quantize(weights, type_num=12)  # 12 = Q4_K

FP8 (offline)

from hip_quant import GGML_TYPE, get_hip_quant

hq = get_hip_quant()
x    = np.random.randn(4096, 4096).astype(np.float32)
grad = (np.random.randn(4096, 4096) * 128).astype(np.float32)

x_e4m3    = hq.quantize_numpy(x,    GGML_TYPE["F8_E4M3"])  # forward
grad_e5m2 = hq.quantize_numpy(grad, GGML_TYPE["F8_E5M2"])  # backward

CLI

hip-quant --help
python -m hip_quant --help

---

PyTorch Training API

Requires: python setup_torch.py build_ext --inplace first.

Element-wise FP8 quant / dequant (Phase 1 & 2)

import torch
from hip_quant.torch_api import quantize_e4m3, dequantize_e4m3
from hip_quant.torch_api import quantize_e5m2, dequantize_e5m2

x = torch.randn(1024, 1024, device="cuda")  # stays on GPU the whole time

x_fp8  = quantize_e4m3(x)          # torch.uint8, same shape, same device
x_back = dequantize_e4m3(x_fp8)    # torch.float32, no CPU transfer

Fake-FP8 Linear (autograd-safe, Phase 3)

Fp8LinearFunction uses E4M3 for forward activations/weights and E5M2 for backward gradients. It accepts torch.float32, torch.float16, and torch.bfloat16 inputs/weights. It also implements Activation Compression, saving uint8 tensors in the autograd graph to cut activation VRAM by 4× versus FP32, and 2× versus FP16/BF16.

BF16/FP16 support applies to:

• quantize_e4m3() and quantize_e5m2() inputs
• Fp8LinearFunction forward/backward
• Fp8Linear, Fp8ScaledLinear, and Fp8ShadowLinear module parameters and gradients
• Fp8ShadowLinear master weights, so user-selected BF16/FP16 master weights reduce persistent parameter and gradient VRAM versus FP32

from hip_quant.torch_api import convert_to_fp8, Adafactor

# Drop-in replacement for all nn.Linear layers in a model
model = MySmallLM(...)

# shadow=True: replaces nn.Linear with Fp8ShadowLinear
# Weights are stored as uint8 in memory, forward pass decompresses on the fly
# Cuts weight VRAM by 4×
convert_to_fp8(model, shadow=True, skip_names={"lm_head"})
model.cuda()

# Adafactor optimizer: adaptive learning rates with sublinear memory cost
# Cuts optimizer state VRAM by ~1000× compared to AdamW
opt = Adafactor(model.parameters(), relative_step=True)

Combined VRAM savings for a 500M-param LLM: Before: ~7.6 GB (Weights 2GB, Acts 1.6GB, AdamW 4GB) After: ~0.9 GB (Weights 0.5GB, Acts 0.4GB, Adafactor 4MB)

Direct autograd.Function

from hip_quant.torch_api import Fp8LinearFunction

out = Fp8LinearFunction.apply(input, weight, bias)  # bias optional

Fused FP8 Linear (gfx12 WMMA kernels)

from hip_quant import (
    fp8_linear_forward,
    fp8_linear_forward_scaled,
    fp8_linear_forward_fp8_weight,
    fp8_linear_backward_input,
    fp8_linear_backward_input_scaled,
    fp8_linear_backward_weight,
    fp8_linear_backward_weight_scaled,
)

# [M,K] @ [N,K].T = [M,N]
# forward: E4M3 x E4M3 WMMA, backward: E5M2/BF8 x E5M2/BF8 WMMA
out        = fp8_linear_forward(input, weight, bias=None)
grad_in    = fp8_linear_backward_input(grad_output, weight)
grad_wt    = fp8_linear_backward_weight(grad_output, input)

# Scaled path used by Fp8ScaledLinear and Fp8ShadowLinear
out_scaled = fp8_linear_forward_scaled(input, weight, bias, input_scale, weight_scale)
grad_in_s  = fp8_linear_backward_input_scaled(grad_output, weight, weight_scale)
grad_wt_s  = fp8_linear_backward_weight_scaled(grad_output, input, input_scale)

These functions are also used by Fp8Linear, Fp8ScaledLinear, and Fp8ShadowLinear after the extension is built.

Scale / amax tracking (Phase 4 scaffold)

from hip_quant.torch_api import Fp8TensorMeta

meta = Fp8TensorMeta(history_len=16, device="cuda")
meta.update(x)                   # records amax, updates scale/inv_scale

x_fp8  = meta.quantize_e4m3(x)  # scaled, then quantized
x_back = meta.dequantize_e4m3(x_fp8)  # dequantized, then rescaled

---

🔒 Memory Safety

All PyTorch extension functions are guarded against:

• Non-CUDA tensors (TORCH_CHECK(is_cuda))
• Non-contiguous layout (TORCH_CHECK(is_contiguous))
• Wrong dtype (float32 / float16 / bfloat16 for floating inputs, uint8 for FP8 buffers)
• Dimension mismatch for GEMM
• int64 → int narrowing — explicit checked_int() with TORCH_CHECK
• Hardware grid limit — gridDim.y ≤ 65535 validated before launch
• Cross-device pointers — input.device() == weight.device() checked
• Empty tensors — numel == 0 early-return before dim3(0) (UB in HIP)

---

🧪 Running Tests

Math tests (no GPU required)

python tests/torch/test_math_fp8.py
# 90/90 pass — validated against ml_dtypes reference

PyTorch GPU tests

# Build extension first
& "C:\venvs\medusa_rocm\Scripts\python.exe" setup_torch.py build_ext --inplace

& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/torch/test_fp8.py -v

Full Pipeline Tests (CPU Mock)

python tests/test_pipeline.py -v
# Tests full integration of Adafactor, Shadow Linear, and activation compression
# Pure-Python mock, 100% test coverage for the API

Compatibility Tests (CPU + DLL)

& "C:\venvs\medusa_rocm\Scripts\python.exe" -m pytest tests/test_compat.py -v

# Device/compat reports
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m hip_quant --info
& "C:\venvs\medusa_rocm\Scripts\python.exe" -m hip_quant --compat

---

🗂️ Project Structure

hip_quant/
├── __init__.py              # NumPy / ctypes offline API
├── __main__.py              # CLI entry point
├── torch_api.py             # PyTorch FP8 training API (Phases 1–4)
├── device_info.py           # GPU/DLL compatibility probe helpers
├── cdna_compat.py           # CDNA feature table, build configs, CPU refs
├── setup_torch.py           # PyTorch C++ extension build script
├── build.ps1                # DLL build script (hipcc)
├── hip_quantize.cpp         # Offline quantization kernels (DLL source)
├── hip_quant_util.h         # Shared FP8 / FP16 device helpers
├── hip_quant_types.h        # GGML block type definitions
├── kernels/                 # Per-format offline HIP kernels (.cu)
├── torch_ext/               # PyTorch extension source
│   ├── pytorch_bindings.cpp # C++ bindings (TORCH_CHECK, pybind11)
│   ├── fp8_quant_kernels.hip# Element-wise quant/dequant kernels
│   └── fp8_linear_kernels.hip# Tiled FP8 GEMM kernels
└── tests/torch/             # GPU test suite (pytest)

---

📋 Architecture Notes

• RDNA4 PyTorch target — FP8 WMMA extension kernels are compiled with --offload-arch=gfx1200 and --offload-arch=gfx1201
• CDNA offline target — build.ps1 -CDNA compiles the portable DLL quantization kernels for gfx90a, gfx942, gfx1200, and gfx1201
• Current validation scope — runtime-tested locally on gfx1201 RX 9070 XT; gfx1200 and CDNA code objects are build-validated and need separate hardware runtime validation
• BF16/FP16 PyTorch support — FP8 quantization and linear kernels accept FP32, FP16, and BF16 tensors, accumulating in FP32 registers and storing results in the input/master dtype
• No CPU transfers — the PyTorch extension operates exclusively on device pointers obtained from tensor.data_ptr() and uses at::cuda::getCurrentCUDAStream()
• Phase 4 GEMM is a correctness-first tiled stub. Replace the inner loop with a rocBLASLt FP8 GEMM path for production throughput once validated on your ROCm stack
• Offline API unchanged — the NumPy/ctypes path is untouched; both APIs coexist cleanly