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] [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