Skip to content

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

Closed
LZ-QWQ wants to merge 1 commit into
THUDM:mainfrom
LZ-QWQ:fix/int4-qat-rocm
Closed

[ROCm] Support the INT4 QAT kernel on ROCm#2273
LZ-QWQ wants to merge 1 commit into
THUDM:mainfrom
LZ-QWQ:fix/int4-qat-rocm

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, worst relative error 0.1051.

@LZ-QWQ
LZ-QWQ force-pushed the fix/int4-qat-rocm branch from 7e64607 to 3197e26 Compare August 15, 2026 07:30
@LZ-QWQ LZ-QWQ closed this Aug 15, 2026
@LZ-QWQ
LZ-QWQ deleted the fix/int4-qat-rocm branch August 15, 2026 07:32
@LZ-QWQ LZ-QWQ changed the title fix(int4_qat): support building the INT4 QAT kernel on ROCm [ROCm] Support the INT4 QAT kernel on ROCm Aug 15, 2026
@LZ-QWQ

LZ-QWQ commented Aug 15, 2026

Copy link
Copy Markdown
Contributor Author

Closed by an accidental branch rename. Superseded by #2274 (same commit, no content change).

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.

1 participant