from collections.abc import Sequence import ctypes as ct 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 def _setup_ctypes(names, argtypes, restype=None): for name in names: fn = getattr(lib, name) fn.argtypes = argtypes fn.restype = restype # 4-bit/8-bit dequantize: (code, A, absmax, out, blocksize, numel, stream) _setup_ctypes( [f"cdequantize_blockwise_{d}_{q}" for d in ("fp32", "bf16", "fp16") for q in ("nf4", "fp4")] + [f"cdequantize_blockwise_{d}" for d in ("fp32", "bf16", "fp16")], [ct.c_void_p] * 4 + [ct.c_int32, ct.c_int32, ct.c_void_p], ) # 4-bit GEMM: (A, B, absmax, absmax_8bit, absmax_code, absmax_offset, out, bias, M, N, K, blocksize, quant_type, stream) _setup_ctypes( [f"cgemm_4bit_{d}" for d in ("bf16", "fp16", "fp32")], [ct.c_void_p] * 8 + [ct.c_int32, ct.c_int32, ct.c_int32, ct.c_int32, ct.c_int32, ct.c_void_p], ) # 4-bit GEMV: (m, n, k, A, B, absmax, code, out, lda, ldb, ldc, blocksize, stream) _setup_ctypes( [f"cgemm_4bit_inference_naive_{d}" for d in ("bf16", "fp16", "fp32")], [ct.c_int32] * 3 + [ct.c_void_p] * 5 + [ct.c_int32] * 3 + [ct.c_int32, ct.c_void_p], ) # int8 igemm: (ctx, m, n, k, A, B, C, rowscale, lda, ldb, ldc, stream) -> int32 _setup_ctypes( ["cigemmlt_32"], [ct.c_void_p] + [ct.c_int32] * 3 + [ct.c_void_p] * 4 + [ct.c_int32] * 3 + [ct.c_void_p], restype=ct.c_int32, ) # int8 mm dequant: (A, row_stats, col_stats, out, bias, numRows, numCols, stream) _setup_ctypes( ["cdequant_mm_int32_fp16"], [ct.c_void_p] * 5 + [ct.c_int32, ct.c_int32, ct.c_void_p], ) # int8 vectorwise quant: (A, out, row_stats, threshold, rows, cols, stream) _setup_ctypes( ["cint8_vector_quant"], [ct.c_void_p] * 3 + [ct.c_float, ct.c_int32, ct.c_int32, ct.c_void_p], ) # 4-bit/8-bit blockwise quantize: (code, A, absmax, out, blocksize, n) _setup_ctypes( [f"cquantize_blockwise_{d}_{q}" for d in ("fp32", "bf16", "fp16") for q in ("nf4", "fp4")] + [f"cquantize_blockwise_{d}" for d in ("fp32", "bf16", "fp16")], [ct.c_void_p] * 4 + [ct.c_int32, ct.c_int32], ) _get_raw_stream = torch._C._cuda_getCurrentRawStream @functools.cache def _gpu_dispatch_props(device_index): props = torch.cuda.get_device_properties(device_index) return props.multi_processor_count, props.major, props.minor @register_kernel("bitsandbytes::int8_linear_matmul", "cuda") def _(A: torch.Tensor, B: torch.Tensor): out = torch.empty((*A.shape[:-1], B.shape[0]), device=A.device, dtype=torch.int32) return _int8_linear_matmul_impl(A, B, out) @register_kernel("bitsandbytes::int8_linear_matmul.out", "cuda") def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): _int8_linear_matmul_impl(A, B, out) def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): A, B = B, A shapeA = A.shape shapeB = B.shape if A.dtype != torch.int8: raise ValueError("B must be int8") if B.dtype != torch.int8: raise ValueError("A must be int8") if A.ndim != 2: raise ValueError("Only two dimensional matrices are supported for argument B") if B.ndim not in (2, 3): raise ValueError("Only two or three dimensional matrices are supported for argument A") if prod(shapeB) <= 0: raise ValueError(f"Input tensor dimensions need to be > 0: {shapeB}") if out.dtype != torch.int32: raise ValueError(f"out must be int32, got {out.dtype}") shapeC = (*shapeB[:-1], shapeA[0]) if out.shape != shapeC: raise ValueError(f"Output shape {out.shape} does not match expected shape {shapeC}") k, m = shapeA n = prod(shapeB[:-1]) lda = shapeA[-1] # Weights (outputs, inputs) ldb = shapeB[-1] # Activations (batch, tokens, inputs) ldc = shapeC[-1] # Output (batch, tokens, outputs) if lda != ldb: raise ValueError( f"int8_linear_matmul only supports B^T @ A. Inner dimensions do not match: B @ A = {shapeB} @ {shapeA}" ) # cuBLASLt does not support int8 matmul with inner dimensions that are not divisible by 4. # We'll fall back to a slower fp32 calculation in this circumstance. # Fortunately, this should not be very common. if lda % 4 != 0: result = torch.matmul(B.float(), A.float().t()).to(torch.int32) return out.copy_(result) with _cuda_device_of(A): ctx = CUBLAS_Context.get_instance().get_context(A.device) has_error = lib.cigemmlt_32( ctx, m, n, k, A.data_ptr(), B.data_ptr(), out.data_ptr(), None, lda, ldb, ldc, _get_raw_stream(A.device.index), ) if has_error: if has_error == 100: # `ERR_NOT_IMPLEMENTED` is defined as 100 in `ops.cu`. The HIP backend # also returns this when no usable hipBLASLt algo exists for the shape # (seen on MI300X for some small-n int8 gemms). Fall back to fp32 — same # path used for the `lda % 4 != 0` case above. import warnings warnings.warn( f"int8_linear_matmul has no usable (hip|cu)blasLt algo for shape " f"{shapeA=} {shapeB=}; falling back to fp32 matmul.", RuntimeWarning, stacklevel=2, ) result = torch.matmul(B.float(), A.float().t()).to(torch.int32) return out.copy_(result) else: raise RuntimeError( f"cublasLt ran into an error!\n\t{shapeA=}, {shapeB=}, {shapeC=}\n\t{(lda, ldb, ldc)=}\n\t{(m, n, k)=}" ) return out @register_kernel("bitsandbytes::int8_mm_dequant", "cuda") def _( A: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, dtype: Optional[torch.dtype] = None, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: if A.dtype != torch.int32: raise ValueError(f"A must be int32, got {A.dtype}") if row_stats.dtype != torch.float32: raise ValueError(f"row_stats must be float32, got {row_stats.dtype}") if col_stats.dtype != torch.float32: raise ValueError(f"col_stats must be float32, got {col_stats.dtype}") # Note: cuda kernel only currently supports fp16 output. # We'll later cast to desired dtype if needed. out = torch.empty_like(A, dtype=torch.float16) # Note: fused bias in the kernel is only supported for fp16 # TODO(matthewdouglas): Consider supporting bf16 fused bias bias_ptr = bias.data_ptr() if bias is not None and bias.dtype == torch.float16 else None with _cuda_device_of(A): lib.cdequant_mm_int32_fp16( A.data_ptr(), row_stats.data_ptr(), col_stats.data_ptr(), out.data_ptr(), bias_ptr, A.numel() // A.shape[-1], A.shape[-1], _get_raw_stream(A.device.index), ) # Add bias separately if not fused in kernel if bias is not None and bias.dtype != torch.float16: out.add_(bias) return out.to(dtype or torch.float16) @register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda") def _(A: torch.Tensor, threshold=0.0): if A.dtype != torch.float16: raise ValueError(f"A must be float16, got {A.dtype}") if threshold < 0.0: raise ValueError("threshold must be non-negative") rows = A.numel() // A.shape[-1] cols = A.shape[-1] row_stats = torch.empty(rows, device=A.device, dtype=torch.float32) out_row = torch.empty(A.shape, device=A.device, dtype=torch.int8) outlier_cols = None if threshold > 0.0: # TODO we could improve perf of this outliers = A.abs() >= threshold if outliers.any(): outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1) else: # Needed for torch.compile support. outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64) with _cuda_device_of(A): lib.cint8_vector_quant( A.data_ptr(), out_row.data_ptr(), row_stats.data_ptr(), threshold, rows, cols, _get_raw_stream(A.device.index), ) # Zero out values from outlier columns across all rows. # The kernel will handle this for outliers themselves, so we can optimize for rows=1. if rows > 1 and outlier_cols is not None: out_row[:, outlier_cols] = 0 return out_row, row_stats, outlier_cols @register_kernel("bitsandbytes::int8_double_quant", "cuda") def _( A: torch.Tensor, threshold=0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: # Use CUDA kernel for rowwise quant and outlier column detection quant_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant.default( A, threshold=threshold, ) # PyTorch impl for colwise col_stats, outlier_mask = _get_col_absmax(A, threshold=threshold) if threshold > 0.0 and outlier_mask is not None: A = A.masked_fill(outlier_mask, 0.0) quant_col = torch.round(A.mul(127.0) / col_stats.unsqueeze(0)).to(torch.int8) return quant_row, quant_col, row_stats, col_stats.flatten().float(), outlier_cols def _get_col_absmax( A: torch.Tensor, threshold=0.0, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: if not A.is_floating_point(): raise ValueError(f"A must be a floating point tensor, got {A.dtype}") outlier_mask = None absA = A.abs().view(-1, A.shape[-1]) if threshold > 0.0: # Filter outliers from stats when enabled outlier_mask = absA >= threshold absA.masked_fill_(outlier_mask, 0.0) # shape [cols]; unsqueeze(0) gives [1,cols] col_stats = absA.amax(dim=0, keepdim=False).float() return col_stats, outlier_mask @register_kernel("bitsandbytes::quantize_blockwise", "cuda") def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: A = A.contiguous() if code.dtype != torch.float32: raise ValueError(f"code must be float32, got {code.dtype}") if blocksize not in (64, 128, 256, 512, 1024, 2048, 4096): raise ValueError(f"invalid blocksize {blocksize}") n = A.numel() blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty_like(A, dtype=torch.uint8) if A.dtype == torch.float32: fn = lib.cquantize_blockwise_fp32 elif A.dtype == torch.float16: fn = lib.cquantize_blockwise_fp16 elif A.dtype == torch.bfloat16: fn = lib.cquantize_blockwise_bf16 else: raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}") with _cuda_device_of(A): fn( code.data_ptr(), A.data_ptr(), absmax.data_ptr(), out.data_ptr(), blocksize, n, ) return out, absmax @register_kernel("bitsandbytes::dequantize_blockwise", "cuda") def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor: out = torch.empty_like(A, dtype=dtype) _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out) return out @register_kernel("bitsandbytes::dequantize_blockwise.out", "cuda") def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor, ) -> None: if out.dtype != dtype: raise ValueError(f"Expected out.dtype == {dtype}, got {out.dtype}") if out.shape != A.shape: raise ValueError(f"Expected out.shape == {A.shape}, got {out.shape}") _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out) def _dequantize_blockwise_impl( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor ) -> None: A = A.contiguous() if dtype == torch.float32: fn = lib.cdequantize_blockwise_fp32 elif dtype == torch.float16: fn = lib.cdequantize_blockwise_fp16 elif dtype == torch.bfloat16: fn = lib.cdequantize_blockwise_bf16 else: raise ValueError(f"Blockwise dequantization only supports 16/32-bit floats, but got {dtype}") with _cuda_device_of(A): fn( code.data_ptr(), A.data_ptr(), absmax.data_ptr(), out.data_ptr(), blocksize, A.numel(), _get_raw_stream(A.device.index), ) @register_kernel("bitsandbytes::quantize_4bit", "cuda") def _( A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: A = A.contiguous() n = A.numel() blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty(((n + 1) // (quant_storage.itemsize * 2), 1), device=A.device, dtype=quant_storage) if A.dtype == torch.bfloat16: if quant_type == "fp4": fn = lib.cquantize_blockwise_bf16_fp4 else: fn = lib.cquantize_blockwise_bf16_nf4 elif A.dtype == torch.float16: if quant_type == "fp4": fn = lib.cquantize_blockwise_fp16_fp4 else: fn = lib.cquantize_blockwise_fp16_nf4 elif A.dtype == torch.float32: if quant_type == "fp4": fn = lib.cquantize_blockwise_fp32_fp4 else: fn = lib.cquantize_blockwise_fp32_nf4 with _cuda_device_of(A): fn( None, A.data_ptr(), absmax.data_ptr(), out.data_ptr(), blocksize, n, ) return out, absmax @register_kernel("bitsandbytes::dequantize_4bit", "cuda") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: out = torch.empty(shape, dtype=dtype, device=A.device) _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) return out @register_kernel("bitsandbytes::dequantize_4bit.out", "cuda") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, out: torch.Tensor, ) -> None: if out.shape != tuple(shape): raise ValueError(f"Expected out.shape == {shape}, got {out.shape}") if out.dtype != dtype: raise ValueError(f"Expected out.dtype == {dtype}, got {out.dtype}") _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) def _dequantize_4bit_impl( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, dtype: torch.dtype, out: torch.Tensor, ) -> None: A = A.contiguous() if dtype == torch.bfloat16: if quant_type == "fp4": fn = lib.cdequantize_blockwise_bf16_fp4 else: fn = lib.cdequantize_blockwise_bf16_nf4 elif dtype == torch.float16: if quant_type == "fp4": fn = lib.cdequantize_blockwise_fp16_fp4 else: fn = lib.cdequantize_blockwise_fp16_nf4 elif dtype == torch.float32: if quant_type == "fp4": fn = lib.cdequantize_blockwise_fp32_fp4 else: fn = lib.cdequantize_blockwise_fp32_nf4 else: raise ValueError(f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}") with _cuda_device_of(A): fn( None, A.data_ptr(), absmax.data_ptr(), out.data_ptr(), blocksize, out.numel(), _get_raw_stream(A.device.index), ) @register_kernel("bitsandbytes::gemv_4bit", "cuda") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int ) -> torch.Tensor: shape = (*A.shape[:-1], shapeB[0]) out = torch.empty(shape, device=A.device, dtype=A.dtype) _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out) return out @register_kernel("bitsandbytes::gemv_4bit.out", "cuda") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: expected_shape = (*A.shape[:-1], shapeB[0]) if out.shape != expected_shape: raise ValueError(f"Expected out.shape == {expected_shape}, got {out.shape}") if out.dtype != A.dtype: raise ValueError(f"Expected out.dtype == {A.dtype}, got {out.dtype}") _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out) def _gemv_4bit_impl( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: if blocksize not in (32, 64, 128, 256, 512, 1024, 2048, 4096): raise ValueError(f"invalid blocksize {blocksize}") # Note: these checks are not strictly necessary, and cost more than they are worth, so they are commented out for now. # torch._check( # A.numel() == A.size(-1), # lambda: f"A must be a vector with leading dimensions of 1, got {A.shape}", # ) # torch._check( # A.dtype in [torch.float16, torch.bfloat16, torch.float32], # lambda: f"A must be float16, bfloat16, or float32, got {A.dtype}", # ) # torch._check( # B.dtype in [torch.uint8, torch.bfloat16, torch.float16, torch.float32], # lambda: f"B must be backed by storage of type uint8, bfloat16, float16, or float32, got {B.dtype}", # ) # torch._check(absmax.dtype == torch.float32, lambda: f"absmax must be float32, got {absmax.dtype}") # torch._check(code.dtype == torch.float32, lambda: f"code must be float32, got {code.dtype}") m = shapeB[0] n = 1 k = shapeB[1] lda = m ldb = (A.shape[-1] + 1) // 2 ldc = m if A.dtype == torch.float16: fn = lib.cgemm_4bit_inference_naive_fp16 elif A.dtype == torch.bfloat16: fn = lib.cgemm_4bit_inference_naive_bf16 elif A.dtype == torch.float32: fn = lib.cgemm_4bit_inference_naive_fp32 with _cuda_device_of(A): fn( m, n, k, A.data_ptr(), B.data_ptr(), absmax.data_ptr(), code.data_ptr(), out.data_ptr(), lda, ldb, ldc, blocksize, _get_raw_stream(A.device.index), ) @functools.cache def _gemm_4bit_use_custom_cuda(device_index, dtype, M, N, K): """Custom kernel vs dequant+F.linear heuristic for M in [5, 1536]. Per-arch notes (bf16/fp16, M >= 8, large weight): sm75 (T4, ~300 GB/s GDDR6): fp16 MMA only; GDDR makes dequant expensive. sm80 (A100, ~2 TB/s HBM2e): mma.sync; HBM thresholds; K-heavy shapes handled explicitly. sm86 (A10, ~600 GB/s GDDR6): dedicated block; wider M caps than sm89 at medium N. sm89 (4090, L40S, GDDR6X): default fallback; tall-K and large-N get higher M caps. sm90 (H100/H200, HBM3/HBM3e): dequant+linear is much faster; thresholds are tight. sm100 (B200/B300, HBM3e): exits early at top of function. sm120 (RTX 5000, GDDR7): dedicated block; medium-N tiers differ from sm89. sm121 (GB10 DGX Spark, LPDDR5X): dedicated block at >=1 wave only; the low memory bandwidth keeps the custom kernel ahead through M=256 there. Sub-wave shapes use the sm89 tiers below. """ if M <= _GEMM_4BIT_CUSTOM_FLOOR_M: return True num_sms, major, minor = _gpu_dispatch_props(device_index) n_blocks = (N + 63) // 64 # fp32 has no MMA kernel; pre-sm75 has no MMA kernel; sm75 has fp16 MMA only. # For all of these, custom only wins in the SIMT range (M<8). if dtype == torch.float32 or major < 7: return M < 8 if major == 7 and (minor < 5 or dtype != torch.float16): return M < 8 # sm87 and sm110: no calibration data, conservative fallback. if (major == 8 and minor == 7) or major == 11: return False # sm100 (B200/B300): dequant+F.linear is significantly faster than our mma.sync kernel. if major == 10: if n_blocks >= num_sms * 3: return M <= 32 if n_blocks >= num_sms: return False if K >= N else M <= 8 return False is_sm75 = major == 7 and minor == 5 is_sm80 = major == 8 and minor == 0 is_sm86 = major == 8 and minor == 6 is_sm90 = major == 9 is_sm120 = major == 12 and minor == 0 is_sm121 = major == 12 and minor == 1 is_hbm = is_sm80 or is_sm90 # sm100 already returned above tall_k_2xn = K > N * 2 # Small-weight path (N*K < 4MB): dequant overhead dominates. if N * K < 4 * 1024 * 1024: if K * 2 < N: # Very short K (K < N/2): latency-dominated, custom 3-9x cheaper. if is_hbm: # Calibrated on A100: custom wins to M=1536 (low wave), M=512 (high wave). # Calibrated on H100/H200: custom wins to M=512 (low wave), M=320 (high wave). low_wave = n_blocks * 3 < num_sms if is_sm80: return M <= (1536 if low_wave else 512) return M <= (512 if low_wave else 320) if is_sm75: # T4: wins require >=3 waves; M cap scales with K depth. if n_blocks >= num_sms * 3: return M <= 320 if K >= 1024: return M <= 64 if K >= 704: return M <= 96 return M <= 320 # sm86/sm89/sm120: well-subscribed wins to M=320; undersubscribed tighter. if n_blocks >= num_sms: return M <= 320 return M <= 192 if n_blocks * K > num_sms * 320 else M <= 320 # K*2 >= N: arch-specific handling at low occupancy. quarter_wave = n_blocks * 4 <= num_sms if is_sm80 and quarter_wave: # A100 <1/4 wave: K>=N loses earlier (K-tiling efficient on HBM2e). if K >= N: return M <= (32 if n_blocks * 8 <= num_sms else 128) return M <= 384 # T4 <1 wave non-short-K: M>8 routes through occupancy caps below. if is_sm75 and n_blocks < num_sms and M > 8: return M <= 64 # General tiers (sm90, sm86, sm89, sm120): # GDDR tall-K (K>=N) at <1/4 wave: K-tiling in default impl wins above M=23. if quarter_wave: return M <= (32 if (K < N or is_hbm) else 23) if n_blocks * 2 <= num_sms: return M <= 16 return False # >=1/2 wave: no validated wins for remaining small-weight shapes # Non-small-weight: custom wins up to M=512; dequant+F.linear wins above that. if M > 512: return False # M=5-7: custom SIMT generally wins because dequant cost dominates. # Exceptions where K-tiling efficiency or MMA occupancy favors dequant+F.linear: # HBM at M=6-7: tall-K (K>N) at ~3/4 MMA wave. # sm90 square (K==N) at specific occupancy bands: arch-specific crossover. if M < 8: hbm_m67_thresh = 36 if is_sm90 else 48 if is_hbm and M >= 6 and n_blocks >= hbm_m67_thresh: lt_75pct_wave = n_blocks * 4 < num_sms * 3 lt_60pct_wave = n_blocks * 5 < num_sms * 3 # Tall-K: K-tiling in default impl wins when under-subscribed. if K > N and lt_75pct_wave: return False # Square: arch-specific crossover around 0.6 wave. # A100 (HBM2e): loses below 0.6 wave. H100/H200 (HBM3/3e): loses above. if K == N: if is_sm80 and lt_60pct_wave: return False if is_sm90 and lt_75pct_wave and not lt_60pct_wave: return False return True # M in [8, 512]: per-arch tier ladders. if is_sm75: # fp16 MMA (m16n8k8). GDDR bandwidth makes dequant relatively expensive. if n_blocks >= num_sms * 3: return M <= (128 if K < N else 64) if n_blocks >= num_sms // 2: return M <= 64 return M <= 32 if is_sm80: # mma.sync (m16n8k16). HBM2e thresholds; K-heavy shapes handled explicitly. if n_blocks >= num_sms * 3: return M <= 128 if n_blocks >= num_sms: return M <= (64 if K < N else 32) # Very tall-K (K>=3N) at >1/4 wave: K-tiling in default impl wins at all M. # Uses >= to catch K==3N (e.g. N=4096,K=12288 M=9-16: measured regression on A100). if K >= N * 3 and n_blocks * 4 > num_sms: return False # Square (K==N) at 0.5-1 wave: K-tiling wins at ~0.6 wave. # n_blocks>=48 excludes small N where SIMT still wins. if K == N and n_blocks >= 48 and n_blocks * 5 < num_sms * 3: return False # <0.5 wave: K<=N custom wins to M=128; K>N default wins above wave threshold. if n_blocks * 2 < num_sms: if K <= N: return M <= 128 if n_blocks * 3 >= num_sms: return False # 0.5-1 wave K= num_sms // 2 and K < N: return M <= 128 return M <= 16 if is_sm86: # ~600-940 GB/s GDDR6/GDDR6X. Dedicated block: sm89 fallback tiers are too # loose for 600 GB/s bandwidth and cause regressions at medium N (~N=4096). if n_blocks >= num_sms: return M <= 128 if n_blocks >= num_sms // 2: return M <= 64 return M <= 16 if is_sm90: # HBM3/HBM3e. dequant+F.linear (WGMMA path) is significantly faster than our # mma.sync kernel; thresholds are calibrated conservatively (H100/H200 share path). if n_blocks >= num_sms * 3: return M <= 64 if n_blocks >= num_sms * 2: return M <= 48 if n_blocks >= num_sms: return M <= 32 if n_blocks >= num_sms // 2: # Square/tall-K at <3/4 wave: K-tiling too efficient on HBM3e. if K >= N and n_blocks * 4 < num_sms * 3: return False return M <= 16 return False if is_sm121: # GB10 (DGX Spark): unified LPDDR5X. Calibrated at >=1 wave only; below one # wave the crossover is strongly K-dependent, so those shapes use the shared # tiers below (which already branch on tall-K). if n_blocks >= num_sms: return M <= 256 if is_sm120: # GDDR7 (~1-1.8 TB/s). Medium-N threshold tiers differ from sm89. if n_blocks >= num_sms * 3: return M <= 256 if n_blocks >= num_sms * 2: return M <= 128 # Short-K (K= num_sms * 4: return M <= (96 if K >= N else 64) if n_blocks >= num_sms: return M <= 64 if n_blocks >= num_sms // 2: # Large-N (n_blocks>=128, N>=8192) with K>=N/2: calibrated on RTX Pro 6000 to M=64. return M <= (64 if (K * 2 >= N and n_blocks >= 128) else 8) if tall_k_2xn and n_blocks > 64: return M <= 16 return M <= 8 # Fallback: sm89 (4090, L40S, L4), sm121 below one wave, unrecognized arches. # GDDR bandwidth makes dequant relatively expensive so custom wins at higher M. if n_blocks >= num_sms * 3: return M <= 256 if n_blocks >= num_sms * 2: return M <= 128 # Near-wave (~0.8x): tall-K and very large N (n_blocks>=200, N>=14336) raise cap to M=128. # N=10240 (n_blocks=160) deliberately excluded to avoid regressions there. if n_blocks * 5 >= num_sms * 4: if tall_k_2xn or n_blocks >= 200: return M <= 128 # Square/tall-K: >=60 SMs wins to M=128; <60 SMs default wins earlier. if K >= N: return M <= (128 if num_sms >= 60 else 32) return M <= 64 if n_blocks >= num_sms // 2: if tall_k_2xn: return M <= 64 if n_blocks >= 64: return M <= 8 return M <= 32 # Tall-K (K>N) at narrow N (n_blocks<=48): M-driven crossover. # K>=3N (e.g. N=2560,K=10240): SIMT wins to M=12. Moderate K>N: M=10. if K > N and n_blocks <= 48: return M <= (12 if K >= N * 3 else 10) return M <= (16 if (tall_k_2xn or n_blocks < 48) else 8) @functools.cache def _gemm_4bit_use_custom_rocm(device_index, dtype, M, N, K): """ Fused SIMT kernel vs dequant+F.linear heuristic for ROCm. RDNA3/RDNA4 calibration keeps the SIMT kernel through ~M=8. gfx125x (CDNA5) is Wave32 but lacks the v_dot2 SIMT instructions, so its kernel takes the fp32-FMA CDNA math path. The crossover below is left at the gfx12 default (M<=8) pending dedicated calibration. CDNA/gfx9 is calibrated on MI308X (gfx942): bf16/fp16 win through M<=4 after the SIMT math-path tuning, while fp32 only has a broad win through M<=2. TODO: revisit once WMMA/MFMA kernels land; calibrate the gfx125x crossover. """ if M <= _GEMM_4BIT_CUSTOM_FLOOR_M and dtype != torch.float32: return True arch = _rocm_gfx_arch(device_index) if arch.startswith("gfx11") or arch.startswith("gfx12"): # RDNA3 / RDNA4 / CDNA5 (gfx125x) return M <= 8 if arch.startswith("gfx9"): # CDNA1-4 / MI-series (Wave64) return M <= (2 if dtype == torch.float32 else 4) return M <= 4 # unknown ROCm arch: conservative tiny-batch floor @functools.cache def _rocm_gfx_arch(device_index): """gfx arch string (e.g. 'gfx1100') for a ROCm device, feature flags stripped.""" name = getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "") or "" return name.split(":")[0] def _gemm_4bit_kernel_impl( A, B, shapeB, absmax, blocksize, quant_type, bias=None, absmax_8bit=None, absmax_code=None, absmax_offset=None ): """Invoke the fused cgemm_4bit_* kernel (shared by the CUDA and ROCm dispatch; the C dispatch in gemm_4bit.cu picks SIMT vs MMA per arch/shape). A is made contiguous because the kernel reads it as row-major (stride K).""" K = A.shape[-1] M = A.numel() // K N = shapeB[0] if K != shapeB[1]: raise RuntimeError(f"A inner dim ({K}) does not match weight ({shapeB[1]})") if absmax.dtype != torch.float32: raise RuntimeError(f"absmax must be float32, got {absmax.dtype}") if bias is not None: if bias.ndim != 1: raise RuntimeError(f"bias must be 1D, got {bias.ndim}D") if bias.dtype != A.dtype: raise RuntimeError(f"bias dtype ({bias.dtype}) must match A dtype ({A.dtype})") A = A.contiguous() quant_type_int = 1 if quant_type == "fp4" else 2 out = torch.empty((*A.shape[:-1], N), dtype=A.dtype, device=A.device) stream = _get_raw_stream(A.device.index) if A.dtype == torch.bfloat16: fn = lib.cgemm_4bit_bf16 elif A.dtype == torch.float16: fn = lib.cgemm_4bit_fp16 elif A.dtype == torch.float32: fn = lib.cgemm_4bit_fp32 else: raise RuntimeError(f"unsupported dtype {A.dtype}") # Offset is expected to be a float32 tensor. absmax_offset_f32 = absmax_offset.to(dtype=torch.float32) if absmax_offset is not None else None with _cuda_device_of(A): fn( A.data_ptr(), B.data_ptr(), absmax.data_ptr(), absmax_8bit.data_ptr() if absmax_8bit is not None else None, absmax_code.data_ptr() if absmax_code is not None else None, absmax_offset_f32.data_ptr() if absmax_offset_f32 is not None else None, out.data_ptr(), bias.data_ptr() if bias is not None else None, M, N, K, blocksize, quant_type_int, stream, ) return out def _dequant_linear_fallback( A, B, shapeB, absmax, blocksize, quant_type, bias=None, absmax_8bit=None, absmax_code=None, absmax_offset=None ): """Unfused fallback shared by CUDA and ROCm: reconstruct the (optionally nested) absmax, dequantize the 4-bit weight via the backend dequant impls (reusing preallocated buffers), then F.linear.""" if absmax_8bit is not None: absmax_dq = torch.empty_like(absmax_8bit, dtype=torch.float32) _dequantize_blockwise_impl(absmax_8bit, absmax, absmax_code, 256, torch.float32, out=absmax_dq) absmax = absmax_dq + absmax_offset B_dq = torch.empty(shapeB, dtype=A.dtype, device=A.device) _dequantize_4bit_impl(B, absmax, blocksize, quant_type, A.dtype, out=B_dq) return torch.nn.functional.linear(A, B_dq, bias) # Unified CUDA/ROCm dispatch for bitsandbytes::gemm_4bit. The choice *among* custom # kernels (CUDA SIMT vs MMA; ROCm SIMT) is made in the C dispatch (csrc/gemm_4bit.cu). _GEMM_4BIT_CUSTOM_FLOOR_M = 4 if torch.version.hip is None: _gemm_4bit_use_custom_fn = _gemm_4bit_use_custom_cuda # CUDA: dequant+F.linear wins past M=1536 (dequant savings negligible at very # large batch). _gemm_4bit_custom_max_m = 1536 else: _gemm_4bit_use_custom_fn = _gemm_4bit_use_custom_rocm # ROCm: the custom path is SIMT-only today; the per-arch heuristic above owns # RDNA/CDNA thresholds. Keep a hard upper cap while WMMA/MFMA paths are absent. _gemm_4bit_custom_max_m = 256 @register_kernel("bitsandbytes::gemm_4bit", "cuda") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, blocksize: int, quant_type: str, bias: Optional[torch.Tensor] = None, absmax_8bit: Optional[torch.Tensor] = None, absmax_code: Optional[torch.Tensor] = None, absmax_offset: Optional[torch.Tensor] = None, ) -> torch.Tensor: K = A.shape[-1] M = A.numel() // K N = shapeB[0] # The backend-specific heuristic owns tiny-M floors and per-arch thresholds. # Past custom_max_m (or for blocksize-misaligned K), use the dequant+F.linear # 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 not use_custom: return _dequant_linear_fallback( A, B, shapeB, absmax, blocksize, quant_type, bias, absmax_8bit=absmax_8bit, absmax_code=absmax_code, absmax_offset=absmax_offset, ) return _gemm_4bit_kernel_impl( A, B, shapeB, absmax, blocksize, quant_type, bias, absmax_8bit, absmax_code, absmax_offset ) """C FUNCTIONS FOR OPTIMIZERS""" str2optimizer32bit = { "adam": ( lib.cadam32bit_grad_fp32, lib.cadam32bit_grad_fp16, lib.cadam32bit_grad_bf16, ), "momentum": ( lib.cmomentum32bit_grad_32, lib.cmomentum32bit_grad_16, ), "rmsprop": ( lib.crmsprop32bit_grad_32, lib.crmsprop32bit_grad_16, ), "lion": ( lib.clion32bit_grad_fp32, lib.clion32bit_grad_fp16, lib.clion32bit_grad_bf16, ), "adagrad": ( lib.cadagrad32bit_grad_32, lib.cadagrad32bit_grad_16, ), "lamb": ( lib.cadam32bit_grad_fp32, lib.cadam32bit_grad_fp16, lib.cadam32bit_grad_bf16, ), "ademamix": ( lib.cademamix32bit_grad_fp32, lib.cademamix32bit_grad_fp16, lib.cademamix32bit_grad_bf16, ), "lars": ( lib.cmomentum32bit_grad_32, lib.cmomentum32bit_grad_16, ), } str2optimizer8bit_blockwise = { "adam": ( lib.cadam_8bit_blockwise_grad_fp32, lib.cadam_8bit_blockwise_grad_fp16, lib.cadam_8bit_blockwise_grad_bf16, ), "momentum": ( lib.cmomentum_8bit_blockwise_grad_fp32, lib.cmomentum_8bit_blockwise_grad_fp16, lib.cmomentum_8bit_blockwise_grad_bf16, ), "rmsprop": ( lib.crmsprop_8bit_blockwise_grad_fp32, lib.crmsprop_8bit_blockwise_grad_fp16, lib.crmsprop_8bit_blockwise_grad_bf16, ), "lion": ( lib.clion_8bit_blockwise_grad_fp32, lib.clion_8bit_blockwise_grad_fp16, lib.clion_8bit_blockwise_grad_bf16, ), "adagrad": ( lib.cadagrad_8bit_blockwise_grad_fp32, lib.cadagrad_8bit_blockwise_grad_fp16, lib.cadagrad_8bit_blockwise_grad_bf16, ), "ademamix": ( lib.cademamix_8bit_blockwise_grad_fp32, lib.cademamix_8bit_blockwise_grad_fp16, lib.cademamix_8bit_blockwise_grad_bf16, ), } def _optimizer_update_32bit_impl( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float, skip_zeros=False, ) -> None: optim_fns = str2optimizer32bit.get(optimizer_name, None) if optim_fns is None: raise ValueError( f"Unsupported optimizer name: {optimizer_name}. Supported optimizers: {list(str2optimizer32bit.keys())}" ) if g.dtype == torch.float32: optim_func = optim_fns[0] elif g.dtype == torch.float16: optim_func = optim_fns[1] elif g.dtype == torch.bfloat16 and len(optim_fns) == 3: optim_func = optim_fns[2] else: raise ValueError( f"Gradient+optimizer bit data type combination not supported: grad {g.dtype}, optimizer {state1.dtype}", ) with _cuda_device_of(g): optim_func( get_ptr(g), get_ptr(p), get_ptr(state1), get_ptr(state2), get_ptr(unorm_vec), ct.c_float(max_unorm), ct.c_float(param_norm), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(beta3), ct.c_float(alpha), ct.c_float(eps), ct.c_float(weight_decay), ct.c_int32(step), ct.c_float(lr), ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel()), ) def _optimizer_update_8bit_blockwise_impl( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, alpha: float, eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float, gnorm_scale: float, skip_zeros=False, ) -> None: # torch._check( # g.numel() == p.numel(), # lambda: f"g and p must have the same number of elements, got {g.numel()} and {p.numel()}", # ) # compute_dtypes = [torch.float16, torch.bfloat16, torch.float32] # torch._check( # g.dtype in compute_dtypes, # lambda: f"g must be bfloat16, float16, or float32, got {g.dtype}", # ) # torch._check( # g.dtype == p.dtype, # lambda: f"Expected all tensors to have the same dtype, got g.dtype={g.dtype}, p.dtype={p.dtype}", # ) # torch._check( # state1.dtype == torch.uint8, # lambda: f"state1 must be uint8, got {state1.dtype}", # ) # torch._check( # qmap1.dtype == absmax1.dtype == torch.float32, # lambda: f"Expected qmap1 and absmax1 to be float32, got qmap1.dtype={qmap1.dtype}, absmax1.dtype={absmax1.dtype}", # ) # if state2 is not None: # torch._check( # state2.dtype == torch.uint8, # lambda: f"state2 must be uint8, got {state2.dtype}", # ) # torch._check( # qmap2.dtype == absmax2.dtype == torch.float32, # lambda: f"Expected qmap2 and absmax2 to be float32, got qmap2.dtype={qmap2.dtype}, absmax2.dtype={absmax2.dtype}", # ) optimizer_fns = str2optimizer8bit_blockwise.get(optimizer_name) if optimizer_fns is None: raise ValueError( f"Unsupported optimizer name: {optimizer_name}. Supported optimizers: {list(str2optimizer8bit_blockwise.keys())}" ) if g.dtype == torch.float32: optimizer_fn = optimizer_fns[0] elif g.dtype == torch.float16: optimizer_fn = optimizer_fns[1] elif g.dtype == torch.bfloat16: optimizer_fn = optimizer_fns[2] else: raise ValueError( f"Unsupported gradient dtype: {g.dtype}. Supported dtypes: torch.float32, torch.float16, torch.bfloat16" ) with _cuda_device_of(g): optimizer_fn( get_ptr(p), get_ptr(g), get_ptr(state1), get_ptr(state2), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(beta3), ct.c_float(alpha), ct.c_float(eps), ct.c_int32(step), ct.c_float(lr), get_ptr(qmap1), get_ptr(qmap2), get_ptr(absmax1), get_ptr(absmax2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel()), ) register_kernel("bitsandbytes::optimizer_update_8bit_blockwise", "cuda")(_optimizer_update_8bit_blockwise_impl) register_kernel("bitsandbytes::optimizer_update_32bit", "cuda")(_optimizer_update_32bit_impl)