Warn once, and only when it matters, for unaligned gemm_4bit

bitsandbytes/backends/cuda/ops.py

@@ -3,14 +3,14 @@
 import functools
 from math import prod
 from typing import Optional
-from warnings import warn
 
 import torch
 
 from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, get_ptr
 
 from ..._ops import register_kernel
 from ...cextension import lib
+from ..utils import _warn_gemm_4bit_unaligned
 
 
 def _setup_ctypes(names, argtypes, restype=None):
@@ -943,15 +943,13 @@ def _(
     # fallback.
     if M > _gemm_4bit_custom_max_m:
         use_custom = False
-    elif K % blocksize != 0:
-        warn(
-            f"inner dimension ({K}) is not aligned for fast kernel "
-            f"with blocksize={blocksize}, falling back to slower implementation.",
-            UserWarning,
-        )
-        use_custom = False
     else:
         use_custom = _gemm_4bit_use_custom_fn(A.device.index, A.dtype, M, N, K)
+        if use_custom and K % blocksize != 0:
+            # Only warn when misalignment is what costs us the fused kernel; when the
+            # heuristic picks the fallback anyway (larger M, e.g. training), it doesn't.
+            _warn_gemm_4bit_unaligned(K, blocksize)
+            use_custom = False
 
     if not use_custom:
         return _dequant_linear_fallback(

bitsandbytes/backends/utils.py

@@ -1,4 +1,6 @@
+import functools
 from importlib.metadata import metadata
+from warnings import warn
 
 from packaging import version
 import torch
@@ -75,6 +77,17 @@ def _get_4bit_code(quant_type: str, device: torch.device) -> torch.Tensor:
     return _code_4bit_cache[key]
 
 
+@functools.cache
+def _warn_gemm_4bit_unaligned(K: int, blocksize: int) -> None:
+    """Warn at most once per (K, blocksize): the misalignment is a fixed property of the
+    model, so repeating it on every gemm_4bit call is pure noise."""
+    warn(
+        f"inner dimension ({K}) is not aligned for fast kernel "
+        f"with blocksize={blocksize}, falling back to slower implementation.",
+        UserWarning,
+    )
+
+
 def get_gaudi_sw_version():
     """
     Returns the installed version of Gaudi SW.

bitsandbytes/backends/xpu/ops.py

@@ -2,7 +2,6 @@
 import ctypes as ct
 import logging
 from typing import Optional
-from warnings import warn
 
 from packaging import version
 import torch
@@ -12,7 +11,7 @@
 from ..._ops import register_kernel
 from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib
 from ..default.ops import _gemm_4bit_default_impl
-from ..utils import _get_4bit_code, triton_available
+from ..utils import _get_4bit_code, _warn_gemm_4bit_unaligned, triton_available
 
 logger = logging.getLogger(__name__)
 
@@ -186,11 +185,7 @@ def _(
                 out = out + bias
             return out
 
-        warn(
-            f"inner dimension ({K}) is not aligned for fast kernel "
-            f"with blocksize={blocksize}, falling back to slower implementation.",
-            UserWarning,
-        )
+        _warn_gemm_4bit_unaligned(K, blocksize)
 
     return _gemm_4bit_default_impl(
         A,

tests/test_ops.py

@@ -1,4 +1,5 @@
 from math import prod
+import warnings
 
 import pytest
 import torch
@@ -333,6 +334,28 @@ def test_gemm_4bit(self, device, dtype, quant_type, compress_statistics, has_bia
                 kwargs={"bias": bias},
             )
 
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_gemm_4bit_unaligned_warning(self, device):
+        """Regression test for #2027: the blocksize-alignment warning must not be emitted
+        on every call, nor at all when the fused kernel was not going to be used anyway."""
+        N, K, blocksize = 128, 3420, 64  # 3420 % 64 != 0 (Qwen2.5-VL vision tower)
+        B = torch.randn(N, K, dtype=torch.float16, device=device)
+        B_q, qs = bitsandbytes.functional.quantize_4bit(B, blocksize=blocksize, quant_type="nf4")
+
+        def run(M):
+            A = torch.randn(M, K, dtype=torch.float16, device=device)
+            with warnings.catch_warnings(record=True) as caught:
+                warnings.simplefilter("always")
+                for _ in range(2):
+                    torch.ops.bitsandbytes.gemm_4bit.default(A, B_q, list(B.shape), qs.absmax, blocksize, "nf4")
+            return [w for w in caught if "not aligned" in str(w.message)]
+
+        # Large M always takes the dequant+F.linear path, aligned or not.
+        assert run(1024) == []
+
+        # When alignment does decide it, warn at most once per (K, blocksize).
+        assert len(run(1)) <= 1
+
     @pytest.mark.parametrize("device", get_available_devices())
     @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype)
     @pytest.mark.parametrize("offset_dtype", [torch.float16, torch.bfloat16], ids=describe_dtype)