[ROCm] Preserve float bits in SIMT DPP reduction

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

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