diff --git a/csrc/gemm_4bit_simt.cu b/csrc/gemm_4bit_simt.cu index d0f5354fb..a7bb5a167 100644 --- a/csrc/gemm_4bit_simt.cu +++ b/csrc/gemm_4bit_simt.cu @@ -77,10 +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 - 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 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 e4cd6a128..c7030cceb 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -1012,6 +1012,33 @@ 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) + @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 + + 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) + A = torch.randn(M, K, dtype=dtype, device=device) + B = torch.randn(N, K, dtype=dtype, device=device) + 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) + 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()) def test_matmul_4bit_weight_orientation(self, device): N, K = 256, 128