Skip to content

[ROCm] Support the INT4 QAT kernel on ROCm - #2274

Merged
zhuzilin merged 1 commit into
THUDM:mainfrom
LZ-QWQ:rocm-int4-qat
Aug 16, 2026
Merged

[ROCm] Support the INT4 QAT kernel on ROCm#2274
zhuzilin merged 1 commit into
THUDM:mainfrom
LZ-QWQ:rocm-int4-qat

Conversation

@LZ-QWQ

@LZ-QWQ LZ-QWQ commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

The int4_qat kernel does not build on a ROCm PyTorch install, so the INT4 QAT path is unusable on AMD GPUs. This fixes the two blockers (a HIP warp-shuffle mask and a const_data_ptr symbol mangling mismatch under hipcc).

Build

  • CUDA: unchanged. Every modification is behind #if defined(__HIP_PLATFORM_AMD__) or torch.version.hip is not None, so the CUDA compile path is byte-for-byte identical.
  • ROCm: pip install . just works. The gfx target comes from PYTORCH_ROCM_ARCH (defaults to gfx950), since hipcc does not accept the nvcc-only flags.

Verification (MI355X / ROCm 7.2, CUDA reference from an RTX 4090)

  • 144 cases covering the block shapes, tensor shapes and dtypes the kernel dispatches on: outputs are bit-exact against CUDA.
  • E2E tools/convert_hf_to_int4_direct.py on Qwen3-0.6B: 196 layers converted, dequantization error within the expected range for group-wise INT4 (worst 0.1051).

@zhuzilin
zhuzilin merged commit 876cd89 into THUDM:main Aug 16, 2026
61 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants