From 4d25e61602531293ede99ef0e22c9616efb03769 Mon Sep 17 00:00:00 2001 From: AMD ROCm Contributor <91481003+liminfei-amd@users.noreply.github.com> Date: Tue, 15 Sep 2026 17:51:34 +0800 Subject: [PATCH 1/2] [ROCm] Preserve float bits in SIMT DPP reduction The AMD mov_dpp builtin operates on integer bits. Explicitly bitcast the floating-point accumulator before and after each lane move so older ROCm compilers do not emit numeric float-to-integer conversions. Add a gfx11 regression test covering the fused Linear4bit path for fp16, bf16, and fp32. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- csrc/gemm_4bit_simt.cu | 9 +++++---- tests/test_functional.py | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/csrc/gemm_4bit_simt.cu b/csrc/gemm_4bit_simt.cu index d0f5354fb..0edc6621c 100644 --- a/csrc/gemm_4bit_simt.cu +++ b/csrc/gemm_4bit_simt.cu @@ -77,10 +77,11 @@ __device__ __forceinline__ float simt_fma_f32(float acc, float a, float b) { // CUDA/others: standard __shfl_down_sync tree. __device__ __forceinline__ float simt_warp_reduce_sum(float v) { #if BNB_HIP - v += __builtin_amdgcn_mov_dpp(v, 0x108, 0xf, 0xf, 1); // row_shr:8 - v += __builtin_amdgcn_mov_dpp(v, 0x104, 0xf, 0xf, 1); // row_shr:4 - v += __builtin_amdgcn_mov_dpp(v, 0x102, 0xf, 0xf, 1); // row_shr:2 - v += __builtin_amdgcn_mov_dpp(v, 0x101, 0xf, 0xf, 1); // row_shr:1 + // mov_dpp is integer-typed; bitcast to preserve the float payload across lanes. + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x108, 0xf, 0xf, 1)); // row_shr:8 + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x104, 0xf, 0xf, 1)); // row_shr:4 + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x102, 0xf, 0xf, 1)); // row_shr:2 + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x101, 0xf, 0xf, 1)); // row_shr:1 return v + __shfl_xor(v, 16, 32); #else #pragma unroll diff --git a/tests/test_functional.py b/tests/test_functional.py index e4cd6a128..8cffdec1b 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -1012,6 +1012,43 @@ def test_matmul_4bit(self, MNK, dtype, blocksize, quant_type, compress_statistic threshold = 0.16 assert mean_err < threshold + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) + def test_matmul_4bit_rocm_gfx11_dpp_reduction(self, device, dtype): + torch_device = torch.device(device) + if torch_device.type != "cuda" or torch.version.hip is None: + pytest.skip("ROCm is required") + from bitsandbytes.backends.cuda.ops import _gemm_4bit_use_custom_rocm, _rocm_gfx_arch + + device_index = torch_device.index + if device_index is None: + device_index = torch.cuda.current_device() + arch = _rocm_gfx_arch(device_index) + if not arch.startswith("gfx11"): + pytest.skip("gfx11 is required") + + torch.manual_seed(0) + M, N, K = 8, 256, 512 + A = torch.randn(M, K, dtype=dtype, device=device) + B = torch.randn(N, K, dtype=dtype, device=device) + + assert _gemm_4bit_use_custom_rocm(device_index, dtype, M, N, K) + linear = bnb.nn.Linear4bit( + K, + N, + bias=False, + quant_type="nf4", + compute_dtype=dtype, + device="meta", + ) + linear.weight = bnb.nn.Params4bit(B, quant_type="nf4", requires_grad=False) + linear = linear.to(device) + dequantized_ref = A @ F.dequantize_4bit(linear.weight.data, linear.weight.quant_state).t() + out = linear(A) + relative_l1 = (out.float() - dequantized_ref.float()).abs().mean() / dequantized_ref.float().abs().mean() + + assert relative_l1.item() < 0.02 + @pytest.mark.parametrize("device", get_available_devices()) def test_matmul_4bit_weight_orientation(self, device): N, K = 256, 128 From eb5f933fec6b702d37b4e2f16db419e7c720d282 Mon Sep 17 00:00:00 2001 From: Laela Zorana Date: Fri, 2 Oct 2026 09:49:30 -0500 Subject: [PATCH 2/2] Keep float bits across the ROCm SIMT DPP reduction mov_dpp takes an int before clang 20, so the ROCm 6.4.4 build converted each partial sum to an integer and back. Pass the bits instead, fix the row_shl labels, and test the fused path on every ROCm arch that selects it. [liminfei-amd: applied on top of #2089, which already passes the float bits; this commit carries the comment, row_shl label, and test changes.] (cherry picked from commit 0b5a20e8ec051eda29a0629cbeb929655d7035df) --- csrc/gemm_4bit_simt.cu | 11 ++++++----- tests/test_functional.py | 36 +++++++++++++----------------------- 2 files changed, 19 insertions(+), 28 deletions(-) diff --git a/csrc/gemm_4bit_simt.cu b/csrc/gemm_4bit_simt.cu index 0edc6621c..a7bb5a167 100644 --- a/csrc/gemm_4bit_simt.cu +++ b/csrc/gemm_4bit_simt.cu @@ -77,11 +77,12 @@ __device__ __forceinline__ float simt_fma_f32(float acc, float a, float b) { // CUDA/others: standard __shfl_down_sync tree. __device__ __forceinline__ float simt_warp_reduce_sum(float v) { #if BNB_HIP - // mov_dpp is integer-typed; bitcast to preserve the float payload across lanes. - v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x108, 0xf, 0xf, 1)); // row_shr:8 - v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x104, 0xf, 0xf, 1)); // row_shr:4 - v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x102, 0xf, 0xf, 1)); // row_shr:2 - v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x101, 0xf, 0xf, 1)); // row_shr:1 + // mov_dpp takes an int before clang 20, so pass the float bits, not the value. + // 0x101 to 0x10f are row_shl, so lane 0 gathers its 16 lane row. + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x108, 0xf, 0xf, 1)); // row_shl:8 + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x104, 0xf, 0xf, 1)); // row_shl:4 + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x102, 0xf, 0xf, 1)); // row_shl:2 + v += __uint_as_float(__builtin_amdgcn_mov_dpp(__float_as_uint(v), 0x101, 0xf, 0xf, 1)); // row_shl:1 return v + __shfl_xor(v, 16, 32); #else #pragma unroll diff --git a/tests/test_functional.py b/tests/test_functional.py index 8cffdec1b..c7030cceb 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -1014,39 +1014,29 @@ def test_matmul_4bit(self, MNK, dtype, blocksize, quant_type, compress_statistic @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) - def test_matmul_4bit_rocm_gfx11_dpp_reduction(self, device, dtype): + @pytest.mark.parametrize("M", [1, 2, 8], ids=id_formatter("M")) + def test_matmul_4bit_rocm_simt_reduction(self, device, dtype, M): + # The fused SIMT kernel sums each warp with DPP moves on every ROCm arch. + # Compare against the same dequantized weight so only the kernel is measured. torch_device = torch.device(device) if torch_device.type != "cuda" or torch.version.hip is None: pytest.skip("ROCm is required") - from bitsandbytes.backends.cuda.ops import _gemm_4bit_use_custom_rocm, _rocm_gfx_arch + from bitsandbytes.backends.cuda.ops import _gemm_4bit_use_custom_rocm - device_index = torch_device.index - if device_index is None: - device_index = torch.cuda.current_device() - arch = _rocm_gfx_arch(device_index) - if not arch.startswith("gfx11"): - pytest.skip("gfx11 is required") + index = torch_device.index if torch_device.index is not None else torch.cuda.current_device() + N, K = 256, 512 + if not _gemm_4bit_use_custom_rocm(index, dtype, M, N, K): + pytest.skip("fused SIMT path is not selected for this shape on this arch") torch.manual_seed(0) - M, N, K = 8, 256, 512 A = torch.randn(M, K, dtype=dtype, device=device) B = torch.randn(N, K, dtype=dtype, device=device) - - assert _gemm_4bit_use_custom_rocm(device_index, dtype, M, N, K) - linear = bnb.nn.Linear4bit( - K, - N, - bias=False, - quant_type="nf4", - compute_dtype=dtype, - device="meta", - ) + linear = bnb.nn.Linear4bit(K, N, bias=False, quant_type="nf4", compute_dtype=dtype, device="meta") linear.weight = bnb.nn.Params4bit(B, quant_type="nf4", requires_grad=False) linear = linear.to(device) - dequantized_ref = A @ F.dequantize_4bit(linear.weight.data, linear.weight.quant_state).t() - out = linear(A) - relative_l1 = (out.float() - dequantized_ref.float()).abs().mean() / dequantized_ref.float().abs().mean() - + ref = A.float() @ F.dequantize_4bit(linear.weight.data, linear.weight.quant_state).float().t() + out = linear(A).float() + relative_l1 = (out - ref).abs().mean() / ref.abs().mean() assert relative_l1.item() < 0.02 @pytest.mark.parametrize("device", get_available_devices())